Merge pull request #22519 from jayfurmanek:nccl2_configure
PiperOrigin-RevId: 215310536
diff --git a/configure.py b/configure.py
index e195d65..7e47175 100644
--- a/configure.py
+++ b/configure.py
@@ -48,16 +48,12 @@
_DEFAULT_PROMPT_ASK_ATTEMPTS = 10
-_TF_WORKSPACE_ROOT = os.path.abspath(os.path.dirname(__file__))
_TF_BAZELRC_FILENAME = '.tf_configure.bazelrc'
-_TF_BAZELRC = os.path.join(_TF_WORKSPACE_ROOT, _TF_BAZELRC_FILENAME)
-_TF_WORKSPACE = os.path.join(_TF_WORKSPACE_ROOT, 'WORKSPACE')
+_TF_WORKSPACE_ROOT = ''
+_TF_BAZELRC = ''
NCCL_LIB_PATHS = [
- "lib64/",
- "lib/powerpc64le-linux-gnu/",
- "lib/x86_64-linux-gnu/",
- ""
+ 'lib64/', 'lib/powerpc64le-linux-gnu/', 'lib/x86_64-linux-gnu/', ''
]
if platform.machine() == 'ppc64le':
@@ -231,7 +227,7 @@
python_lib_path = default_python_lib_path
environ_cp['PYTHON_LIB_PATH'] = python_lib_path
- python_major_version = get_python_major_version(python_bin_path)
+ _ = get_python_major_version(python_bin_path)
# Convert python path to Windows style before writing into bazel.rc
if is_windows() or is_cygwin():
@@ -250,10 +246,10 @@
f.write('export PYTHON_BIN_PATH="%s"' % python_bin_path)
-def reset_tf_configure_bazelrc(workspace_path):
+def reset_tf_configure_bazelrc():
"""Reset file that contains customized config settings."""
open(_TF_BAZELRC, 'w').close()
- bazelrc_path = os.path.join(workspace_path, '.bazelrc')
+ bazelrc_path = os.path.join(_TF_WORKSPACE_ROOT, '.bazelrc')
data = []
if os.path.exists(bazelrc_path):
@@ -892,7 +888,7 @@
"""Set CUDNN_INSTALL_PATH and TF_CUDNN_VERSION."""
ask_cudnn_version = (
'Please specify the cuDNN version you want to use. '
- '[Leave empty to default to cuDNN %s.0]: ') % _DEFAULT_CUDNN_VERSION
+ '[Leave empty to default to cuDNN %s]: ') % _DEFAULT_CUDNN_VERSION
for _ in range(_DEFAULT_PROMPT_ASK_ATTEMPTS):
tf_cudnn_version = get_from_env_or_user_or_default(
@@ -1049,7 +1045,7 @@
for lib_file in possible_files:
if is_cuda_compatible(lib_file, cuda_ver, cudnn_ver):
matches = nvinfer_pattern.search(lib_file)
- if len(matches.groups()) == 0:
+ if not matches.groups():
continue
ver_str = matches.group(1)
ver = convert_version_to_int(ver_str) if len(ver_str) else 0
@@ -1147,18 +1143,18 @@
if os.path.exists('%s.%s' % (nccl2_path_from_ldconfig, tf_nccl_version)):
nccl_install_path = os.path.dirname(nccl2_path_from_ldconfig)
print('NCCL libraries found in ' + nccl2_path_from_ldconfig)
-
+
# Check if this is the main system lib location
if re.search('.*linux-gnu', nccl_install_path):
- trunc_nccl_install_path = "/usr"
- print("This looks like a system path.")
+ trunc_nccl_install_path = '/usr'
+ print('This looks like a system path.')
else:
- trunc_nccl_install_path = nccl_install_path + "/.."
-
+ trunc_nccl_install_path = nccl_install_path + '/..'
+
# Look for header
- nccl_hdr_path = trunc_nccl_install_path + "/include"
- print("Assuming NCCL header path is " + nccl_hdr_path)
- if os.path.exists(nccl_hdr_path + "/nccl.h"):
+ nccl_hdr_path = trunc_nccl_install_path + '/include'
+ print('Assuming NCCL header path is ' + nccl_hdr_path)
+ if os.path.exists(nccl_hdr_path + '/nccl.h'):
# Set NCCL_INSTALL_PATH
environ_cp['NCCL_INSTALL_PATH'] = nccl_install_path
write_action_env_to_bazelrc('NCCL_INSTALL_PATH', nccl_install_path)
@@ -1168,22 +1164,26 @@
write_action_env_to_bazelrc('NCCL_HDR_PATH', nccl_hdr_path)
break
else:
- print('The header for NCCL2 cannot be found. Please install the libnccl-dev package.')
+ print(
+ 'The header for NCCL2 cannot be found. Please install the libnccl-dev package.'
+ )
else:
- print('NCCL2 is listed by ldconfig but the library is not found. '
- 'Your ldconfig is out of date. Please run sudo ldconfig.')
+ print('NCCL2 is listed by ldconfig but the library is not found. '
+ 'Your ldconfig is out of date. Please run sudo ldconfig.')
else:
# NCCL is not found in ldconfig. Ask the user for the location.
default_nccl_path = environ_cp.get('CUDA_TOOLKIT_PATH')
- ask_nccl_path = (r'Please specify the location where NCCL %s library is '
- 'installed. Refer to README.md for more details. [Default '
- 'is %s]:') % (tf_nccl_version, default_nccl_path)
+ ask_nccl_path = (
+ r'Please specify the location where NCCL %s library is '
+ 'installed. Refer to README.md for more details. [Default '
+ 'is %s]:') % (tf_nccl_version, default_nccl_path)
nccl_install_path = get_from_env_or_user_or_default(
- environ_cp, 'NCCL_INSTALL_PATH', ask_nccl_path, default_nccl_path)
+ environ_cp, 'NCCL_INSTALL_PATH', ask_nccl_path, default_nccl_path)
# Result returned from "read" will be used unexpanded. That make "~"
# unusable. Going through one more level of expansion to handle that.
- nccl_install_path = os.path.realpath(os.path.expanduser(nccl_install_path))
+ nccl_install_path = os.path.realpath(
+ os.path.expanduser(nccl_install_path))
if is_windows() or is_cygwin():
nccl_install_path = cygpath(nccl_install_path)
@@ -1194,9 +1194,10 @@
nccl_lpath = '%s/lib/%s' % (nccl_install_path, nccl_lib_filename)
if not os.path.exists(nccl_lpath):
for relative_path in NCCL_LIB_PATHS:
- path = '%s/%s%s' % (nccl_install_path, relative_path, nccl_lib_filename)
+ path = '%s/%s%s' % (nccl_install_path, relative_path,
+ nccl_lib_filename)
if os.path.exists(path):
- print("NCCL found at " + path)
+ print('NCCL found at ' + path)
nccl_lib_path = path
break
else:
@@ -1205,21 +1206,25 @@
nccl_lib_path = 'lib/libnccl.%s.dylib' % tf_nccl_version
nccl_lib_path = os.path.join(nccl_install_path, nccl_lib_path)
- nccl_hdr_path = os.path.join(os.path.dirname(nccl_lib_path), '../include/nccl.h')
- print("Assuming NCCL header path is "+nccl_hdr_path)
+ nccl_hdr_path = os.path.join(
+ os.path.dirname(nccl_lib_path), '../include/nccl.h')
+ print('Assuming NCCL header path is ' + nccl_hdr_path)
if os.path.exists(nccl_lib_path) and os.path.exists(nccl_hdr_path):
# Set NCCL_INSTALL_PATH
environ_cp['NCCL_INSTALL_PATH'] = os.path.dirname(nccl_lib_path)
- write_action_env_to_bazelrc('NCCL_INSTALL_PATH', os.path.dirname(nccl_lib_path))
+ write_action_env_to_bazelrc('NCCL_INSTALL_PATH',
+ os.path.dirname(nccl_lib_path))
# Set NCCL_HDR_PATH
environ_cp['NCCL_HDR_PATH'] = os.path.dirname(nccl_hdr_path)
- write_action_env_to_bazelrc('NCCL_HDR_PATH', os.path.dirname(nccl_hdr_path))
+ write_action_env_to_bazelrc('NCCL_HDR_PATH',
+ os.path.dirname(nccl_hdr_path))
break
# Reset and Retry
- print('Invalid path to NCCL %s toolkit, %s or %s not found. Please use the '
- 'O/S agnostic package of NCCL 2' % (tf_nccl_version, nccl_lib_path,
+ print(
+ 'Invalid path to NCCL %s toolkit, %s or %s not found. Please use the '
+ 'O/S agnostic package of NCCL 2' % (tf_nccl_version, nccl_lib_path,
nccl_hdr_path))
environ_cp['TF_NCCL_VERSION'] = ''
@@ -1467,13 +1472,9 @@
raise ValueError('Cannot find the MPI library file in %s/lib' % mpi_home)
-def set_grpc_build_flags():
- write_to_bazelrc('build --define grpc_no_ares=true')
-
-
def set_system_libs_flag(environ_cp):
syslibs = environ_cp.get('TF_SYSTEM_LIBS', '')
- if syslibs and syslibs != '':
+ if syslibs:
if ',' in syslibs:
syslibs = ','.join(sorted(syslibs.split(',')))
else:
@@ -1503,6 +1504,14 @@
# TODO(pcloudy): Remove this flag when upgrading Bazel to 0.16.0
# Short object file path will be enabled by default.
write_to_bazelrc('build --experimental_shortened_obj_file_path=true')
+ # When building zip file for some py_binary and py_test targets, don't
+ # include its dependencies. This is for:
+ # 1. Running python tests against the system installed TF pip package.
+ # 2. Avoiding redundant files in
+ # //tensorflow/tools/pip_package:simple_console_windows,
+ # which is a py_binary used during creating TF pip package.
+ # See https://github.com/tensorflow/tensorflow/issues/22390
+ write_to_bazelrc('build --define=no_tensorflow_py_deps=true')
if get_var(
environ_cp, 'TF_OVERRIDE_EIGEN_STRONG_INLINE', 'Eigen strong inline',
@@ -1524,30 +1533,31 @@
def main():
+ global _TF_WORKSPACE_ROOT
+ global _TF_BAZELRC
+
parser = argparse.ArgumentParser()
parser.add_argument(
'--workspace',
type=str,
- default=_TF_WORKSPACE_ROOT,
+ default=os.path.abspath(os.path.dirname(__file__)),
help='The absolute path to your active Bazel workspace.')
args = parser.parse_args()
+ _TF_WORKSPACE_ROOT = args.workspace
+ _TF_BAZELRC = os.path.join(_TF_WORKSPACE_ROOT, _TF_BAZELRC_FILENAME)
+
# Make a copy of os.environ to be clear when functions and getting and setting
# environment variables.
environ_cp = dict(os.environ)
check_bazel_version('0.15.0')
- reset_tf_configure_bazelrc(args.workspace)
+ reset_tf_configure_bazelrc()
cleanup_makefile()
setup_python(environ_cp)
if is_windows():
- environ_cp['TF_NEED_AWS'] = '0'
- environ_cp['TF_NEED_GCP'] = '0'
- environ_cp['TF_NEED_HDFS'] = '0'
- environ_cp['TF_NEED_JEMALLOC'] = '0'
- environ_cp['TF_NEED_KAFKA'] = '0'
environ_cp['TF_NEED_OPENCL_SYCL'] = '0'
environ_cp['TF_NEED_COMPUTECPP'] = '0'
environ_cp['TF_NEED_OPENCL'] = '0'
@@ -1557,14 +1567,12 @@
# Windows.
environ_cp['TF_DOWNLOAD_CLANG'] = '0'
environ_cp['TF_ENABLE_XLA'] = '0'
- environ_cp['TF_NEED_GDR'] = '0'
- environ_cp['TF_NEED_VERBS'] = '0'
environ_cp['TF_NEED_MPI'] = '0'
environ_cp['TF_SET_ANDROID_WORKSPACE'] = '0'
if is_macos():
- environ_cp['TF_NEED_JEMALLOC'] = '0'
environ_cp['TF_NEED_TENSORRT'] = '0'
+ environ_cp['TF_ENABLE_XLA'] = '0'
# The numpy package on ppc64le uses OpenBLAS which has multi-threading
# issues that lead to incorrect answers. Set OMP_NUM_THREADS=1 at
@@ -1573,24 +1581,10 @@
if is_ppc64le():
write_action_env_to_bazelrc('OMP_NUM_THREADS', 1)
- set_build_var(environ_cp, 'TF_NEED_JEMALLOC', 'jemalloc as malloc',
- 'with_jemalloc', True)
- set_build_var(environ_cp, 'TF_NEED_GCP', 'Google Cloud Platform',
- 'with_gcp_support', True, 'gcp')
- set_build_var(environ_cp, 'TF_NEED_HDFS', 'Hadoop File System',
- 'with_hdfs_support', True, 'hdfs')
- set_build_var(environ_cp, 'TF_NEED_AWS', 'Amazon AWS Platform',
- 'with_aws_support', True, 'aws')
- set_build_var(environ_cp, 'TF_NEED_KAFKA', 'Apache Kafka Platform',
- 'with_kafka_support', True, 'kafka')
+ set_build_var(environ_cp, 'TF_NEED_IGNITE', 'Apache Ignite',
+ 'with_ignite_support', True, 'ignite')
set_build_var(environ_cp, 'TF_ENABLE_XLA', 'XLA JIT', 'with_xla_support',
- False, 'xla')
- set_build_var(environ_cp, 'TF_NEED_GDR', 'GDR', 'with_gdr_support', False,
- 'gdr')
- set_build_var(environ_cp, 'TF_NEED_VERBS', 'VERBS', 'with_verbs_support',
- False, 'verbs')
- set_build_var(environ_cp, 'TF_NEED_NGRAPH', 'nGraph', 'with_ngraph_support',
- False, 'ngraph')
+ True, 'xla')
set_action_env_var(environ_cp, 'TF_NEED_OPENCL_SYCL', 'OpenCL SYCL', False)
if environ_cp.get('TF_NEED_OPENCL_SYCL') == '1':
@@ -1602,6 +1596,13 @@
else:
set_trisycl_include_dir(environ_cp)
+ set_action_env_var(environ_cp, 'TF_NEED_ROCM', 'ROCm', False)
+ if (environ_cp.get('TF_NEED_ROCM') == '1' and
+ 'LD_LIBRARY_PATH' in environ_cp and
+ environ_cp.get('LD_LIBRARY_PATH') != '1'):
+ write_action_env_to_bazelrc('LD_LIBRARY_PATH',
+ environ_cp.get('LD_LIBRARY_PATH'))
+
set_action_env_var(environ_cp, 'TF_NEED_CUDA', 'CUDA', False)
if (environ_cp.get('TF_NEED_CUDA') == '1' and
'TF_CUDA_CONFIG_REPO' not in environ_cp):
@@ -1642,12 +1643,24 @@
write_to_bazelrc('build --config=download_clang')
write_to_bazelrc('test --config=download_clang')
+ # SYCL / ROCm / CUDA are mutually exclusive.
+ # At most 1 GPU platform can be configured.
+ gpu_platform_count = 0
+ if environ_cp.get('TF_NEED_OPENCL_SYCL') == '1':
+ gpu_platform_count += 1
+ if environ_cp.get('TF_NEED_ROCM') == '1':
+ gpu_platform_count += 1
+ if environ_cp.get('TF_NEED_CUDA') == '1':
+ gpu_platform_count += 1
+ if gpu_platform_count >= 2:
+ raise UserInputError('SYCL / CUDA / ROCm are mututally exclusive. '
+ 'At most 1 GPU platform can be configured.')
+
set_build_var(environ_cp, 'TF_NEED_MPI', 'MPI', 'with_mpi_support', False)
if environ_cp.get('TF_NEED_MPI') == '1':
set_mpi_home(environ_cp)
set_other_mpi_vars(environ_cp)
- set_grpc_build_flags()
set_cc_opt_flags(environ_cp)
set_system_libs_flag(environ_cp)
if is_windows():
@@ -1672,6 +1685,9 @@
'more details.')
config_info_line('mkl', 'Build with MKL support.')
config_info_line('monolithic', 'Config for mostly static monolithic build.')
+ config_info_line('gdr', 'Build with GDR support.')
+ config_info_line('verbs', 'Build with libverbs support.')
+ config_info_line('ngraph', 'Build with Intel nGraph support.')
if __name__ == '__main__':
diff --git a/tensorflow/BUILD b/tensorflow/BUILD
index 3610eea..9b62a50 100644
--- a/tensorflow/BUILD
+++ b/tensorflow/BUILD
@@ -203,21 +203,6 @@
visibility = ["//visibility:public"],
)
-# TODO(jhseu): Enable on other platforms other than Linux.
-config_setting(
- name = "with_jemalloc_linux_x86_64",
- define_values = {"with_jemalloc": "true"},
- values = {"cpu": "k8"},
- visibility = ["//visibility:public"],
-)
-
-config_setting(
- name = "with_jemalloc_linux_ppc64le",
- define_values = {"with_jemalloc": "true"},
- values = {"cpu": "ppc"},
- visibility = ["//visibility:public"],
-)
-
config_setting(
name = "with_default_optimizations",
define_values = {"with_default_optimizations": "true"},
@@ -225,56 +210,8 @@
)
config_setting(
- name = "with_gcp_support",
- define_values = {"with_gcp_support": "true"},
- visibility = ["//visibility:public"],
-)
-
-config_setting(
- name = "with_hdfs_support",
- define_values = {"with_hdfs_support": "true"},
- visibility = ["//visibility:public"],
-)
-
-config_setting(
- name = "with_aws_support",
- define_values = {"with_aws_support": "true"},
- visibility = ["//visibility:public"],
-)
-
-config_setting(
- name = "with_kafka_support",
- define_values = {"with_kafka_support": "true"},
- visibility = ["//visibility:public"],
-)
-
-# Crosses between platforms and file system libraries not supported on those
-# platforms due to limitations in nested select() statements.
-config_setting(
- name = "with_gcp_support_windows_override",
- define_values = {"with_gcp_support": "true"},
- values = {"cpu": "x64_windows"},
- visibility = ["//visibility:public"],
-)
-
-config_setting(
- name = "with_hdfs_support_windows_override",
- define_values = {"with_hdfs_support": "true"},
- values = {"cpu": "x64_windows"},
- visibility = ["//visibility:public"],
-)
-
-config_setting(
- name = "with_aws_support_windows_override",
- define_values = {"with_aws_support": "true"},
- values = {"cpu": "x64_windows"},
- visibility = ["//visibility:public"],
-)
-
-config_setting(
- name = "with_kafka_support_windows_override",
- define_values = {"with_kafka_support": "true"},
- values = {"cpu": "x64_windows"},
+ name = "with_ignite_support",
+ define_values = {"with_ignite_support": "true"},
visibility = ["//visibility:public"],
)
@@ -286,48 +223,6 @@
)
config_setting(
- name = "with_gcp_support_android_override",
- define_values = {"with_gcp_support": "true"},
- values = {"crosstool_top": "//external:android/crosstool"},
- visibility = ["//visibility:public"],
-)
-
-config_setting(
- name = "with_hdfs_support_android_override",
- define_values = {"with_hdfs_support": "true"},
- values = {"crosstool_top": "//external:android/crosstool"},
- visibility = ["//visibility:public"],
-)
-
-config_setting(
- name = "with_aws_support_android_override",
- define_values = {"with_aws_support": "true"},
- values = {"crosstool_top": "//external:android/crosstool"},
- visibility = ["//visibility:public"],
-)
-
-config_setting(
- name = "with_gcp_support_ios_override",
- define_values = {"with_gcp_support": "true"},
- values = {"crosstool_top": "//tools/osx/crosstool:crosstool"},
- visibility = ["//visibility:public"],
-)
-
-config_setting(
- name = "with_hdfs_support_ios_override",
- define_values = {"with_hdfs_support": "true"},
- values = {"crosstool_top": "//tools/osx/crosstool:crosstool"},
- visibility = ["//visibility:public"],
-)
-
-config_setting(
- name = "with_aws_support_ios_override",
- define_values = {"with_aws_support": "true"},
- values = {"crosstool_top": "//tools/osx/crosstool:crosstool"},
- visibility = ["//visibility:public"],
-)
-
-config_setting(
name = "with_xla_support",
define_values = {"with_xla_support": "true"},
visibility = ["//visibility:public"],
@@ -356,30 +251,6 @@
)
config_setting(
- name = "with_jemalloc_linux_x86_64_dynamic",
- define_values = {
- "with_jemalloc": "true",
- "framework_shared_object": "true",
- },
- values = {
- "cpu": "k8",
- },
- visibility = ["//visibility:public"],
-)
-
-config_setting(
- name = "with_jemalloc_linux_ppc64le_dynamic",
- define_values = {
- "with_jemalloc": "true",
- "framework_shared_object": "true",
- },
- values = {
- "cpu": "ppc",
- },
- visibility = ["//visibility:public"],
-)
-
-config_setting(
name = "using_cuda_clang",
define_values = {
"using_cuda_clang": "true",
diff --git a/tensorflow/c/BUILD b/tensorflow/c/BUILD
index 43c279b..17e2e29 100644
--- a/tensorflow/c/BUILD
+++ b/tensorflow/c/BUILD
@@ -246,6 +246,7 @@
":c_api_experimental",
":c_test_util",
"//tensorflow/core:lib",
+ "//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
],
diff --git a/tensorflow/c/c_api_experimental.cc b/tensorflow/c/c_api_experimental.cc
index 3bcc62c..d4b7813 100644
--- a/tensorflow/c/c_api_experimental.cc
+++ b/tensorflow/c/c_api_experimental.cc
@@ -23,6 +23,7 @@
#include "tensorflow/core/lib/strings/strcat.h"
#include "tensorflow/core/platform/platform.h"
#include "tensorflow/core/protobuf/config.pb.h"
+#include "tensorflow/core/protobuf/tensorflow_server.pb.h"
using tensorflow::FunctionDef;
using tensorflow::Node;
@@ -8508,6 +8509,20 @@
VLOG(1) << "Enqueuing is done.";
}
+TF_Buffer* TFE_GetServerDef(const char* text_proto, TF_Status* status) {
+ tensorflow::ServerDef server_def;
+ if (!tensorflow::protobuf::TextFormat::ParseFromString(text_proto,
+ &server_def)) {
+ status->status = tensorflow::errors::Internal(
+ "Invalid text proto for ServerDef: ", text_proto);
+ return nullptr;
+ }
+ status->status = tensorflow::Status();
+ TF_Buffer* ret = TF_NewBuffer();
+ TF_CHECK_OK(MessageToBuffer(server_def, ret));
+ return ret;
+}
+
TFE_Context* TFE_CreateContextFromSession(TF_Session* session,
TF_Status* status) {
auto* opts = TFE_NewContextOptions();
@@ -8723,35 +8738,7 @@
TF_DeleteStatus(status);
}
-TFE_TensorHandle* TFE_RunConstOp(TFE_Context* ctx) {
- // Intentionally LOG into INFO below for ease of debugging.
- VLOG(1) << "TFE_RunConstOp called";
-
- auto* status = TF_NewStatus();
- auto* op = TFE_NewOp(ctx, "Const", status);
- CheckOk(status);
- TFE_OpSetAttrType(op, "dtype", TF_FLOAT);
-
- auto* tensor =
- TF_AllocateTensor(TF_FLOAT, /*shape.data()*/ nullptr, /*shape.size()*/ 0,
- TF_DataTypeSize(TF_FLOAT) * 1);
- auto* ptr = reinterpret_cast<char*>(TF_TensorData(tensor));
- *reinterpret_cast<float*>(ptr) = 17.0;
-
- TFE_OpSetAttrTensor(op, "value", tensor, status);
- CheckOk(status);
- TF_DeleteTensor(tensor);
- VLOG(1) << "New op created";
-
- TFE_TensorHandle* retval;
- int num_retvals = 1;
- TFE_Execute(op, &retval, &num_retvals, status);
- CheckOk(status);
- CHECK_EQ(num_retvals, 1);
- VLOG(1) << "Op executed";
-
- TFE_DeleteOp(op);
- TF_DeleteStatus(status);
-
- return retval;
+TF_CAPI_EXPORT extern void TF_MakeInternalErrorStatus(TF_Status* status,
+ const char* errMsg) {
+ status->status = tensorflow::errors::Internal(errMsg);
}
diff --git a/tensorflow/c/c_api_experimental.h b/tensorflow/c/c_api_experimental.h
index a3ca847..d98d532 100644
--- a/tensorflow/c/c_api_experimental.h
+++ b/tensorflow/c/c_api_experimental.h
@@ -131,6 +131,8 @@
int tensor_id,
TF_Tensor* tensor,
TF_Status* status);
+// Create a serialized tensorflow.ServerDef proto.
+TF_Buffer* TFE_GetServerDef(const char* text_proto, TF_Status* status);
// TODO: remove this API in favor of the next one.
TF_CAPI_EXPORT extern TFE_Context* TFE_NewContextFromSession(
@@ -178,10 +180,8 @@
TF_CAPI_EXPORT extern void TFE_TensorHandlePrintDebugString(
TFE_TensorHandle* handle);
-// Returns a const scalar tensor.
-// Caller owns both the input and the output tensor handles.
-// TODO: Remove this API with hard-coded tensor computation.
-TF_CAPI_EXPORT extern TFE_TensorHandle* TFE_RunConstOp(TFE_Context* ctx);
+TF_CAPI_EXPORT extern void TF_MakeInternalErrorStatus(TF_Status* status,
+ const char* errMsg);
#ifdef __cplusplus
} /* end extern "C" */
diff --git a/tensorflow/c/c_api_experimental_test.cc b/tensorflow/c/c_api_experimental_test.cc
index 30fcfd4..c6effd3 100644
--- a/tensorflow/c/c_api_experimental_test.cc
+++ b/tensorflow/c/c_api_experimental_test.cc
@@ -16,8 +16,10 @@
#include "tensorflow/c/c_api_experimental.h"
#include "tensorflow/c/c_test_util.h"
#include "tensorflow/core/lib/io/path.h"
+#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/protobuf/tensorflow_server.pb.h"
namespace tensorflow {
namespace {
@@ -116,5 +118,49 @@
TF_DeleteStatus(s);
}
+TEST(CAPI_EXPERIMENTAL, GetServerDefTest) {
+ const string expected_text_proto(R"(cluster {
+ job {
+ name: "worker"
+ tasks {
+ key: 0
+ value: "tpuserver:0"
+ }
+ tasks {
+ key: 1
+ value: "localhost:1"
+ }
+ }
+}
+job_name: "worker"
+task_index: 1
+protocol: "grpc"
+)");
+
+ TF_Status* status = TF_NewStatus();
+ TF_Buffer* result = TFE_GetServerDef(expected_text_proto.c_str(), status);
+ EXPECT_EQ(TF_GetCode(status), TF_OK);
+
+ ServerDef actual;
+ ASSERT_TRUE(actual.ParseFromArray(result->data, result->length));
+ string actual_text_proto;
+ tensorflow::protobuf::TextFormat::PrintToString(actual, &actual_text_proto);
+ EXPECT_EQ(expected_text_proto, actual_text_proto);
+
+ const string malformed_text_proto(R"(cluster {
+ job {
+ name: "worker")");
+ TF_Buffer* null_result =
+ TFE_GetServerDef(malformed_text_proto.c_str(), status);
+ EXPECT_NE(TF_GetCode(status), TF_OK);
+ EXPECT_TRUE(tensorflow::str_util::StrContains(
+ TF_Message(status), "Invalid text proto for ServerDef"));
+ EXPECT_EQ(null_result, nullptr);
+
+ // Cleanup
+ TF_DeleteBuffer(result);
+ TF_DeleteStatus(status);
+}
+
} // namespace
} // namespace tensorflow
diff --git a/tensorflow/c/eager/c_api.cc b/tensorflow/c/eager/c_api.cc
index 0bf3d95..3554ec0 100755
--- a/tensorflow/c/eager/c_api.cc
+++ b/tensorflow/c/eager/c_api.cc
@@ -578,6 +578,14 @@
op->operation.MutableAttrs()->Set(attr_name, attr_value);
}
+void TFE_OpSetAttrFunctionName(TFE_Op* op, const char* attr_name,
+ const char* data, size_t length) {
+ tensorflow::AttrValue attr_value;
+ tensorflow::NameAttrList* func = attr_value.mutable_func();
+ func->set_name(data, length);
+ op->operation.MutableAttrs()->Set(attr_name, attr_value);
+}
+
void TFE_OpSetAttrTensor(TFE_Op* op, const char* attr_name, TF_Tensor* tensor,
TF_Status* status) {
tensorflow::Tensor t;
diff --git a/tensorflow/c/eager/c_api.h b/tensorflow/c/eager/c_api.h
index 6323f8a..b2454d8 100755
--- a/tensorflow/c/eager/c_api.h
+++ b/tensorflow/c/eager/c_api.h
@@ -313,6 +313,9 @@
const char* attr_name,
const TFE_Op* value);
+TF_CAPI_EXPORT void TFE_OpSetAttrFunctionName(TFE_Op* op, const char* attr_name,
+ const char* data, size_t length);
+
TF_CAPI_EXPORT extern void TFE_OpSetAttrTensor(TFE_Op* op,
const char* attr_name,
TF_Tensor* tensor,
diff --git a/tensorflow/c/eager/tape.h b/tensorflow/c/eager/tape.h
index 41b5b8f..5ba55a2 100644
--- a/tensorflow/c/eager/tape.h
+++ b/tensorflow/c/eager/tape.h
@@ -130,7 +130,7 @@
const string& op_type, std::vector<TapeTensor>& output_tensors,
gtl::ArraySlice<int64> input_tensor_id,
gtl::ArraySlice<tensorflow::DataType> input_dtypes,
- BackwardFunction* backward_function,
+ const std::function<BackwardFunction*()>& backward_function_getter,
const std::function<void(BackwardFunction*)>& backward_function_deleter);
void DeleteTrace(int64 tensor_id);
@@ -206,10 +206,9 @@
const string& op_type, std::vector<TapeTensor>& output_tensors,
gtl::ArraySlice<int64> input_tensor_id,
gtl::ArraySlice<tensorflow::DataType> input_dtypes,
- BackwardFunction* backward_function,
+ const std::function<BackwardFunction*()>& backward_function_getter,
const std::function<void(BackwardFunction*)>& backward_function_deleter) {
if (!ShouldRecord(input_tensor_id, input_dtypes)) {
- backward_function_deleter(backward_function);
return;
}
std::vector<int64> ids;
@@ -229,7 +228,7 @@
tensors.push_back(o);
}
op_tape_[op_id] = OpTapeEntry<BackwardFunction, TapeTensor>{
- op_type, std::move(tensors), ids, backward_function,
+ op_type, std::move(tensors), std::move(ids), backward_function_getter(),
backward_function_deleter};
}
diff --git a/tensorflow/cc/framework/cc_op_gen.cc b/tensorflow/cc/framework/cc_op_gen.cc
index a32d1b1..3959337 100644
--- a/tensorflow/cc/framework/cc_op_gen.cc
+++ b/tensorflow/cc/framework/cc_op_gen.cc
@@ -853,11 +853,7 @@
}
}
- strings::StrAppend(&class_decl, "\n");
-
- if (output_types.empty()) {
- strings::StrAppend(&class_decl, " Operation operation;\n");
- }
+ strings::StrAppend(&class_decl, "\n Operation operation;\n");
for (int i = 0; i < output_types.size(); ++i) {
strings::StrAppend(&class_decl, " ", output_types[i], " ", output_names[i],
";\n");
@@ -878,9 +874,11 @@
string return_on_error =
strings::StrCat("if (!", scope_str, ".ok()) return;");
+ strings::StrAppend(out, " this->operation = Operation(ret);\n");
+
// No outputs.
if (graph_op_def.output_arg_size() == 0) {
- strings::StrAppend(out, " this->operation = Operation(ret);\n return;\n");
+ strings::StrAppend(out, " return;\n");
return;
}
if (graph_op_def.output_arg_size() == 1) {
diff --git a/tensorflow/cc/framework/scope.cc b/tensorflow/cc/framework/scope.cc
index 7f6ac4c..6abc9e2 100644
--- a/tensorflow/cc/framework/scope.cc
+++ b/tensorflow/cc/framework/scope.cc
@@ -62,7 +62,7 @@
refiner_(refiner),
scope_used_(nullptr),
colocation_constraints_(),
- disable_shape_inference_(false) {}
+ disable_shape_inference_(refiner_ == nullptr) {}
Scope Scope::NewRootScope() {
Graph* graph = new Graph(OpRegistry::Global());
@@ -94,6 +94,7 @@
exit_on_error_(other.impl()->exit_on_error_),
kernel_label_(other.impl()->kernel_label_),
device_(other.impl()->device_),
+ assigned_device_(other.impl()->assigned_device_),
colocation_constraints_(other.impl()->colocation_constraints_),
disable_shape_inference_(other.impl()->disable_shape_inference_) {}
@@ -110,6 +111,7 @@
exit_on_error_(other.impl()->exit_on_error_),
kernel_label_(other.impl()->kernel_label_),
device_(other.impl()->device_),
+ assigned_device_(other.impl()->assigned_device_),
colocation_constraints_(other.impl()->colocation_constraints_),
disable_shape_inference_(other.impl()->disable_shape_inference_) {}
@@ -132,6 +134,7 @@
exit_on_error_(other.impl()->exit_on_error_),
kernel_label_(other.impl()->kernel_label_),
device_(other.impl()->device_),
+ assigned_device_(other.impl()->assigned_device_),
colocation_constraints_(other.impl()->colocation_constraints_),
disable_shape_inference_(other.impl()->disable_shape_inference_) {}
@@ -163,6 +166,7 @@
exit_on_error_(other.impl()->exit_on_error_),
kernel_label_(other.impl()->kernel_label_),
device_(other.impl()->device_),
+ assigned_device_(other.impl()->assigned_device_),
colocation_constraints_(other.impl()->colocation_constraints_),
disable_shape_inference_(other.impl()->disable_shape_inference_) {}
@@ -178,6 +182,7 @@
exit_on_error_(true),
kernel_label_(other.impl()->kernel_label_),
device_(other.impl()->device_),
+ assigned_device_(other.impl()->assigned_device_),
colocation_constraints_(other.impl()->colocation_constraints_),
disable_shape_inference_(other.impl()->disable_shape_inference_) {}
@@ -194,6 +199,7 @@
exit_on_error_(other.impl()->exit_on_error_),
kernel_label_(kernel_label),
device_(other.impl()->device_),
+ assigned_device_(other.impl()->assigned_device_),
colocation_constraints_(other.impl()->colocation_constraints_),
disable_shape_inference_(other.impl()->disable_shape_inference_) {}
@@ -210,12 +216,30 @@
exit_on_error_(other.impl()->exit_on_error_),
kernel_label_(other.impl()->kernel_label_),
device_(other.impl()->device_),
+ assigned_device_(other.impl()->assigned_device_),
colocation_constraints_(
clear_colocations
? std::unordered_set<string>()
: other.impl()->GetColocationConstraints(colocate_with_op)),
disable_shape_inference_(other.impl()->disable_shape_inference_) {}
+Scope::Impl::Impl(const Scope& other, Tags::AssignedDevice,
+ const string& assigned_device)
+ : graph_(other.impl()->graph_),
+ status_(other.impl()->status_),
+ name_map_(other.impl()->name_map_),
+ refiner_(other.impl()->refiner_),
+ scope_used_(other.impl()->scope_used_),
+ control_deps_(other.impl()->control_deps_),
+ name_(other.impl()->name_),
+ op_name_(other.impl()->op_name_),
+ exit_on_error_(other.impl()->exit_on_error_),
+ kernel_label_(other.impl()->kernel_label_),
+ device_(other.impl()->device_),
+ assigned_device_(assigned_device),
+ colocation_constraints_(other.impl()->colocation_constraints_),
+ disable_shape_inference_(other.impl()->disable_shape_inference_) {}
+
std::unordered_set<string> Scope::Impl::GetColocationConstraints(
const Operation& colocate_with_op) const {
std::unordered_set<string> current_constraints(colocation_constraints_);
@@ -299,6 +323,9 @@
if (!impl()->device_.empty()) {
builder->Device(impl()->device_);
}
+ if (!impl()->assigned_device_.empty()) {
+ builder->AssignedDevice(impl()->assigned_device_);
+ }
}
string Scope::Impl::GetUniqueName(const string& prefix,
@@ -394,6 +421,10 @@
return Scope(new Impl(*this, Impl::Tags::Device(), device));
}
+Scope Scope::WithAssignedDevice(const string& assigned_device) const {
+ return Scope(new Impl(*this, Impl::Tags::AssignedDevice(), assigned_device));
+}
+
Scope Scope::ColocateWith(const Operation& op) const {
return Scope(new Impl(*this, Impl::Tags::Colocate(), op,
/* clear_colocations */ false));
diff --git a/tensorflow/cc/framework/scope.h b/tensorflow/cc/framework/scope.h
index 30c32bd..e307d89 100644
--- a/tensorflow/cc/framework/scope.h
+++ b/tensorflow/cc/framework/scope.h
@@ -133,6 +133,10 @@
/// the device field set to 'device'.
Scope WithDevice(const string& device) const;
+ /// Returns a new scope. All ops created within the returned scope will have
+ /// their assigned device set to `assigned_device`.
+ Scope WithAssignedDevice(const string& assigned_device) const;
+
/// Return a new scope. All ops created within the returned scope will be
/// co-located on the device where op is placed.
/// NOTE: This function is intended to be use internal libraries only for
diff --git a/tensorflow/cc/framework/scope_internal.h b/tensorflow/cc/framework/scope_internal.h
index 58adaef..514e02e 100644
--- a/tensorflow/cc/framework/scope_internal.h
+++ b/tensorflow/cc/framework/scope_internal.h
@@ -26,6 +26,8 @@
// graph, status, name_map, and refiner.
// This is intended to enable the C API (which are used by other language
// bindings) to create a Scope and access C++ functionality (i.e. gradients).
+//
+// Shape inference is disabled if `refiner` is nullptr.
Scope NewInternalScope(Graph* graph, Status* status, ShapeRefiner* refiner);
class Scope::Impl {
@@ -58,6 +60,7 @@
enum class ExitOnError;
enum class KernelLabel;
enum class Colocate;
+ enum class AssignedDevice;
};
Impl(Graph* graph, Status* status, NameMap* name_map, ShapeRefiner* refiner,
@@ -74,6 +77,7 @@
Impl(const Scope& other, Tags::KernelLabel, const string& kernel_label);
Impl(const Scope& other, Tags::Colocate, const Operation& colocate_with_op,
bool clear_colocations);
+ Impl(const Scope& other, Tags::AssignedDevice, const string& assigned_device);
std::unordered_set<string> GetColocationConstraints(
const Operation& colocate_with_op) const;
@@ -107,6 +111,7 @@
const bool exit_on_error_ = false;
const string kernel_label_ = "";
const string device_ = "";
+ const string assigned_device_ = "";
const std::unordered_set<string> colocation_constraints_;
// If true, Scope::DoShapeInference() always returns Status:OK().
diff --git a/tensorflow/compiler/aot/tests/make_test_graphs.py b/tensorflow/compiler/aot/tests/make_test_graphs.py
index de135d7..64b861a 100644
--- a/tensorflow/compiler/aot/tests/make_test_graphs.py
+++ b/tensorflow/compiler/aot/tests/make_test_graphs.py
@@ -47,7 +47,7 @@
def tfadd_with_ckpt(out_dir):
x = array_ops.placeholder(dtypes.int32, name='x_hold')
- y = variables.Variable(constant_op.constant([0]), name='y_saved')
+ y = variables.VariableV1(constant_op.constant([0]), name='y_saved')
math_ops.add(x, y, name='x_y_sum')
init_op = variables.initialize_all_variables()
@@ -62,7 +62,7 @@
def tfadd_with_ckpt_saver(out_dir):
x = array_ops.placeholder(dtypes.int32, name='x_hold')
- y = variables.Variable(constant_op.constant([0]), name='y_saved')
+ y = variables.VariableV1(constant_op.constant([0]), name='y_saved')
math_ops.add(x, y, name='x_y_sum')
init_op = variables.initialize_all_variables()
diff --git a/tensorflow/compiler/jit/BUILD b/tensorflow/compiler/jit/BUILD
index 4e18472..f202709 100644
--- a/tensorflow/compiler/jit/BUILD
+++ b/tensorflow/compiler/jit/BUILD
@@ -258,6 +258,7 @@
"//tensorflow/core:lib_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/kernels:variable_ops",
+ "@com_google_absl//absl/container:flat_hash_map",
],
)
@@ -323,6 +324,7 @@
"//tensorflow/core:graph",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:optional",
@@ -383,12 +385,16 @@
":shape_inference_helpers",
":union_find",
":xla_cluster_util",
+ "//tensorflow/cc:cc_ops",
+ "//tensorflow/cc:ops",
+ "//tensorflow/cc:scope_internal",
"//tensorflow/compiler/jit/graphcycles",
"//tensorflow/compiler/jit/legacy_flags:mark_for_compilation_pass_flags",
"//tensorflow/compiler/jit/ops:xla_ops",
"//tensorflow/compiler/tf2xla:dump_graph",
"//tensorflow/compiler/tf2xla:resource_operation_table",
"//tensorflow/compiler/tf2xla:xla_compiler",
+ "//tensorflow/compiler/tf2xla/cc:xla_jit_ops",
"//tensorflow/compiler/xla:status_macros",
"//tensorflow/compiler/xla:util",
"//tensorflow/core:core_cpu",
@@ -400,6 +406,7 @@
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/kernels:bounds_check",
"@com_google_absl//absl/algorithm:container",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
],
@@ -471,6 +478,7 @@
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core:testlib",
+ "@com_google_absl//absl/container:flat_hash_map",
],
)
@@ -478,6 +486,7 @@
name = "compilation_passes_test",
size = "small",
srcs = [
+ "build_xla_ops_pass_test.cc",
"encapsulate_subgraphs_pass_test.cc",
"encapsulate_xla_computations_pass_test.cc",
"mark_for_compilation_pass_test.cc",
@@ -486,6 +495,7 @@
deps = [
":common",
":compilation_passes",
+ ":node_matchers",
":xla_cluster_util",
":xla_gpu_device",
"//tensorflow/cc:cc_ops",
@@ -507,6 +517,7 @@
"//tensorflow/core:test_main",
"//tensorflow/core:testlib",
"//tensorflow/core/grappler/optimizers/data:graph_utils",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
],
diff --git a/tensorflow/compiler/jit/build_xla_ops_pass.cc b/tensorflow/compiler/jit/build_xla_ops_pass.cc
index 13a518d..5974696 100644
--- a/tensorflow/compiler/jit/build_xla_ops_pass.cc
+++ b/tensorflow/compiler/jit/build_xla_ops_pass.cc
@@ -14,8 +14,12 @@
==============================================================================*/
#include "tensorflow/compiler/jit/build_xla_ops_pass.h"
+#include "absl/algorithm/container.h"
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/scope_internal.h"
#include "tensorflow/compiler/jit/defs.h"
#include "tensorflow/compiler/jit/encapsulate_subgraphs_pass.h"
+#include "tensorflow/compiler/tf2xla/cc/ops/xla_jit_ops.h"
#include "tensorflow/compiler/tf2xla/dump_graph.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "tensorflow/core/common_runtime/function.h"
@@ -31,139 +35,108 @@
#include "tensorflow/core/public/version.h"
namespace tensorflow {
+namespace {
+void MoveOutgoingEdges(Graph* g, Node* old_node, Node* new_node) {
+ std::vector<const Edge*> out_edges(old_node->out_edges().begin(),
+ old_node->out_edges().end());
+ for (const Edge* edge : out_edges) {
+ // TODO(sanjoy): This does not update NodeDef inputs. To be able to update
+ // NodeDef inputs we first need to fix encapsulate_subgraphs_pass to fix up
+ // the NodeDef inputs to the function call nodes.
+ g->AddEdge(new_node, edge->src_output(), edge->dst(), edge->dst_input());
+ g->RemoveEdge(edge);
+ }
+}
-static Status BuildXlaCompileNode(
- const string& nodename, const string& function_name,
- const AttrValueMap& function_attr, const string& device_name,
- const DataTypeVector& constant_dtypes, int num_resources,
- const DataTypeVector& arg_dtypes, Graph* graph, Node** node) {
- NodeDef def;
- def.set_name(graph->NewName(nodename));
- def.set_op("_XlaCompile");
- def.set_device(device_name);
- AddNodeAttr("Tconstants", constant_dtypes, &def);
- AddNodeAttr("Targs", arg_dtypes, &def);
- AddNodeAttr("Nresources", num_resources, &def);
+struct XlaClusterInfo {
+ std::vector<Output> constant_inputs;
+ std::vector<Output> non_constant_inputs;
+ std::vector<Output> resource_inputs;
NameAttrList function;
- function.set_name(function_name);
- *function.mutable_attr() = function_attr;
- AddNodeAttr("function", function, &def);
+};
- Status status;
- *node = graph->AddNode(def, &status);
- return status;
+Output IncomingEdgeAsOutput(const Edge* e) {
+ return Output(e->src(), e->src_output());
}
-static Status BuildXlaRunNode(const string& nodename, const string& device_name,
- const DataTypeVector& arg_dtypes,
- const DataTypeVector& result_dtypes, Graph* graph,
- Node** node) {
- NodeDef def;
- def.set_name(graph->NewName(nodename));
- def.set_op("_XlaRun");
- def.set_device(device_name);
- AddNodeAttr("Targs", arg_dtypes, &def);
- AddNodeAttr("Tresults", result_dtypes, &def);
-
- Status status;
- *node = graph->AddNode(def, &status);
- return status;
-}
-
-static Status GetXlaAttrs(Node* node, int* num_constant_args,
- int* num_resource_args, DataTypeVector* const_dtypes,
- DataTypeVector* arg_dtypes) {
+Status GetXlaClusterInfo(Node* n, XlaClusterInfo* result) {
+ int num_constant_inputs, num_resource_inputs;
TF_RETURN_IF_ERROR(
- GetNodeAttr(node->attrs(), kXlaNumConstantArgsAttr, num_constant_args));
+ GetNodeAttr(n->attrs(), kXlaNumConstantArgsAttr, &num_constant_inputs));
TF_RETURN_IF_ERROR(
- GetNodeAttr(node->attrs(), kXlaNumResourceArgsAttr, num_resource_args));
+ GetNodeAttr(n->attrs(), kXlaNumResourceArgsAttr, &num_resource_inputs));
- if (*num_constant_args < 0 || *num_resource_args < 0 ||
- *num_constant_args + *num_resource_args > node->num_inputs()) {
+ if (num_constant_inputs < 0 || num_resource_inputs < 0 ||
+ num_constant_inputs + num_resource_inputs > n->num_inputs()) {
return errors::InvalidArgument(
"Invalid number of constant/resource arguments to XLA kernel.");
}
- const int num_nonconst_args =
- node->num_inputs() - *num_constant_args - *num_resource_args;
+ int num_non_constant_inputs =
+ n->num_inputs() - num_constant_inputs - num_resource_inputs;
- const DataTypeVector& input_types = node->input_types();
- std::copy(input_types.begin(), input_types.begin() + *num_constant_args,
- std::back_inserter(*const_dtypes));
- std::copy(input_types.begin() + *num_constant_args,
- input_types.begin() + *num_constant_args + num_nonconst_args,
- std::back_inserter(*arg_dtypes));
+ std::vector<const Edge*> input_edges_vector;
+ TF_RETURN_IF_ERROR(n->input_edges(&input_edges_vector));
+ absl::Span<const Edge*> input_edges(input_edges_vector);
+
+ absl::c_transform(input_edges.subspan(0, num_constant_inputs),
+ std::back_inserter(result->constant_inputs),
+ IncomingEdgeAsOutput);
+
+ absl::c_transform(
+ input_edges.subspan(num_constant_inputs, num_non_constant_inputs),
+ std::back_inserter(result->non_constant_inputs), IncomingEdgeAsOutput);
+
+ absl::c_transform(
+ input_edges.subspan(num_constant_inputs + num_non_constant_inputs,
+ num_resource_inputs),
+ std::back_inserter(result->resource_inputs), IncomingEdgeAsOutput);
+
+ result->function.set_name(n->type_string());
+ *result->function.mutable_attr() = n->def().attr();
return Status::OK();
}
-static void CopyIncomingEdges(Graph* g, Node* old_node, Node* new_node,
- int prefix_to_ignore) {
- for (const Edge* edge : old_node->in_edges()) {
- if (edge->IsControlEdge()) {
- g->AddControlEdge(edge->src(), new_node);
- } else if (edge->dst_input() >= prefix_to_ignore) {
- g->AddEdge(edge->src(), edge->src_output(), new_node,
- edge->dst_input() - prefix_to_ignore);
+Status CopyIncomingControlEdges(Graph* g, Node* from, Node* to) {
+ for (const Edge* e : from->in_edges()) {
+ if (e->IsControlEdge()) {
+ g->AddControlEdge(e->src(), to);
}
}
+
+ return Status::OK();
}
-static void MoveOutgoingEdges(Graph* g, Node* old_node, Node* new_node) {
- std::vector<const Edge*> out_edges(old_node->out_edges().begin(),
- old_node->out_edges().end());
- for (const Edge* edge : out_edges) {
- Node* dst = edge->dst();
- int src_output = edge->src_output();
- int dst_input = edge->dst_input();
- g->RemoveEdge(edge);
+Status ReplaceNodeWithXlaCompileAndXlaRun(Graph* g, Node* n) {
+ Status status;
+ Scope root = NewInternalScope(g, &status, /*refiner=*/nullptr)
+ .NewSubScope(n->name())
+ .WithDevice(n->requested_device())
+ .WithAssignedDevice(n->assigned_device_name());
- if (edge->IsControlEdge()) {
- g->AddControlEdge(new_node, dst);
- } else {
- g->AddEdge(new_node, src_output, dst, dst_input);
- }
- }
-}
+ XlaClusterInfo cluster_info;
+ TF_RETURN_IF_ERROR(GetXlaClusterInfo(n, &cluster_info));
-static Status ReplaceNodeWithXlaCompileAndRun(Graph* g, Node* n) {
- int num_constant_args, num_resource_args;
- DataTypeVector const_dtypes;
- DataTypeVector arg_dtypes;
+ ops::_XlaCompile xla_compile(root.WithOpName("xla_compile"),
+ /*constants=*/cluster_info.constant_inputs,
+ /*args=*/cluster_info.non_constant_inputs,
+ /*resources=*/cluster_info.resource_inputs,
+ cluster_info.function);
+ TF_RETURN_IF_ERROR(
+ CopyIncomingControlEdges(g, /*from=*/n, /*to=*/xla_compile.key.node()));
- TF_RETURN_IF_ERROR(GetXlaAttrs(n, &num_constant_args, &num_resource_args,
- &const_dtypes, &arg_dtypes));
+ std::vector<Output> xla_run_args = cluster_info.non_constant_inputs;
+ absl::c_copy(cluster_info.resource_inputs, std::back_inserter(xla_run_args));
+ ops::_XlaRun xla_run(root.WithOpName("xla_run"), xla_run_args,
+ xla_compile.key, n->output_types());
- Node *compile_node, *run_node;
-
- TF_RETURN_IF_ERROR(BuildXlaCompileNode(
- n->name(), n->type_string(), n->def().attr(), n->requested_device(),
- const_dtypes, num_resource_args, arg_dtypes, g, &compile_node));
-
- DataTypeVector arg_dtypes_with_resources = arg_dtypes;
- for (int i = 0; i < num_resource_args; i++) {
- arg_dtypes_with_resources.push_back(DT_RESOURCE);
- }
-
- TF_RETURN_IF_ERROR(BuildXlaRunNode(n->name(), n->requested_device(),
- arg_dtypes_with_resources,
- n->output_types(), g, &run_node));
-
- compile_node->set_assigned_device_name(n->assigned_device_name());
- run_node->set_assigned_device_name(n->assigned_device_name());
-
- CopyIncomingEdges(g, /*old_node=*/n, /*new_node=*/compile_node,
- /*prefix_to_ignore=*/0);
- CopyIncomingEdges(g, /*old_node=*/n, /*new_node=*/run_node,
- /*prefix_to_ignore=*/num_constant_args);
-
- // The compilation_key output.
- g->AddEdge(compile_node, 0, run_node, n->num_inputs() - num_constant_args);
-
- MoveOutgoingEdges(g, /*old_node=*/n, /*new_node=*/run_node);
+ MoveOutgoingEdges(g, /*old_node=*/n,
+ /*new_node=*/xla_run.operation.node());
g->RemoveNode(n);
return Status::OK();
}
+} // namespace
Status BuildXlaOpsPass::Run(const GraphOptimizationPassOptions& options) {
Graph* graph = options.graph->get();
@@ -177,7 +150,7 @@
// Only compile nodes that are marked for compilation by the
// compilation-marking pass (via 'attr_name').
if (IsXlaCompiledKernel(*n)) {
- TF_RETURN_IF_ERROR(ReplaceNodeWithXlaCompileAndRun(graph, n));
+ TF_RETURN_IF_ERROR(ReplaceNodeWithXlaCompileAndXlaRun(graph, n));
}
}
diff --git a/tensorflow/compiler/jit/build_xla_ops_pass_test.cc b/tensorflow/compiler/jit/build_xla_ops_pass_test.cc
new file mode 100644
index 0000000..9d56db7
--- /dev/null
+++ b/tensorflow/compiler/jit/build_xla_ops_pass_test.cc
@@ -0,0 +1,138 @@
+/* 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/jit/build_xla_ops_pass.h"
+
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/ops/array_ops.h"
+#include "tensorflow/cc/ops/resource_variable_ops.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/compiler/jit/defs.h"
+#include "tensorflow/compiler/jit/encapsulate_subgraphs_pass.h"
+#include "tensorflow/compiler/jit/node_matchers.h"
+#include "tensorflow/core/graph/algorithm.h"
+#include "tensorflow/core/grappler/optimizers/data/graph_utils.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace {
+
+using ::tensorflow::testing::FindNodeByName;
+using ::tensorflow::testing::matchers::CtrlDeps;
+using ::tensorflow::testing::matchers::NodeWith;
+using ::tensorflow::testing::matchers::Op;
+
+Status BuildXlaOps(const Scope& s, std::unique_ptr<Graph>* result) {
+ auto graph = absl::make_unique<Graph>(OpRegistry::Global());
+ TF_RETURN_IF_ERROR(s.ToGraph(graph.get()));
+
+ // Assign all nodes to the CPU device.
+ static const char* kCpuDevice = "/job:localhost/replica:0/task:0/cpu:0";
+ for (Node* n : graph->nodes()) {
+ if (n->assigned_device_name().empty()) {
+ n->set_assigned_device_name(kCpuDevice);
+ }
+ }
+
+ GraphOptimizationPassOptions opt_options;
+ opt_options.graph = &graph;
+ BuildXlaOpsPass pass;
+ TF_RETURN_IF_ERROR(pass.Run(opt_options));
+ *result = std::move(graph);
+ return Status::OK();
+}
+
+Status MakeXlaCompiledKernel(Graph* graph, const string& callee_name,
+ const string& node_name, int num_constant_args,
+ int num_resource_args, Node** result) {
+ NodeDef call_node;
+ call_node.set_name(node_name);
+ call_node.set_op(callee_name);
+ AddNodeAttr(kXlaCompiledKernelAttr, true, &call_node);
+ AddNodeAttr(kXlaNumConstantArgsAttr, num_constant_args, &call_node);
+ AddNodeAttr(kXlaNumResourceArgsAttr, num_resource_args, &call_node);
+ Status s;
+ *result = graph->AddNode(call_node, &s);
+ return s;
+}
+
+Status MakeXlaCompiledKernel(Graph* graph, const string& callee_name,
+ const string& node_name, Node** result) {
+ return MakeXlaCompiledKernel(graph, callee_name, node_name,
+ /*num_constant_args=*/0, /*num_resource_args=*/0,
+ result);
+}
+
+Node* MakeWrite(const Scope& scope, const string& id) {
+ Output var_handle =
+ ops::VarHandleOp(scope.WithOpName("Var" + id), DT_FLOAT, TensorShape({}));
+ Output value_to_write =
+ ops::Const(scope.WithOpName("ValueToAssign" + id), 1.0f);
+ ops::AssignVariableOp assign_op(scope.WithOpName("Assignee" + id), var_handle,
+ value_to_write);
+ return assign_op.operation.node();
+}
+
+FunctionDefLibrary CreateFunctionDefLibWithConstFunction(const string& name) {
+ FunctionDefLibrary flib_def;
+ FunctionDef func = FunctionDefHelper::Create(
+ /*function_name=*/name, /*in_def=*/{}, /*out_def=*/{"out: float"},
+ /*attr_def*/
+ {}, /*node_def=*/{FunctionDefHelper::Const("one", 1.0f)},
+ /*ret_def=*/{{"out", "out:output:0"}});
+ *flib_def.add_function() = std::move(func);
+ return flib_def;
+}
+
+TEST(BuildXlaOps, ControlDepsPreserved) {
+ Scope root = Scope::NewRootScope().ExitOnError();
+
+ FunctionDefLibrary flib_def =
+ CreateFunctionDefLibWithConstFunction("cluster_0");
+ TF_ASSERT_OK(root.graph()->AddFunctionLibrary(flib_def));
+ Node* call;
+ TF_ASSERT_OK(MakeXlaCompiledKernel(root.graph(), "cluster_0", "C", &call));
+ Node* write_op = MakeWrite(root, "write");
+ root.graph()->AddControlEdge(call, write_op);
+
+ std::unique_ptr<Graph> graph;
+ TF_ASSERT_OK(BuildXlaOps(root, &graph));
+
+ Node* write_op_new = FindNodeByName(graph.get(), write_op->name());
+ ASSERT_NE(write_op_new, nullptr);
+ EXPECT_THAT(write_op_new, NodeWith(CtrlDeps(NodeWith(Op("_XlaRun")))));
+}
+
+TEST(BuildXlaOps, CleanFailureOnBogusAttr) {
+ Scope root = Scope::NewRootScope().ExitOnError();
+
+ FunctionDefLibrary flib_def =
+ CreateFunctionDefLibWithConstFunction("cluster_0");
+ TF_ASSERT_OK(root.graph()->AddFunctionLibrary(flib_def));
+ Node* call;
+ TF_ASSERT_OK(
+ MakeXlaCompiledKernel(root.graph(), "cluster_0", "C", 100, 100, &call));
+ Node* write_op = MakeWrite(root, "write");
+ root.graph()->AddControlEdge(call, write_op);
+
+ std::unique_ptr<Graph> graph;
+ Status failure_status = BuildXlaOps(root, &graph);
+ ASSERT_FALSE(failure_status.ok());
+ EXPECT_EQ(failure_status.code(), error::INVALID_ARGUMENT);
+}
+
+} // namespace
+} // namespace tensorflow
diff --git a/tensorflow/compiler/jit/deadness_analysis.cc b/tensorflow/compiler/jit/deadness_analysis.cc
index 9128b48..e63d4b7 100644
--- a/tensorflow/compiler/jit/deadness_analysis.cc
+++ b/tensorflow/compiler/jit/deadness_analysis.cc
@@ -14,6 +14,8 @@
==============================================================================*/
#include "tensorflow/compiler/jit/deadness_analysis.h"
+#include "absl/algorithm/container.h"
+#include "absl/container/flat_hash_map.h"
#include "absl/strings/str_join.h"
#include "tensorflow/compiler/jit/deadness_analysis_internal.h"
#include "tensorflow/core/graph/algorithm.h"
@@ -383,6 +385,8 @@
}
Predicate* MakeAndOrImpl(absl::Span<Predicate* const> operands, bool is_and);
+ Predicate* MakeInternedAndOr(std::vector<Predicate*> simplified_ops,
+ Predicate::Kind pred_kind);
// Predicate instances are interned, meaning that there is only a single
// instance of a Predicate object with a given content. This makes checking
@@ -417,23 +421,52 @@
}
};
- gtl::FlatMap<SignatureForAndOr, std::unique_ptr<Predicate>,
- HashSignatureForAndOr>
+ absl::flat_hash_map<SignatureForAndOr, std::unique_ptr<Predicate>,
+ HashSignatureForAndOr>
interned_and_or_instances_;
- gtl::FlatMap<SignatureForNot, std::unique_ptr<Predicate>>
+ absl::flat_hash_map<SignatureForNot, std::unique_ptr<Predicate>>
interned_not_instances_;
- gtl::FlatMap<SignatureForAndRec, std::unique_ptr<Predicate>>
+ absl::flat_hash_map<SignatureForAndRec, std::unique_ptr<Predicate>>
interned_and_rec_instances_;
- gtl::FlatMap<SignatureForSymbol, std::unique_ptr<Predicate>,
- HashSignatureForSymbol>
+ absl::flat_hash_map<SignatureForSymbol, std::unique_ptr<Predicate>,
+ HashSignatureForSymbol>
interned_symbol_instances_;
};
+Predicate* PredicateFactory::MakeInternedAndOr(
+ std::vector<Predicate*> simplified_ops, Predicate::Kind pred_kind) {
+ std::stable_sort(
+ simplified_ops.begin(), simplified_ops.end(),
+ [](Predicate* a, Predicate* b) { return a->hash() < b->hash(); });
+
+ auto it = interned_and_or_instances_.find({pred_kind, simplified_ops});
+ if (it != interned_and_or_instances_.end()) {
+ return it->second.get();
+ }
+
+ simplified_ops.shrink_to_fit();
+ // NB! Because we'll use a non-owning reference to simplified_ops in the
+ // key for interned_and_or_instances_ we need to be careful to std::move()
+ // it all the way through.
+ absl::Span<Predicate* const> operands_slice = simplified_ops;
+ std::unique_ptr<Predicate> new_pred =
+ pred_kind == Predicate::Kind::kAnd
+ ? Make<AndPredicate>(std::move(simplified_ops))
+ : Make<OrPredicate>(std::move(simplified_ops));
+
+ Predicate* new_pred_ptr = new_pred.get();
+ interned_and_or_instances_.emplace(
+ SignatureForAndOr(pred_kind, operands_slice), std::move(new_pred));
+ return new_pred_ptr;
+}
+
// Common code to create AndPredicate or OrPredicate instances.
Predicate* PredicateFactory::MakeAndOrImpl(
absl::Span<Predicate* const> operands, bool is_and) {
Predicate::Kind pred_kind =
is_and ? Predicate::Kind::kAnd : Predicate::Kind::kOr;
+ Predicate::Kind other_pred_kind =
+ is_and ? Predicate::Kind::kOr : Predicate::Kind::kAnd;
gtl::FlatSet<Predicate*> simplified_ops_set;
std::vector<Predicate*> simplified_ops;
for (Predicate* op : operands) {
@@ -472,30 +505,63 @@
}
}
- std::stable_sort(
- simplified_ops.begin(), simplified_ops.end(),
- [](Predicate* a, Predicate* b) { return a->hash() < b->hash(); });
+ // If all ops contain the same subop, then factor it out thanks to the
+ // distributive property. Such as:
+ // - (A & B) | (A & C) | (A & D) => A & (B | C | D)
+ // - (A | B) & (A | C) & (A | D) => A | (B & C & D)
+ //
+ // First find any predicates contained in all subops.
+ std::vector<Predicate*> common_inner_operands;
+ gtl::FlatSet<Predicate*> common_inner_operands_set;
+ for (Predicate* op : simplified_ops) {
+ if (op->kind() != other_pred_kind) {
+ common_inner_operands.clear();
+ break;
+ }
- auto it = interned_and_or_instances_.find({pred_kind, simplified_ops});
- if (it == interned_and_or_instances_.end()) {
- simplified_ops.shrink_to_fit();
- // NB! Because we'll use a non-owning reference to simplified_ops in the
- // key for interned_and_or_instances_ we need to be careful to std::move()
- // it all the way through.
- absl::Span<Predicate* const> operands_slice = simplified_ops;
- std::unique_ptr<Predicate> new_pred =
- is_and ? Make<AndPredicate>(std::move(simplified_ops))
- : Make<OrPredicate>(std::move(simplified_ops));
-
- Predicate* new_pred_ptr = new_pred.get();
- CHECK(interned_and_or_instances_
- .emplace(SignatureForAndOr(pred_kind, operands_slice),
- std::move(new_pred))
- .second);
- return new_pred_ptr;
- } else {
- return it->second.get();
+ if (common_inner_operands.empty()) {
+ common_inner_operands.insert(common_inner_operands.end(),
+ op->GetOperands().begin(),
+ op->GetOperands().end());
+ } else {
+ std::vector<Predicate*> sub_ops_intersection;
+ common_inner_operands.clear();
+ absl::c_copy_if(op->GetOperands(),
+ std::back_inserter(common_inner_operands),
+ [&](Predicate* sub_op) {
+ return common_inner_operands_set.count(sub_op) == 1;
+ });
+ }
+ if (common_inner_operands.empty()) break;
+ common_inner_operands_set.clear();
+ common_inner_operands_set.insert(common_inner_operands.begin(),
+ common_inner_operands.end());
}
+
+ if (common_inner_operands.empty()) {
+ return MakeInternedAndOr(std::move(simplified_ops), pred_kind);
+ }
+
+ // For all predicates that can be factored out, remove them and recreate the
+ // subops.
+ std::vector<Predicate*> factored_ops;
+ for (Predicate* op : simplified_ops) {
+ std::vector<Predicate*> new_sub_op_ops;
+ absl::c_copy_if(op->GetOperands(), std::back_inserter(new_sub_op_ops),
+ [&](Predicate* sub_op) {
+ return std::find(common_inner_operands.begin(),
+ common_inner_operands.end(),
+ sub_op) == common_inner_operands.end();
+ });
+ factored_ops.push_back(MakeAndOrImpl(new_sub_op_ops, !is_and));
+ }
+
+ Predicate* new_inner_op = MakeAndOrImpl(factored_ops, is_and);
+ std::vector<Predicate*> outer_ops;
+ outer_ops.push_back(new_inner_op);
+ outer_ops.insert(outer_ops.end(), common_inner_operands.begin(),
+ common_inner_operands.end());
+ return MakeAndOrImpl(outer_ops, !is_and);
}
class DeadnessAnalysisImpl : public DeadnessAnalysis {
@@ -507,7 +573,8 @@
Status PopulateWithReversePostOrder(absl::Span<Node* const> rpo);
bool HasInputsWithMismatchingDeadness(const Node& node) override;
void Print() const override;
- gtl::FlatMap<TensorId, string, TensorId::Hasher> PredicateMapAsString() const;
+ absl::flat_hash_map<TensorId, string, TensorId::Hasher> PredicateMapAsString()
+ const;
private:
enum class EdgeKind { kDataAndControl, kDataOnly, kControlOnly };
@@ -549,7 +616,7 @@
Status HandleNode(Node* n, std::vector<bool>* should_revisit);
const Graph& graph_;
- gtl::FlatMap<TensorId, Predicate*, TensorId::Hasher> predicate_map_;
+ absl::flat_hash_map<TensorId, Predicate*, TensorId::Hasher> predicate_map_;
PredicateFactory predicate_factory_;
bool vlog_;
};
@@ -912,9 +979,9 @@
return Status::OK();
}
-gtl::FlatMap<TensorId, string, TensorId::Hasher>
+absl::flat_hash_map<TensorId, string, TensorId::Hasher>
DeadnessAnalysisImpl::PredicateMapAsString() const {
- gtl::FlatMap<TensorId, string, TensorId::Hasher> result;
+ absl::flat_hash_map<TensorId, string, TensorId::Hasher> result;
std::vector<TensorId> tensor_ids;
for (const auto& kv_pair : predicate_map_) {
CHECK(result.insert({kv_pair.first, kv_pair.second->ToString()}).second);
diff --git a/tensorflow/compiler/jit/deadness_analysis_internal.h b/tensorflow/compiler/jit/deadness_analysis_internal.h
index 3df2679..3547823 100644
--- a/tensorflow/compiler/jit/deadness_analysis_internal.h
+++ b/tensorflow/compiler/jit/deadness_analysis_internal.h
@@ -16,15 +16,15 @@
#ifndef TENSORFLOW_COMPILER_JIT_DEADNESS_ANALYSIS_INTERNAL_H_
#define TENSORFLOW_COMPILER_JIT_DEADNESS_ANALYSIS_INTERNAL_H_
+#include "absl/container/flat_hash_map.h"
#include "tensorflow/core/graph/tensor_id.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
namespace tensorflow {
namespace deadness_analysis_internal {
// Returns a map describing the predicate each Tensor was mapped to. For
// testing purposes only.
-using PredicateMapTy = gtl::FlatMap<TensorId, string, TensorId::Hasher>;
+using PredicateMapTy = absl::flat_hash_map<TensorId, string, TensorId::Hasher>;
Status ComputePredicates(const Graph& graph, PredicateMapTy* out_predicate_map);
// Returns a map describing the predicate each Tensor was mapped to. For
diff --git a/tensorflow/compiler/jit/deadness_analysis_test.cc b/tensorflow/compiler/jit/deadness_analysis_test.cc
index 28a5604..617e314 100644
--- a/tensorflow/compiler/jit/deadness_analysis_test.cc
+++ b/tensorflow/compiler/jit/deadness_analysis_test.cc
@@ -384,10 +384,31 @@
EXPECT_FALSE(result->HasInputsWithMismatchingDeadness(*add2.node()));
}
-TEST(DeadnessAnalysisTest, NEGATIVE_AndOrDistributive) {
- // This demonstrates one of the weaknesses in the current approach -- since we
- // only do some basic simplifications we can't see that "(A|B)&C" ==
- // "(A&C)|(B&C)".
+TEST(DeadnessAnalysisTest, AndOrDistributiveSimplified) {
+ // (*A | (~*A & ((~*B & ~*A) | (~*A & *B)))) == #true
+ Scope root = Scope::NewRootScope().ExitOnError();
+
+ ops::Switch sw_0 = CreateSwitch(root, "A");
+ ops::Switch sw_1 = CreateSwitch(root, "B");
+ Output add0 =
+ ops::Add(root.WithOpName("and0"), sw_0.output_false, sw_1.output_true);
+ Output add1 =
+ ops::Add(root.WithOpName("and1"), sw_0.output_false, sw_1.output_false);
+ ops::Merge or2(root.WithOpName("or2"), {add0, add1});
+ Output add3 =
+ ops::Add(root.WithOpName("and3"), or2.output, sw_0.output_false);
+ ops::Merge or4(root.WithOpName("or4"), {add3, sw_0.output_true});
+
+ std::unique_ptr<DeadnessAnalysis> result;
+ TF_ASSERT_OK(AnalyzeDeadness(root.graph(), &result));
+
+ PredicateMapTy predicate_map;
+ TF_ASSERT_OK(ComputePredicates(*root.graph(), &predicate_map));
+ EXPECT_EQ(predicate_map[ControlOutputFor(or4.output)], "#true");
+}
+
+TEST(DeadnessAnalysisTest, AndOrDistributive) {
+ // (A|B)&C == (A&C)|(B&C)
Scope root = Scope::NewRootScope().ExitOnError();
ops::Switch sw_0 = CreateSwitch(root, "0");
@@ -408,7 +429,7 @@
std::unique_ptr<DeadnessAnalysis> result;
TF_ASSERT_OK(AnalyzeDeadness(root.graph(), &result));
- EXPECT_TRUE(result->HasInputsWithMismatchingDeadness(*add2.node()));
+ EXPECT_FALSE(result->HasInputsWithMismatchingDeadness(*add3.node()));
}
TEST(DeadnessAnalysisTest, Ternary) {
diff --git a/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc b/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc
index e0632ff..d165341 100644
--- a/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc
+++ b/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc
@@ -748,6 +748,12 @@
graph_->set_versions(graph_in->versions());
}
+ // TODO(b/116981129): Enhance how the device for the encapsulated subgraph is
+ // determined. In case of hard placement, ensure all the encapsulated nodes
+ // have the same requested device, which in turn will be the requested device
+ // for the entire encapsulated subgraph. In case of soft placement, use a
+ // deterministic approach to fill in the requested device. Handle co-location
+ // constraints similarly if they exist.
if (device_.empty()) {
device_ = node->assigned_device_name().empty()
? node->requested_device()
@@ -1357,28 +1363,31 @@
Status Encapsulator::GetFunctionNameAttr(
Node const* node, string* attr, string* outside_compilation_attr) const {
- Status s = GetNodeAttr(node->attrs(), group_attribute_, attr);
- if (s.code() == error::Code::NOT_FOUND) {
- // Return empty attr if there's no group_attribute.
- attr->clear();
- } else {
- TF_RETURN_IF_ERROR(s);
- }
- bool has_group_attr = s.ok();
- s = GetNodeAttr(node->attrs(), outside_compilation_attribute_,
- outside_compilation_attr);
- if (s.code() == error::Code::NOT_FOUND) {
- // Return empty attr if there's no outside_compilation attribute.
- outside_compilation_attr->clear();
- } else {
- TF_RETURN_IF_ERROR(s);
- if (!has_group_attr) {
- return errors::InvalidArgument(
- "Node ", node->name(), " has ", outside_compilation_attribute_,
- " attribute but no ", group_attribute_, " attribute.");
+ AttrSlice attrs = node->attrs();
+ attr->clear();
+ outside_compilation_attr->clear();
+ bool found_group_attribute = false;
+ bool found_outside_compilation_attribute = false;
+ for (const auto& node_attr : attrs) {
+ if (node_attr.first == group_attribute_) {
+ TF_RETURN_IF_ERROR(AttrValueHasType(node_attr.second, "string"));
+ *attr = node_attr.second.s();
+ found_group_attribute = true;
+ } else if (node_attr.first == outside_compilation_attribute_) {
+ TF_RETURN_IF_ERROR(AttrValueHasType(node_attr.second, "string"));
+ *outside_compilation_attr = node_attr.second.s();
+ found_outside_compilation_attribute = true;
}
+ if (found_group_attribute && found_outside_compilation_attribute) break;
}
- return Status::OK();
+
+ if (found_outside_compilation_attribute && !found_group_attribute) {
+ return errors::InvalidArgument(
+ "Node ", node->name(), " has ", outside_compilation_attribute_,
+ " attribute but no ", group_attribute_, " attribute.");
+ } else {
+ return Status::OK();
+ }
}
bool IsInSubgraph(const string& func_id, const string& outside_compilation_id) {
diff --git a/tensorflow/compiler/jit/encapsulate_xla_computations_pass.cc b/tensorflow/compiler/jit/encapsulate_xla_computations_pass.cc
index 97ef8cd..755c364 100644
--- a/tensorflow/compiler/jit/encapsulate_xla_computations_pass.cc
+++ b/tensorflow/compiler/jit/encapsulate_xla_computations_pass.cc
@@ -297,7 +297,9 @@
// Target the XLA CPU/GPU backends.
VLOG(2) << "Replacing with XlaLaunch";
+ VLOG(2) << "Device is " << launch->requested_device();
def.set_op("XlaLaunch");
+ def.set_device(launch->requested_device());
AddNodeAttr("Tconstants", DataTypeVector{}, &def);
AddNodeAttr("Targs", arg_types, &def);
AddNodeAttr("Nresources", num_variables, &def);
diff --git a/tensorflow/compiler/jit/encapsulate_xla_computations_pass_test.cc b/tensorflow/compiler/jit/encapsulate_xla_computations_pass_test.cc
index f643fb0..22531a4 100644
--- a/tensorflow/compiler/jit/encapsulate_xla_computations_pass_test.cc
+++ b/tensorflow/compiler/jit/encapsulate_xla_computations_pass_test.cc
@@ -19,7 +19,7 @@
#include "tensorflow/cc/ops/resource_variable_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/compiler/jit/encapsulate_subgraphs_pass.h"
-#include "tensorflow/compiler/tf2xla/cc/ops/xla_jit_op.h"
+#include "tensorflow/compiler/tf2xla/cc/ops/xla_jit_ops.h"
#include "tensorflow/compiler/tf2xla/test_util.h"
#include "tensorflow/core/framework/graph_to_functiondef.h"
#include "tensorflow/core/graph/graph_constructor.h"
@@ -55,6 +55,7 @@
.Input(u.node()->name(), 0, DT_RESOURCE)
.Input(v.node()->name(), 0, DT_RESOURCE)
.Input(w.node()->name(), 0, DT_RESOURCE)
+ .Device("/gpu:0")
.Attr(EncapsulateXlaComputationsPass::kXlaClusterAttr, "launch0")
.Attr("_variable_start_index", 4)
.Finalize(&def));
@@ -107,10 +108,11 @@
auto add_attrs = [](Node* node) {
node->AddAttr(EncapsulateXlaComputationsPass::kXlaClusterAttr, "launch0");
+ node->set_requested_device("/gpu:0");
};
auto b_identity = ops::Identity(scope.WithOpName("B_identity"), arg1);
-
+ add_attrs(b_identity.node());
auto read_u = ops::ReadVariableOp(scope.WithOpName("ReadU"), arg4, DT_FLOAT);
add_attrs(read_u.node());
auto read_v = ops::ReadVariableOp(scope.WithOpName("ReadV"), arg5, DT_FLOAT);
@@ -215,6 +217,7 @@
auto add_attrs = [](Node* node) {
node->AddAttr(EncapsulateXlaComputationsPass::kXlaClusterAttr, "launch0");
+ node->set_requested_device("/gpu:0");
};
auto b_identity = ops::Identity(scope.WithOpName("B_identity"), b);
@@ -317,8 +320,8 @@
NameAttrList function;
function.set_name("launch0");
auto launch = ops::XlaLaunch(
- scope.WithOpName("launch0"), std::initializer_list<Input>{},
- std::initializer_list<Input>{a, b, c, d},
+ scope.WithOpName("launch0").WithDevice("/gpu:0"),
+ std::initializer_list<Input>{}, std::initializer_list<Input>{a, b, c, d},
std::initializer_list<Input>{u, v, w},
DataTypeVector{DT_FLOAT, DT_INT32, DT_FLOAT, DT_FLOAT}, function);
diff --git a/tensorflow/compiler/jit/kernels/BUILD b/tensorflow/compiler/jit/kernels/BUILD
index 0839f1c..26cb3af 100644
--- a/tensorflow/compiler/jit/kernels/BUILD
+++ b/tensorflow/compiler/jit/kernels/BUILD
@@ -26,6 +26,7 @@
"//tensorflow/core:lib",
"//tensorflow/core:stream_executor_no_cuda",
"//tensorflow/core/kernels:variable_ops",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/memory",
],
alwayslink = 1,
diff --git a/tensorflow/compiler/jit/kernels/xla_ops.cc b/tensorflow/compiler/jit/kernels/xla_ops.cc
index a85006e..cfd27a6 100644
--- a/tensorflow/compiler/jit/kernels/xla_ops.cc
+++ b/tensorflow/compiler/jit/kernels/xla_ops.cc
@@ -15,6 +15,7 @@
#include "tensorflow/compiler/jit/kernels/xla_ops.h"
+#include "absl/container/flat_hash_map.h"
#include "absl/memory/memory.h"
#include "tensorflow/compiler/jit/defs.h"
#include "tensorflow/compiler/tf2xla/shape_util.h"
@@ -163,7 +164,7 @@
private:
mutex mutex_;
int64 key_counter_ GUARDED_BY(mutex_);
- gtl::FlatMap<KeyT, XlaExecutableClosure> closures_ GUARDED_BY(mutex_);
+ absl::flat_hash_map<KeyT, XlaExecutableClosure> closures_ GUARDED_BY(mutex_);
TF_DISALLOW_COPY_AND_ASSIGN(XlaExecutableClosureStore);
};
diff --git a/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc b/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc
index 4f9145b..2a80c74 100644
--- a/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc
+++ b/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc
@@ -15,6 +15,7 @@
#include "tensorflow/compiler/jit/mark_for_compilation_pass_test_helper.h"
+#include "absl/container/flat_hash_map.h"
#include "absl/memory/memory.h"
#include "absl/strings/match.h"
#include "tensorflow/cc/framework/ops.h"
@@ -61,10 +62,10 @@
return ids;
}
-gtl::FlatMap<string, std::vector<string>> GetClusterSets(
+absl::flat_hash_map<string, std::vector<string>> GetClusterSets(
const Graph& g, std::vector<string>* cluster_names = nullptr) {
CHECK(cluster_names == nullptr || cluster_names->empty());
- gtl::FlatMap<string, std::vector<string>> cluster_sets;
+ absl::flat_hash_map<string, std::vector<string>> cluster_sets;
for (const auto& p : GetClusters(g)) {
cluster_sets[p.second].push_back(p.first);
}
@@ -566,7 +567,7 @@
std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
TF_EXPECT_OK(root.ToGraph(graph.get()));
TF_ASSERT_OK(MarkForCompilationPassTestHelper::MarkForCompilation(&graph));
- gtl::FlatMap<string, std::vector<string>> cluster_sets =
+ absl::flat_hash_map<string, std::vector<string>> cluster_sets =
GetClusterSets(*graph);
ASSERT_EQ(cluster_sets.size(), 1);
std::vector<string> expected_clustered_nodes = {"AssignmentW", "ReadR",
@@ -586,7 +587,7 @@
std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
TF_EXPECT_OK(root.ToGraph(graph.get()));
TF_ASSERT_OK(MarkForCompilationPassTestHelper::MarkForCompilation(&graph));
- gtl::FlatMap<string, std::vector<string>> cluster_sets =
+ absl::flat_hash_map<string, std::vector<string>> cluster_sets =
GetClusterSets(*graph);
ASSERT_EQ(cluster_sets.size(), 1);
std::vector<string> expected_clustered_nodes = {"AssignmentW",
@@ -616,7 +617,7 @@
TF_ASSERT_OK(MarkForCompilationPassTestHelper::MarkForCompilation(&graph));
std::vector<string> cluster_names;
- gtl::FlatMap<string, std::vector<string>> cluster_sets =
+ absl::flat_hash_map<string, std::vector<string>> cluster_sets =
GetClusterSets(*graph, &cluster_names);
ASSERT_EQ(cluster_sets.size(), 2);
diff --git a/tensorflow/compiler/jit/resource_operation_safety_analysis.cc b/tensorflow/compiler/jit/resource_operation_safety_analysis.cc
index 56e35c0..657bb40 100644
--- a/tensorflow/compiler/jit/resource_operation_safety_analysis.cc
+++ b/tensorflow/compiler/jit/resource_operation_safety_analysis.cc
@@ -89,7 +89,6 @@
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/graph/algorithm.h"
#include "tensorflow/core/graph/tensor_id.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/lib/gtl/flatset.h"
#include "tensorflow/core/lib/hash/hash.h"
#include "tensorflow/core/util/ptr_util.h"
diff --git a/tensorflow/compiler/jit/xla_compilation_cache.h b/tensorflow/compiler/jit/xla_compilation_cache.h
index 10ad87e..17c0321 100644
--- a/tensorflow/compiler/jit/xla_compilation_cache.h
+++ b/tensorflow/compiler/jit/xla_compilation_cache.h
@@ -16,6 +16,7 @@
#ifndef TENSORFLOW_COMPILER_JIT_XLA_COMPILATION_CACHE_H_
#define TENSORFLOW_COMPILER_JIT_XLA_COMPILATION_CACHE_H_
+#include "absl/container/flat_hash_map.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "tensorflow/compiler/tf2xla/xla_context.h"
#include "tensorflow/compiler/xla/client/local_client.h"
@@ -24,7 +25,6 @@
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/lib/core/threadpool.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/thread_annotations.h"
@@ -152,7 +152,7 @@
};
mutex compile_cache_mu_;
- gtl::FlatMap<Signature, std::unique_ptr<Entry>, Signature::Hash> cache_
+ absl::flat_hash_map<Signature, std::unique_ptr<Entry>, Signature::Hash> cache_
GUARDED_BY(compile_cache_mu_);
struct CompileStats {
@@ -165,7 +165,7 @@
mutex compile_stats_mu_;
// Maps cluster names to compilation statistics for said cluster.
- gtl::FlatMap<string, CompileStats> compile_stats_
+ absl::flat_hash_map<string, CompileStats> compile_stats_
GUARDED_BY(compile_stats_mu_);
TF_DISALLOW_COPY_AND_ASSIGN(XlaCompilationCache);
diff --git a/tensorflow/compiler/jit/xla_device_ops.h b/tensorflow/compiler/jit/xla_device_ops.h
index 2ccee79..6967ad1 100644
--- a/tensorflow/compiler/jit/xla_device_ops.h
+++ b/tensorflow/compiler/jit/xla_device_ops.h
@@ -100,9 +100,15 @@
Name("VarHandleOp").Device(DEVICE).HostMemory("resource"), \
ResourceHandleOp<Var>); \
REGISTER_KERNEL_BUILDER( \
+ Name("_VarHandlesOp").Device(DEVICE).HostMemory("resources"), \
+ ResourceHandlesOp<Var>); \
+ REGISTER_KERNEL_BUILDER( \
Name("ReadVariableOp").Device(DEVICE).HostMemory("resource"), \
ReadVariableOp); \
REGISTER_KERNEL_BUILDER( \
+ Name("_ReadVariablesOp").Device(DEVICE).HostMemory("resources"), \
+ ReadVariablesOp); \
+ REGISTER_KERNEL_BUILDER( \
Name("DestroyResourceOp").Device(DEVICE).HostMemory("resource"), \
DestroyResourceOp); \
REGISTER_KERNEL_BUILDER(Name("Shape") \
diff --git a/tensorflow/compiler/tests/binary_ops_test.py b/tensorflow/compiler/tests/binary_ops_test.py
index e219cf3..1b39d53 100644
--- a/tensorflow/compiler/tests/binary_ops_test.py
+++ b/tensorflow/compiler/tests/binary_ops_test.py
@@ -1445,6 +1445,13 @@
np.array([4, 0], dtype=np.int32),
expected=np.zeros([4, 0], dtype=dtype))
+ x = np.arange(3).reshape((3, 1, 1, 1)).astype(dtype)
+ self._testBinary(
+ array_ops.broadcast_to,
+ x,
+ np.array((3, 7, 8, 9), dtype=np.int32),
+ expected=np.tile(x, (1, 7, 8, 9)))
+
if __name__ == "__main__":
googletest.main()
diff --git a/tensorflow/compiler/tests/fused_batchnorm_test.py b/tensorflow/compiler/tests/fused_batchnorm_test.py
index 8c018cc..374942a 100644
--- a/tensorflow/compiler/tests/fused_batchnorm_test.py
+++ b/tensorflow/compiler/tests/fused_batchnorm_test.py
@@ -29,6 +29,11 @@
from tensorflow.python.ops import nn
from tensorflow.python.platform import test
+DATA_FORMATS = (
+ ("_data_format_NHWC", "NHWC"),
+ ("_data_format_NCHW", "NCHW"),
+)
+
class FusedBatchNormTest(xla_test.XLATestCase, parameterized.TestCase):
@@ -65,12 +70,7 @@
grad_offset = np.sum(grad_y, axis=(0, 1, 2))
return grad_x, grad_scale, grad_offset
- @parameterized.named_parameters(
- ("_data_format_NHWC", "NHWC"),
- ("_data_format_NCHW", "NCHW"),
- ("_data_format_HWNC", "HWNC"),
- ("_data_format_HWCN", "HWCN"),
- )
+ @parameterized.named_parameters(*DATA_FORMATS)
def testInference(self, data_format):
channel = 3
x_shape = [2, 2, 6, channel]
@@ -170,30 +170,15 @@
self.assertAllClose(y_val, y_ref_converted, atol=1e-3)
self.assertAllClose(var_val, var_ref, atol=1e-3)
- @parameterized.named_parameters(
- ("_data_format_NHWC", "NHWC"),
- ("_data_format_NCHW", "NCHW"),
- ("_data_format_HWNC", "HWNC"),
- ("_data_format_HWCN", "HWCN"),
- )
+ @parameterized.named_parameters(*DATA_FORMATS)
def testLearning(self, data_format):
self._testLearning(False, data_format)
- @parameterized.named_parameters(
- ("_data_format_NHWC", "NHWC"),
- ("_data_format_NCHW", "NCHW"),
- ("_data_format_HWNC", "HWNC"),
- ("_data_format_HWCN", "HWCN"),
- )
+ @parameterized.named_parameters(*DATA_FORMATS)
def testLearningWithGradientChecker(self, data_format):
self._testLearning(True, data_format)
- @parameterized.named_parameters(
- ("_data_format_NHWC", "NHWC"),
- ("_data_format_NCHW", "NCHW"),
- ("_data_format_HWNC", "HWNC"),
- ("_data_format_HWCN", "HWCN"),
- )
+ @parameterized.named_parameters(*DATA_FORMATS)
def testGradientTraining(self, data_format):
# TODO(b/64270657): Use gradient_checker here in addition to comparing with
# this reference implementation.
@@ -241,12 +226,7 @@
self.assertAllClose(grad_scale_val, grad_scale_ref, atol=1e-2)
self.assertAllClose(grad_offset_val, grad_offset_ref, atol=1e-3)
- @parameterized.named_parameters(
- ("_data_format_NHWC", "NHWC"),
- ("_data_format_NCHW", "NCHW"),
- ("_data_format_HWNC", "HWNC"),
- ("_data_format_HWCN", "HWCN"),
- )
+ @parameterized.named_parameters(*DATA_FORMATS)
def testGradientInference(self, data_format):
# TODO(b/64270657): Use gradient_checker here in addition to comparing with
# this reference implementation.
diff --git a/tensorflow/compiler/tests/image_ops_test.py b/tensorflow/compiler/tests/image_ops_test.py
index bbe746e..68fdb5c 100644
--- a/tensorflow/compiler/tests/image_ops_test.py
+++ b/tensorflow/compiler/tests/image_ops_test.py
@@ -724,6 +724,49 @@
self.assertEqual(num_valid, 2)
self.assertAllClose(indices_tf[:num_valid], [3, 0])
+ def testNMS3Then1WithScoreMaxThresh(self):
+ # Three boxes are selected based on IOU.
+ # One is filtered out by score threshold.
+ # One is filtered out by max_output_size.
+
+ with compat.forward_compatibility_horizon(2018, 8, 8):
+ boxes_data = [[0, 0, 1, 1], [0, 0.1, 1, 1.1], [0, -0.1, 1, 0.9],
+ [0, 10, 1, 11], [0, 10.1, 1, 11.1], [0, 100, 1, 101]]
+ boxes_np = np.array(boxes_data, dtype=np.float32)
+
+ scores_data = [0.9, 0.75, 0.6, 0.95, 0.5, 0.3]
+ scores_np = np.array(scores_data, dtype=np.float32)
+ max_output_size = 1
+ iou_threshold_np = np.array(0.5, dtype=np.float32)
+ score_threshold_np = np.array(0.4, dtype=np.float32)
+
+ with self.cached_session() as sess:
+ boxes = array_ops.placeholder(boxes_np.dtype, shape=boxes_np.shape)
+ scores = array_ops.placeholder(scores_np.dtype, shape=scores_np.shape)
+ iou_threshold = array_ops.placeholder(iou_threshold_np.dtype,
+ iou_threshold_np.shape)
+ score_threshold = array_ops.placeholder(score_threshold_np.dtype,
+ score_threshold_np.shape)
+ with self.test_scope():
+ selected_indices = image_ops.non_max_suppression_padded(
+ boxes=boxes,
+ scores=scores,
+ max_output_size=max_output_size,
+ iou_threshold=iou_threshold,
+ score_threshold=score_threshold,
+ pad_to_max_output_size=True)
+ inputs_feed = {
+ boxes: boxes_np,
+ scores: scores_np,
+ iou_threshold: iou_threshold_np,
+ score_threshold: score_threshold_np
+ }
+ (indices_tf, num_valid) = sess.run(
+ selected_indices, feed_dict=inputs_feed)
+
+ self.assertEqual(indices_tf.size, max_output_size)
+ self.assertEqual(num_valid, 1)
+ self.assertAllClose(indices_tf[:num_valid], [3])
if __name__ == "__main__":
test.main()
diff --git a/tensorflow/compiler/tests/lstm.py b/tensorflow/compiler/tests/lstm.py
index 43c469d..73b3638 100644
--- a/tensorflow/compiler/tests/lstm.py
+++ b/tensorflow/compiler/tests/lstm.py
@@ -117,7 +117,7 @@
def RandomVar(shape, name=None):
"""Returns a variable of the given shape initialized to random values."""
- return variables.Variable(
+ return variables.VariableV1(
random_ops.random_uniform(shape), dtype=dtypes.float32, name=name)
diff --git a/tensorflow/compiler/tf2xla/BUILD b/tensorflow/compiler/tf2xla/BUILD
index ba1e3b2..3f631f9 100644
--- a/tensorflow/compiler/tf2xla/BUILD
+++ b/tensorflow/compiler/tf2xla/BUILD
@@ -635,6 +635,7 @@
"//tensorflow/core:lib",
"//tensorflow/core:ops",
"@com_google_absl//absl/algorithm:container",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/strings",
],
)
@@ -649,6 +650,7 @@
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"@com_google_absl//absl/algorithm:container",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/strings",
],
)
diff --git a/tensorflow/compiler/tf2xla/cc/BUILD b/tensorflow/compiler/tf2xla/cc/BUILD
index ea8d1b3..adcdb6c 100644
--- a/tensorflow/compiler/tf2xla/cc/BUILD
+++ b/tensorflow/compiler/tf2xla/cc/BUILD
@@ -30,14 +30,15 @@
tf_gen_op_wrapper_cc(
name = "xla_jit_op_gen",
- out_ops_file = "ops/xla_jit_op",
+ include_internal_ops = 1,
+ out_ops_file = "ops/xla_jit_ops",
deps = ["//tensorflow/compiler/jit/ops:xla_ops"],
)
cc_library(
name = "xla_jit_ops",
- srcs = ["ops/xla_jit_op.cc"],
- hdrs = ["ops/xla_jit_op.h"],
+ srcs = ["ops/xla_jit_ops.cc"],
+ hdrs = ["ops/xla_jit_ops.h"],
deps = [
"//tensorflow/cc:const_op",
"//tensorflow/cc:ops",
diff --git a/tensorflow/compiler/tf2xla/functionalize_control_flow.cc b/tensorflow/compiler/tf2xla/functionalize_control_flow.cc
index 98b333a..36c6f5d 100644
--- a/tensorflow/compiler/tf2xla/functionalize_control_flow.cc
+++ b/tensorflow/compiler/tf2xla/functionalize_control_flow.cc
@@ -92,45 +92,6 @@
});
const FunctionBody* body = flr->GetFunctionBody(handle);
- // If any node has associated functions, functionalize them first.
- // Gather nodes with associated functions first, because rewriting those nodes
- // might involve node deletion/addition. Avoid modifying nodes while iterating
- // it.
- std::vector<std::pair<Node*, std::vector<AssociatedFunctionInfo>>>
- nodes_to_associated_functions;
- for (auto* n : body->graph->nodes()) {
- auto associated_functions = GetAssociatedFunctions(*n, flr);
- if (!associated_functions.empty()) {
- nodes_to_associated_functions.push_back({n, associated_functions});
- }
- }
- for (auto iter : nodes_to_associated_functions) {
- Node* n = iter.first;
- auto associated_functions = iter.second;
- for (auto& associated_function : associated_functions) {
- string name = associated_function.func_name();
- string canonicalized_name = Canonicalize(name, AttrSlice(&attrs));
- auto iter = canonicalized_name_to_new_name->find(canonicalized_name);
- string new_name;
- if (iter != canonicalized_name_to_new_name->end()) {
- // If we already functionalized this function, skip functionalization
- // but still rewrite the node.
- new_name = iter->second;
- } else {
- new_name = fld->UniqueFunctionName(absl::StrCat(name, "_f15n_"));
- TF_RETURN_IF_ERROR(FunctionalizeControlFlowForFunction(
- name, new_name, attrs, fld, flr, canonicalized_name_to_new_name));
- (*canonicalized_name_to_new_name)[canonicalized_name] = new_name;
- }
- // Notice that if "n" is a function call, RewriteAssociatedFunction() will
- // delete it and create a new node instead, making "n" an invalid pointer.
- // That's fine because in that case, associated_functions will only have
- // one member and the loop will only run once.
- TF_RETURN_IF_ERROR(RewriteAssociatedFunction(
- body->graph, n, fld, associated_function, new_name));
- }
- }
-
// Call graph optimizer. The most important optimization we need is constant
// folding, which will replace ops like Shape/BroadcastGradientArgs with
// constant shape input. Without this optimization, those ops might become
@@ -152,9 +113,70 @@
opts.set_do_function_inlining(true);
opts.set_do_constant_folding(true);
GraphOptimizer optimizer(opts);
+ auto cf_consider_fn = [](const Node* n) {
+ // Skip SymbolicGradient op when doing constant folding.
+ // Enabling SymbolicGradient op in constant folding requires
+ // flr->device() to be non-null, and here we have not constructed
+ // proper Device object yet (it will be constructed in XlaCompiler).
+ return n->type_string() != FunctionLibraryDefinition::kGradientOp;
+ };
optimizer.Optimize(flr, flr->env(),
/*device=*/nullptr, &optimized_graph,
- /*shape_map=*/nullptr);
+ /*shape_map=*/nullptr, /*cse_consider_fn=*/nullptr,
+ cf_consider_fn);
+ if (VLOG_IS_ON(4)) {
+ dump_graph::DumpGraphToFile(
+ absl::StrCat("functionalize_control_flow_after_opt_", func_name),
+ *optimized_graph, fld);
+ }
+
+ // If any node has associated functions, functionalize them first.
+ // Gather nodes with associated functions first, because rewriting those nodes
+ // might involve node deletion/addition. Avoid modifying nodes while iterating
+ // it.
+ std::vector<std::pair<Node*, std::vector<AssociatedFunctionInfo>>>
+ nodes_to_associated_functions;
+ for (auto* n : optimized_graph->nodes()) {
+ auto associated_functions = GetAssociatedFunctions(*n, flr);
+ if (!associated_functions.empty()) {
+ nodes_to_associated_functions.push_back({n, associated_functions});
+ }
+ }
+ for (auto iter : nodes_to_associated_functions) {
+ Node* n = iter.first;
+ auto associated_functions = iter.second;
+ for (auto& associated_function : associated_functions) {
+ string name = associated_function.func_name();
+ string canonicalized_name =
+ Canonicalize(name, AttrSlice(&associated_function.attrs()));
+ auto iter = canonicalized_name_to_new_name->find(canonicalized_name);
+ string new_name;
+ if (iter != canonicalized_name_to_new_name->end()) {
+ // If we already functionalized this function, skip functionalization
+ // but still rewrite the node.
+ new_name = iter->second;
+ } else {
+ if (associated_function.type() ==
+ AssociatedFunctionInfo::AssociatedFunctionType::kSymbolicGradient) {
+ // For SymbolicGradient, `name` is always "SymbolicGradient",
+ // which is not very informative. Use node name instead.
+ new_name = fld->UniqueFunctionName(absl::StrCat(n->name(), "_f15n_"));
+ } else {
+ new_name = fld->UniqueFunctionName(absl::StrCat(name, "_f15n_"));
+ }
+ TF_RETURN_IF_ERROR(FunctionalizeControlFlowForFunction(
+ name, new_name, associated_function.attrs(), fld, flr,
+ canonicalized_name_to_new_name));
+ (*canonicalized_name_to_new_name)[canonicalized_name] = new_name;
+ }
+ // Notice that if "n" is a function call, RewriteAssociatedFunction() will
+ // delete it and create a new node instead, making "n" an invalid pointer.
+ // That's fine because in that case, associated_functions will only have
+ // one member and the loop will only run once.
+ TF_RETURN_IF_ERROR(RewriteAssociatedFunction(
+ optimized_graph.get(), n, fld, associated_function, new_name));
+ }
+ }
// Functionalize the function body.
if (VLOG_IS_ON(4)) {
diff --git a/tensorflow/compiler/tf2xla/kernels/batch_norm_op.cc b/tensorflow/compiler/tf2xla/kernels/batch_norm_op.cc
index b3ad0ae..a267c0c 100644
--- a/tensorflow/compiler/tf2xla/kernels/batch_norm_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/batch_norm_op.cc
@@ -34,12 +34,6 @@
OP_REQUIRES(
ctx, FormatFromString(data_format_str, &data_format_),
errors::InvalidArgument("Invalid data format: ", data_format_str));
- OP_REQUIRES(ctx,
- (data_format_ == FORMAT_NHWC || data_format_ == FORMAT_NCHW ||
- data_format_ == FORMAT_HWNC || data_format_ == FORMAT_HWCN),
- errors::InvalidArgument(
- "Unsupported data format ", ToString(data_format_),
- "; supported formats are NHWC, NCHW, HWNC and HWCN"));
}
void Compile(XlaOpKernelContext* ctx) override {
@@ -110,12 +104,6 @@
OP_REQUIRES(
ctx, FormatFromString(data_format_str, &data_format_),
errors::InvalidArgument("Invalid data format: ", data_format_str));
- OP_REQUIRES(ctx,
- (data_format_ == FORMAT_NHWC || data_format_ == FORMAT_NCHW ||
- data_format_ == FORMAT_HWNC || data_format_ == FORMAT_HWCN),
- errors::InvalidArgument(
- "Unsupported data format ", ToString(data_format_),
- "; supported formats are NHWC, NCHW, HWNC and HWCN"));
}
void Compile(XlaOpKernelContext* ctx) override {
diff --git a/tensorflow/compiler/tf2xla/kernels/binary_ops.cc b/tensorflow/compiler/tf2xla/kernels/binary_ops.cc
index 6667645..a988d3c 100644
--- a/tensorflow/compiler/tf2xla/kernels/binary_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/binary_ops.cc
@@ -103,6 +103,24 @@
XLA_MAKE_BINARY(FloorDiv,
FloorDivImpl(b, input_type(0), lhs, rhs, broadcast_helper));
+static xla::XlaOp XlogyImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x,
+ xla::XlaOp y, const BCast& broadcast_helper) {
+ std::tie(x, y) = XlaBinaryOp::Broadcast(b, x, y, broadcast_helper);
+ auto zero = XlaHelpers::Zero(b, dtype);
+ auto is_zero = xla::Eq(x, zero);
+ return xla::Select(is_zero, zero, xla::Mul(x, xla::Log(y)));
+}
+XLA_MAKE_BINARY(Xlogy, XlogyImpl(b, input_type(0), lhs, rhs, broadcast_helper));
+
+static xla::XlaOp XdivyImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x,
+ xla::XlaOp y, const BCast& broadcast_helper) {
+ std::tie(x, y) = XlaBinaryOp::Broadcast(b, x, y, broadcast_helper);
+ auto zero = XlaHelpers::Zero(b, dtype);
+ auto is_zero = xla::Eq(x, zero);
+ return xla::Select(is_zero, zero, xla::Div(x, y));
+}
+XLA_MAKE_BINARY(Xdivy, XdivyImpl(b, input_type(0), lhs, rhs, broadcast_helper));
+
// Implementation of FloorMod. Pseudo-code:
// T trunc_mod = std::fmod(x, y);
// return (x < T(0)) == (y < T(0)) ? trunc_mod : std::fmod(trunc_mod + y, y);
diff --git a/tensorflow/compiler/tf2xla/kernels/broadcast_to_op.cc b/tensorflow/compiler/tf2xla/kernels/broadcast_to_op.cc
index 4bd7c74..696c1c3 100644
--- a/tensorflow/compiler/tf2xla/kernels/broadcast_to_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/broadcast_to_op.cc
@@ -64,10 +64,9 @@
output_shape.DebugString()));
broadcast_dims.push_back(broadcast_shape.size());
- if (output_dims[i] == input_dims[i] || input_dims[i] == 1) {
+ if (output_dims[i] == input_dims[i]) {
broadcast_shape.push_back(output_dims[i]);
- }
- if (output_dims[i] != input_dims[i]) {
+ } else if (output_dims[i] != input_dims[i]) {
// Add dimensions [I, O/I], which we will later flatten to just
// [O]. We must do this in two phases since XLA broadcasting does not
// support tiling.
diff --git a/tensorflow/compiler/tf2xla/kernels/image_ops.cc b/tensorflow/compiler/tf2xla/kernels/image_ops.cc
index 33a73fe..921b434 100644
--- a/tensorflow/compiler/tf2xla/kernels/image_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/image_ops.cc
@@ -355,6 +355,9 @@
OP_REQUIRES(
context, output_size >= 0,
errors::InvalidArgument("Need output_size >= 0, got ", output_size));
+ OP_REQUIRES(context, output_size <= kint32max,
+ errors::InvalidArgument("Need output_size <= kint32Max, got ",
+ output_size));
xla::XlaOp score_thresh = context->Input("score_threshold");
xla::XlaOp iou_thresh = context->Input("iou_threshold");
@@ -439,12 +442,14 @@
xla::Broadcast(xla::ConstantR0<int32>(builder, 1), {num_boxes}),
xla::Broadcast(xla::ConstantR0<int32>(builder, 0), {num_boxes}));
- // num_valid is scalar.
- xla::XlaOp num_valid = xla::Reduce(
+ // num_valid is scalar. Value should be bound by output_size.
+ xla::XlaOp num_valid_total = xla::Reduce(
ones_included,
/*init_value=*/xla::ConstantR0<int>(builder, 0),
/*computation=*/CreateScalarAddComputation(xla::S32, builder),
/*dimensions_to_reduce=*/{0});
+ xla::XlaOp num_valid =
+ xla::Min(num_valid_total, xla::ConstantR0<int32>(builder, output_size));
xla::XlaOp output_tuple = TopK(scores_included, output_size);
xla::XlaOp selected_indices = xla::GetTupleElement(output_tuple, 1);
diff --git a/tensorflow/compiler/tf2xla/resource_operation_table.cc b/tensorflow/compiler/tf2xla/resource_operation_table.cc
index 20f2ce2..72b2409 100644
--- a/tensorflow/compiler/tf2xla/resource_operation_table.cc
+++ b/tensorflow/compiler/tf2xla/resource_operation_table.cc
@@ -15,7 +15,7 @@
#include "tensorflow/compiler/tf2xla/resource_operation_table.h"
#include "absl/algorithm/container.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
+#include "absl/container/flat_hash_map.h"
namespace tensorflow {
/*static*/ absl::string_view XlaResourceOpInfo::XlaResourceOpKindToString(
@@ -30,9 +30,9 @@
}
}
-static gtl::FlatMap<absl::string_view, XlaResourceOpInfo>*
+static absl::flat_hash_map<absl::string_view, XlaResourceOpInfo>*
CreateResourceOpInfoMap() {
- auto* result = new gtl::FlatMap<absl::string_view, XlaResourceOpInfo>;
+ auto* result = new absl::flat_hash_map<absl::string_view, XlaResourceOpInfo>;
auto add = [&](absl::string_view op, XlaResourceOpKind op_kind,
XlaResourceKind resource_kind) {
@@ -103,15 +103,15 @@
return result;
}
-static const gtl::FlatMap<absl::string_view, XlaResourceOpInfo>&
+static const absl::flat_hash_map<absl::string_view, XlaResourceOpInfo>&
GetStaticResourceOpInfoMap() {
- static gtl::FlatMap<absl::string_view, XlaResourceOpInfo>* op_info_map =
- CreateResourceOpInfoMap();
+ static absl::flat_hash_map<absl::string_view, XlaResourceOpInfo>*
+ op_info_map = CreateResourceOpInfoMap();
return *op_info_map;
}
const XlaResourceOpInfo* GetResourceOpInfoForOp(absl::string_view op) {
- const gtl::FlatMap<absl::string_view, XlaResourceOpInfo>& op_infos =
+ const absl::flat_hash_map<absl::string_view, XlaResourceOpInfo>& op_infos =
GetStaticResourceOpInfoMap();
auto it = op_infos.find(op);
return it == op_infos.end() ? nullptr : &it->second;
diff --git a/tensorflow/compiler/tf2xla/resource_operation_table_test.cc b/tensorflow/compiler/tf2xla/resource_operation_table_test.cc
index a85ef04..956f597 100644
--- a/tensorflow/compiler/tf2xla/resource_operation_table_test.cc
+++ b/tensorflow/compiler/tf2xla/resource_operation_table_test.cc
@@ -16,6 +16,7 @@
#include "tensorflow/compiler/tf2xla/resource_operation_table.h"
#include "absl/algorithm/container.h"
+#include "absl/container/flat_hash_map.h"
#include "absl/strings/str_join.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "tensorflow/core/lib/core/status_test_util.h"
@@ -33,7 +34,7 @@
}
TEST(ResourceOperationTableTest, HaveAllResourceOps) {
- gtl::FlatMap<string, bool> known_resource_ops;
+ absl::flat_hash_map<string, bool> known_resource_ops;
for (absl::string_view known_resource_op :
resource_op_table_internal::GetKnownResourceOps()) {
ASSERT_TRUE(
diff --git a/tensorflow/compiler/tf2xla/tf2xla_util.cc b/tensorflow/compiler/tf2xla/tf2xla_util.cc
index d6f42ba..01dd3ba 100644
--- a/tensorflow/compiler/tf2xla/tf2xla_util.cc
+++ b/tensorflow/compiler/tf2xla/tf2xla_util.cc
@@ -336,9 +336,9 @@
}
if (node_def.op() == FunctionLibraryDefinition::kGradientOp) {
- // Skip gradient op. Gradient op has "f" attr, which is set to the function
- // we are getting gradient for. That function is not associated with the op.
- return false;
+ // Gradient op has "f" attr, which is set to the function we are getting
+ // gradient for. We need to functionalize the gradient function.
+ return true;
}
for (const auto& iter : node_def.attr()) {
@@ -357,17 +357,18 @@
if (flr->GetFunctionLibraryDefinition()->Contains(op)) {
// This is a function call node.
AttrValueMap attrs(node.attrs().begin(), node.attrs().end());
- results.emplace_back(AssociatedFunctionInfo(op, attrs));
+ results.emplace_back(AssociatedFunctionInfo::FunctionCall(op, attrs));
} else if (node.type_string() == FunctionLibraryDefinition::kGradientOp) {
- // Skip gradient op. Gradient op has "f" attr, which is set to the function
- // we are getting gradient for. That function is not associated with the op.
+ // This is a SymbolicGradient op.
+ AttrValueMap attrs(node.attrs().begin(), node.attrs().end());
+ results.emplace_back(AssociatedFunctionInfo::SymbolicGradient(op, attrs));
} else {
// Collect all function attrs for the node.
for (auto& iter : node.attrs()) {
if (iter.second.has_func()) {
VLOG(2) << "Found function attr for node " << node.name() << ": "
<< iter.first << " = " << iter.second.func().name();
- results.emplace_back(AssociatedFunctionInfo(
+ results.emplace_back(AssociatedFunctionInfo::FunctionAttr(
iter.second.func().name(), iter.second.func().attr(), iter.first));
}
}
@@ -410,6 +411,21 @@
graph->RemoveNode(node);
break;
}
+ case AssociatedFunctionInfo::kSymbolicGradient: {
+ NameAttrList func;
+ TF_RETURN_IF_ERROR(GetNodeAttr(
+ node->attrs(), FunctionLibraryDefinition::kFuncAttr, &func));
+ GradientDef gradient_def;
+ gradient_def.set_function_name(func.name());
+ gradient_def.set_gradient_func(rewritten_function_name);
+ string original_grad_func = fld->FindGradient(func.name());
+ if (original_grad_func.empty()) {
+ TF_RETURN_IF_ERROR(fld->AddGradientDef(gradient_def));
+ } else if (original_grad_func != rewritten_function_name) {
+ TF_RETURN_IF_ERROR(fld->ReplaceGradient(gradient_def));
+ }
+ break;
+ }
case AssociatedFunctionInfo::kFunctionAttr: {
// Change function attr to rewritten functions.
NameAttrList func;
diff --git a/tensorflow/compiler/tf2xla/tf2xla_util.h b/tensorflow/compiler/tf2xla/tf2xla_util.h
index 6065d0b..53eab8b 100644
--- a/tensorflow/compiler/tf2xla/tf2xla_util.h
+++ b/tensorflow/compiler/tf2xla/tf2xla_util.h
@@ -65,21 +65,33 @@
class AssociatedFunctionInfo {
public:
enum AssociatedFunctionType {
- kFunctionCallNode = 0,
- kFunctionAttr = 1,
+ kFunctionAttr = 0,
+ kFunctionCallNode = 1,
+ kSymbolicGradient = 2,
};
- // The node is a function call.
- AssociatedFunctionInfo(const string& func_name, const AttrValueMap& attrs)
- : type_(kFunctionCallNode), func_name_(func_name), attrs_(attrs) {}
-
// The function is an attr of the node.
- AssociatedFunctionInfo(const string& func_name, const AttrValueMap& attrs,
- const string& attr_name)
- : type_(kFunctionAttr),
- func_name_(func_name),
- attrs_(attrs),
- attr_name_(attr_name) {}
+ static AssociatedFunctionInfo FunctionAttr(const string& func_name,
+ const AttrValueMap& attrs,
+ const string& attr_name) {
+ return AssociatedFunctionInfo(kFunctionAttr, func_name, attrs, attr_name);
+ }
+
+ // The node is a function call.
+ static AssociatedFunctionInfo FunctionCall(const string& func_name,
+ const AttrValueMap& attrs) {
+ // attr_name will not be used in this case.
+ return AssociatedFunctionInfo(kFunctionCallNode, func_name, attrs,
+ /*attr_name=*/"");
+ }
+
+ // The node is a SymbolicGradient op.
+ static AssociatedFunctionInfo SymbolicGradient(const string& func_name,
+ const AttrValueMap& attrs) {
+ // attr_name will not be used in this case.
+ return AssociatedFunctionInfo(kSymbolicGradient, func_name, attrs,
+ /*attr_name=*/"");
+ }
AssociatedFunctionType type() const { return type_; }
@@ -90,6 +102,13 @@
const AttrValueMap& attrs() const { return attrs_; }
private:
+ AssociatedFunctionInfo(AssociatedFunctionType type, const string& func_name,
+ const AttrValueMap& attrs, const string& attr_name)
+ : type_(type),
+ func_name_(func_name),
+ attrs_(attrs),
+ attr_name_(attr_name) {}
+
// Available for all instances.
AssociatedFunctionType type_;
string func_name_;
@@ -105,14 +124,18 @@
// Gets functions associated with the node. Current cases:
// 1. For function call node, its function name;
-// 2. For nodes like XlaWhile/XlaIf, all their function attributes.
+// 2. For SymbolicGradient op, returned func_name will be "SymbolicGradient",
+// and returned attrs will be this node's attributes;
+// 3. For nodes like XlaWhile/XlaIf, all their function attributes.
std::vector<AssociatedFunctionInfo> GetAssociatedFunctions(
const Node& node, FunctionLibraryRuntime* flr);
// Changes associated functions for the node. Current cases:
// 1. For function call node, creates a new node with the new function name and
// remove the old node;
-// 2. For nodes like XlaWhile/XlaIf, modify their function attributes.
+// 2. For SymbolicGradient op, add or replace GradientDef in
+// FunctionLibraryDefinition;
+// 3. For nodes like XlaWhile/XlaIf, modify their function attributes.
Status RewriteAssociatedFunction(
Graph* graph, Node* node, FunctionLibraryDefinition* fld,
const AssociatedFunctionInfo& associated_function,
diff --git a/tensorflow/compiler/tf2xla/type_util.h b/tensorflow/compiler/tf2xla/type_util.h
index bda667e..6354216 100644
--- a/tensorflow/compiler/tf2xla/type_util.h
+++ b/tensorflow/compiler/tf2xla/type_util.h
@@ -25,6 +25,14 @@
// Converts a Tensorflow DataType to an XLA PrimitiveType.
Status DataTypeToPrimitiveType(DataType data_type, xla::PrimitiveType* type);
+// N.B.: there is intentionally no function to convert an XLA PrimitiveType to
+// a TensorFlow DataType. The mapping from TF types to XLA types is not
+// one-to-one: for example, both DT_INT8 and DT_QINT8 map to xla::S8. So the
+// inverse would not be a well-defined function. If you find that you want the
+// inverse mapping, then most likely you should be preserving the original
+// TensorFlow type, rather than trying to convert an XLA type into a TensorFlow
+// type.
+
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_TF2XLA_TYPE_UTIL_H_
diff --git a/tensorflow/compiler/xla/client/BUILD b/tensorflow/compiler/xla/client/BUILD
index f825f67b..1191cff 100644
--- a/tensorflow/compiler/xla/client/BUILD
+++ b/tensorflow/compiler/xla/client/BUILD
@@ -220,6 +220,7 @@
"//tensorflow/compiler/xla/service:shape_inference",
"//tensorflow/core:lib",
"@com_google_absl//absl/algorithm:container",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:span",
diff --git a/tensorflow/compiler/xla/client/lib/testing.cc b/tensorflow/compiler/xla/client/lib/testing.cc
index 25cc37e..ff0ec76 100644
--- a/tensorflow/compiler/xla/client/lib/testing.cc
+++ b/tensorflow/compiler/xla/client/lib/testing.cc
@@ -97,13 +97,11 @@
<< "Computation should have progran shape.";
auto program_shape = computation.proto().program_shape();
- // Create and run a program which produces a tuple with one element per
- // parameter, then return the tuple's constituent buffers.
- std::vector<Shape> param_shapes(program_shape.parameters().begin(),
- program_shape.parameters().end());
- auto fake_input_tuple =
- MakeFakeDataOrDie(ShapeUtil::MakeTupleShape(param_shapes), client);
- return client->DeconstructTuple(*fake_input_tuple).ValueOrDie();
+ std::vector<std::unique_ptr<GlobalData>> results;
+ for (const Shape& shape : program_shape.parameters()) {
+ results.push_back(MakeFakeDataOrDie(shape, client));
+ }
+ return results;
}
} // namespace xla
diff --git a/tensorflow/compiler/xla/client/xla_builder.cc b/tensorflow/compiler/xla/client/xla_builder.cc
index 95ff643..5277de6 100644
--- a/tensorflow/compiler/xla/client/xla_builder.cc
+++ b/tensorflow/compiler/xla/client/xla_builder.cc
@@ -1278,7 +1278,7 @@
XlaOp XlaBuilder::CustomCall(const string& call_target_name,
absl::Span<const XlaOp> operands,
- const Shape& shape) {
+ const Shape& shape, const string& opaque) {
return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> {
HloInstructionProto instr;
if (absl::StartsWith(call_target_name, "$")) {
@@ -1289,6 +1289,7 @@
}
*instr.mutable_shape() = shape;
instr.set_custom_call_target(call_target_name);
+ instr.set_custom_call_opaque(opaque);
return AddInstruction(std::move(instr), HloOpcode::kCustomCall, operands);
});
}
@@ -2681,8 +2682,9 @@
}
XlaOp CustomCall(XlaBuilder* builder, const string& call_target_name,
- absl::Span<const XlaOp> operands, const Shape& shape) {
- return builder->CustomCall(call_target_name, operands, shape);
+ absl::Span<const XlaOp> operands, const Shape& shape,
+ const string& opaque) {
+ return builder->CustomCall(call_target_name, operands, shape, opaque);
}
XlaOp Complex(const XlaOp& real, const XlaOp& imag,
diff --git a/tensorflow/compiler/xla/client/xla_builder.h b/tensorflow/compiler/xla/client/xla_builder.h
index d0c59fa..b7295e8 100644
--- a/tensorflow/compiler/xla/client/xla_builder.h
+++ b/tensorflow/compiler/xla/client/xla_builder.h
@@ -21,6 +21,7 @@
#include <type_traits>
#include <utility>
+#include "absl/container/flat_hash_map.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "tensorflow/compiler/xla/client/padding.h"
@@ -34,7 +35,6 @@
#include "tensorflow/compiler/xla/statusor.h"
#include "tensorflow/compiler/xla/types.h"
#include "tensorflow/compiler/xla/xla_data.pb.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/lib/gtl/flatset.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/stacktrace.h"
@@ -577,11 +577,9 @@
absl::Span<const XlaOp> operands);
// Enqueues a custom call instruction onto the computation.
- // During code generation, a call instruction is emitted which targets a
- // symbol with the name |call_target_name|. The |operands| are passed to the
- // call instruction. |shape| is the resultant shape.
XlaOp CustomCall(const string& call_target_name,
- absl::Span<const XlaOp> operands, const Shape& shape);
+ absl::Span<const XlaOp> operands, const Shape& shape,
+ const string& opaque);
// The following methods enqueue element-wise binary arithmetic operations
// onto the computation. The shapes of the operands have to match unless one
@@ -1029,7 +1027,7 @@
// A map from XlaOp::Handle to the index in the instructions_ vector where the
// instruction is held.
- tensorflow::gtl::FlatMap<int64, int64> handle_to_index_;
+ absl::flat_hash_map<int64, int64> handle_to_index_;
// The embedded computations used by this computation. Each computation was
// the entry computation of some XlaComputation, the key is the unique id of
@@ -1195,7 +1193,8 @@
friend XlaOp Call(XlaBuilder* builder, const XlaComputation& computation,
absl::Span<const XlaOp> operands);
friend XlaOp CustomCall(XlaBuilder* builder, const string& call_target_name,
- absl::Span<const XlaOp> operands, const Shape& shape);
+ absl::Span<const XlaOp> operands, const Shape& shape,
+ const string& opaque);
friend XlaOp Complex(const XlaOp& real, const XlaOp& imag,
absl::Span<const int64> broadcast_dimensions);
friend XlaOp Conj(const XlaOp& operand);
@@ -1717,12 +1716,17 @@
XlaOp Call(XlaBuilder* builder, const XlaComputation& computation,
absl::Span<const XlaOp> operands);
-// Enqueues a custom call instruction onto the computation.
-// During code generation, a call instruction is emitted which targets a
-// symbol with the name |call_target_name|. The |operands| are passed to the
-// call instruction. |shape| is the resultant shape.
+// Enqueues a custom call instruction onto the computation. A custom call
+// invokes code external to XLA. The |operands| are passed to the external code,
+// and the external code is expected to produce a result of the given
+// |shape|. The exact mechanism is backend-specific. For example, in the CPU
+// backend, a call instruction is emitted which targets a symbol with the name
+// |call_target_name|. |call_target_name| and |opaque| can arbitrary strings,
+// but |call_target_name| should be short as it may be used in labels. |opaque|
+// can encode arbitrarily large amounts of information.
XlaOp CustomCall(XlaBuilder* builder, const string& call_target_name,
- absl::Span<const XlaOp> operands, const Shape& shape);
+ absl::Span<const XlaOp> operands, const Shape& shape,
+ const string& opaque = "");
// The following methods enqueue element-wise binary arithmetic operations
// onto the computation. The shapes of the operands have to match unless one
diff --git a/tensorflow/compiler/xla/executable_run_options.cc b/tensorflow/compiler/xla/executable_run_options.cc
index a472747..0f9b591 100644
--- a/tensorflow/compiler/xla/executable_run_options.cc
+++ b/tensorflow/compiler/xla/executable_run_options.cc
@@ -45,6 +45,16 @@
return stream_;
}
+ExecutableRunOptions& ExecutableRunOptions::set_host_to_device_stream(
+ stream_executor::Stream* stream) {
+ host_to_device_stream_ = stream;
+ return *this;
+}
+
+stream_executor::Stream* ExecutableRunOptions::host_to_device_stream() const {
+ return host_to_device_stream_;
+}
+
ExecutableRunOptions& ExecutableRunOptions::set_intra_op_thread_pool(
const Eigen::ThreadPoolDevice* intra_op_thread_pool) {
intra_op_thread_pool_ = intra_op_thread_pool;
diff --git a/tensorflow/compiler/xla/executable_run_options.h b/tensorflow/compiler/xla/executable_run_options.h
index 416131b..ba3217f 100644
--- a/tensorflow/compiler/xla/executable_run_options.h
+++ b/tensorflow/compiler/xla/executable_run_options.h
@@ -65,6 +65,13 @@
ExecutableRunOptions& set_stream(stream_executor::Stream* stream);
stream_executor::Stream* stream() const;
+ // If set, this is the stream to perform any pre-computation transfers on.
+ // The platform of the stream must match the platform the executable was
+ // built for. A value of nullptr indicates the option has not been set.
+ ExecutableRunOptions& set_host_to_device_stream(
+ stream_executor::Stream* stream);
+ stream_executor::Stream* host_to_device_stream() const;
+
// Sets the thread pool device on which to run Eigen subcomputations.
// Does not take ownership.
ExecutableRunOptions& set_intra_op_thread_pool(
@@ -90,6 +97,7 @@
const Eigen::ThreadPoolDevice* intra_op_thread_pool_ = nullptr;
ExecutionProfile* execution_profile_ = nullptr;
int rng_seed_ = 0;
+ stream_executor::Stream* host_to_device_stream_ = nullptr;
};
} // namespace xla
diff --git a/tensorflow/compiler/xla/service/BUILD b/tensorflow/compiler/xla/service/BUILD
index e800cf4..8da6364 100644
--- a/tensorflow/compiler/xla/service/BUILD
+++ b/tensorflow/compiler/xla/service/BUILD
@@ -146,6 +146,7 @@
"//tensorflow/compiler/xla:shape_util",
"//tensorflow/compiler/xla:util",
"//tensorflow/core:lib",
+ "@com_google_absl//absl/container:flat_hash_map",
],
)
@@ -250,6 +251,7 @@
"//tensorflow/compiler/xla:xla_data_proto",
"//tensorflow/core:lib",
"@com_google_absl//absl/algorithm:container",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:inlined_vector",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
@@ -333,6 +335,7 @@
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"@com_google_absl//absl/algorithm:container",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:inlined_vector",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
@@ -395,6 +398,7 @@
"//tensorflow/compiler/xla:util",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/types:span",
],
)
@@ -485,6 +489,7 @@
"//tensorflow/compiler/xla:status_macros",
"//tensorflow/compiler/xla:util",
"//tensorflow/core:lib",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
@@ -903,6 +908,7 @@
"//tensorflow/compiler/xla:util",
"//tensorflow/compiler/xla:xla_data_proto",
"//tensorflow/core:lib",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
],
@@ -952,6 +958,7 @@
deps = [
"//tensorflow/compiler/xla:types",
"//tensorflow/core:lib",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/strings",
],
)
@@ -987,6 +994,7 @@
"//tensorflow/compiler/xla:types",
"//tensorflow/compiler/xla:util",
"//tensorflow/core:lib",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
],
@@ -1034,6 +1042,7 @@
"//tensorflow/compiler/xla:util",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
@@ -1087,6 +1096,7 @@
"//tensorflow/compiler/xla:types",
"//tensorflow/compiler/xla:util",
"//tensorflow/core:lib",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
],
@@ -1125,6 +1135,7 @@
"//tensorflow/compiler/xla:statusor",
"//tensorflow/compiler/xla:util",
"//tensorflow/core:lib",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/memory",
],
)
@@ -1146,6 +1157,7 @@
"//tensorflow/compiler/xla/tests:xla_internal_test_main",
"//tensorflow/core:lib",
"//tensorflow/core:test",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/memory",
],
)
@@ -1196,6 +1208,7 @@
"//tensorflow/compiler/xla:statusor",
"//tensorflow/compiler/xla:util",
"//tensorflow/core:lib",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/types:optional",
],
@@ -1216,6 +1229,7 @@
"//tensorflow/compiler/xla:types",
"//tensorflow/compiler/xla:util",
"//tensorflow/core:lib",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:span",
@@ -1260,6 +1274,7 @@
"//tensorflow/compiler/xla:util",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
+ "@com_google_absl//absl/container:flat_hash_map",
],
)
@@ -1280,6 +1295,7 @@
"//tensorflow/compiler/xla/tests:xla_internal_test_main",
"//tensorflow/core:test",
"@com_google_absl//absl/algorithm:container",
+ "@com_google_absl//absl/container:flat_hash_map",
],
)
@@ -1304,6 +1320,7 @@
"//tensorflow/compiler/xla:util",
"//tensorflow/core:lib",
"@com_google_absl//absl/algorithm:container",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/memory",
],
)
@@ -1330,6 +1347,7 @@
"//tensorflow/compiler/xla/service:hlo",
"//tensorflow/compiler/xla/service:hlo_pass",
"//tensorflow/core:lib",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/strings",
],
)
@@ -1385,6 +1403,7 @@
"//tensorflow/compiler/xla:util",
"//tensorflow/compiler/xla:xla_data_proto",
"//tensorflow/core:lib",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/types:optional",
"@com_google_absl//absl/types:span",
],
@@ -1640,6 +1659,7 @@
":while_loop_analysis",
"//tensorflow/compiler/xla:statusor",
"//tensorflow/core:lib",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:optional",
],
@@ -1671,6 +1691,7 @@
"//tensorflow/compiler/xla:types",
"//tensorflow/compiler/xla:util",
"//tensorflow/core:lib",
+ "@com_google_absl//absl/container:flat_hash_map",
],
)
@@ -2203,6 +2224,7 @@
"//tensorflow/compiler/xla:util",
"//tensorflow/compiler/xla:xla_data_proto",
"//tensorflow/core:lib",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:span",
],
@@ -2263,6 +2285,7 @@
"//tensorflow/compiler/xla:util",
"//tensorflow/compiler/xla:xla_data_proto",
"//tensorflow/core:lib",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:inlined_vector",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
@@ -2319,6 +2342,7 @@
"//tensorflow/compiler/xla:util",
"//tensorflow/compiler/xla:xla_data_proto",
"//tensorflow/core:lib",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
@@ -2345,6 +2369,7 @@
"//tensorflow/compiler/xla:types",
"//tensorflow/compiler/xla:util",
"//tensorflow/core:lib",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/strings",
],
)
@@ -2416,6 +2441,7 @@
"//tensorflow/compiler/xla:status_macros",
"//tensorflow/compiler/xla:util",
"//tensorflow/core:lib",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
],
@@ -2460,6 +2486,7 @@
"//tensorflow/compiler/xla:util",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:inlined_vector",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
@@ -2588,6 +2615,7 @@
"//tensorflow/compiler/xla:types",
"//tensorflow/compiler/xla:util",
"//tensorflow/core:lib",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
@@ -2701,6 +2729,7 @@
"//tensorflow/compiler/xla:types",
"//tensorflow/compiler/xla:util",
"//tensorflow/core:lib",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/memory",
],
)
@@ -3147,6 +3176,7 @@
":hlo_pass_pipeline",
"//tensorflow/compiler/xla:shape_util",
"//tensorflow/core:lib",
+ "@com_google_absl//absl/container:flat_hash_map",
],
)
@@ -3269,6 +3299,7 @@
"//tensorflow/compiler/xla:util",
"//tensorflow/core:lib",
"@com_google_absl//absl/algorithm:container",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:inlined_vector",
],
)
@@ -3298,6 +3329,7 @@
"//tensorflow/compiler/xla:util",
"//tensorflow/core:lib",
"@com_google_absl//absl/algorithm:container",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:inlined_vector",
],
)
@@ -3354,6 +3386,7 @@
"//tensorflow/core:lib",
"//tensorflow/core:ptr_util",
"@com_google_absl//absl/algorithm:container",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:inlined_vector",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:optional",
diff --git a/tensorflow/compiler/xla/service/allocation_tracker.h b/tensorflow/compiler/xla/service/allocation_tracker.h
index a7d8927..af227fe 100644
--- a/tensorflow/compiler/xla/service/allocation_tracker.h
+++ b/tensorflow/compiler/xla/service/allocation_tracker.h
@@ -22,6 +22,7 @@
#include <string>
#include <vector>
+#include "absl/container/flat_hash_map.h"
#include "tensorflow/compiler/xla/service/backend.h"
#include "tensorflow/compiler/xla/statusor.h"
#include "tensorflow/compiler/xla/types.h"
@@ -110,7 +111,7 @@
// A map from device memory opaque value to allocation. One such map is
// maintained per device ordinal.
- using AllocationMap = tensorflow::gtl::FlatMap<const void*, Allocation>;
+ using AllocationMap = absl::flat_hash_map<const void*, Allocation>;
tensorflow::mutex mutex_;
@@ -146,7 +147,7 @@
// non-owning "view" into a tuple's sub-buffers. The sub-buffers are then
// free'd when both the view *and* the original tuple are Unregistered. This
// refcounting is managed in opaque_to_allocation_map_.
- tensorflow::gtl::FlatMap<int64, std::vector<std::unique_ptr<ShapedBuffer>>>
+ absl::flat_hash_map<int64, std::vector<std::unique_ptr<ShapedBuffer>>>
handle_to_shaped_buffers_ GUARDED_BY(mutex_);
TF_DISALLOW_COPY_AND_ASSIGN(AllocationTracker);
diff --git a/tensorflow/compiler/xla/service/batchnorm_expander.cc b/tensorflow/compiler/xla/service/batchnorm_expander.cc
index 30d33e0..f70f6dd 100644
--- a/tensorflow/compiler/xla/service/batchnorm_expander.cc
+++ b/tensorflow/compiler/xla/service/batchnorm_expander.cc
@@ -35,7 +35,6 @@
#include "tensorflow/compiler/xla/xla_data.pb.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/status.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/types.h"
diff --git a/tensorflow/compiler/xla/service/bfloat16_propagation.h b/tensorflow/compiler/xla/service/bfloat16_propagation.h
index 6a62439..c74326f 100644
--- a/tensorflow/compiler/xla/service/bfloat16_propagation.h
+++ b/tensorflow/compiler/xla/service/bfloat16_propagation.h
@@ -21,6 +21,7 @@
#include <unordered_set>
#include <vector>
+#include "absl/container/flat_hash_map.h"
#include "tensorflow/compiler/xla/service/bfloat16_support.h"
#include "tensorflow/compiler/xla/service/hlo_dataflow_analysis.h"
#include "tensorflow/compiler/xla/service/hlo_instruction.h"
@@ -186,7 +187,7 @@
// Mapping from each HloComputation to the number of callers to it in the
// module. Populated at the beginning of this pass.
- tensorflow::gtl::FlatMap<const HloComputation*, int64> caller_counts_;
+ absl::flat_hash_map<const HloComputation*, int64> caller_counts_;
// We first store the potential F32-to-BF16 changes to changes_to_bf16_, which
// are subject to further adjustment, then finally applied to the HLOs. This
@@ -195,8 +196,7 @@
//
// For each HloInstruction, changes_to_bf16_ stores the affected buffers in
// the output as a map from in-place pointers to subshapes to shape indices.
- tensorflow::gtl::FlatMap<HloInstruction*,
- tensorflow::gtl::FlatMap<Shape*, ShapeIndex>>
+ absl::flat_hash_map<HloInstruction*, absl::flat_hash_map<Shape*, ShapeIndex>>
changes_to_bf16_;
// Whether the last processed HLO module has been changed by this pass.
diff --git a/tensorflow/compiler/xla/service/buffer_assignment.cc b/tensorflow/compiler/xla/service/buffer_assignment.cc
index 34a7be0..3efa0b1 100644
--- a/tensorflow/compiler/xla/service/buffer_assignment.cc
+++ b/tensorflow/compiler/xla/service/buffer_assignment.cc
@@ -22,6 +22,7 @@
#include <ostream>
#include <utility>
+#include "absl/container/flat_hash_map.h"
#include "absl/memory/memory.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_format.h"
@@ -41,9 +42,9 @@
namespace xla {
namespace {
+using absl::flat_hash_map;
using absl::StrAppend;
using absl::StrAppendFormat;
-using ::tensorflow::gtl::FlatMap;
using ::tensorflow::gtl::FlatSet;
using ::tensorflow::strings::HumanReadableNumBytes;
@@ -519,7 +520,8 @@
// BufferAllocation.
void BufferAssignment::CombineTempAllocations() {
VLOG(1) << "CombineTempAllocations()";
- FlatMap<LogicalBuffer::Color, BufferAllocation, LogicalBuffer::Color::Hasher>
+ flat_hash_map<LogicalBuffer::Color, BufferAllocation,
+ LogicalBuffer::Color::Hasher>
combined_allocation_map;
// Move all temp allocations into a single run at the end of the allocations
@@ -582,7 +584,8 @@
}
// Update allocation indices to their new positions.
- allocation_index_for_buffer_.clear_no_resize();
+ allocation_index_for_buffer_.erase(allocation_index_for_buffer_.begin(),
+ allocation_index_for_buffer_.end());
for (size_t index = 0; index < allocations_.size(); ++index) {
BufferAllocation* allocation = &allocations_[index];
allocation->set_index(index);
@@ -814,7 +817,7 @@
const HloComputation* computation, bool is_thread_local,
const FlatSet<const LogicalBuffer*>& colocated_buffers,
const FlatSet<BufferAllocation::Index>& colocated_allocations,
- FlatMap<const HloComputation*, FlatSet<const LogicalBuffer*>>*
+ flat_hash_map<const HloComputation*, FlatSet<const LogicalBuffer*>>*
buffers_to_assign_sequentially,
BufferAssignment* assignment) {
// Buffers are sorted and assigned to BufferAllocations in decreasing order of
@@ -833,7 +836,7 @@
// Generate a post order sort of instructions for sorting of the
// LogicalBuffers.
- FlatMap<const HloInstruction*, int> post_order_position;
+ flat_hash_map<const HloInstruction*, int> post_order_position;
int position = 0;
for (auto* instruction : computation->MakeInstructionPostOrder()) {
post_order_position.emplace(instruction, position);
@@ -1043,12 +1046,12 @@
return Status::OK();
}
-FlatMap<LogicalBuffer::Color, FlatSet<const LogicalBuffer*>,
- LogicalBuffer::Color::Hasher>
+flat_hash_map<LogicalBuffer::Color, FlatSet<const LogicalBuffer*>,
+ LogicalBuffer::Color::Hasher>
BufferAssigner::SplitBuffersByColor(
const FlatSet<const LogicalBuffer*>& buffers) {
- FlatMap<LogicalBuffer::Color, FlatSet<const LogicalBuffer*>,
- LogicalBuffer::Color::Hasher>
+ flat_hash_map<LogicalBuffer::Color, FlatSet<const LogicalBuffer*>,
+ LogicalBuffer::Color::Hasher>
color_map;
for (auto buffer : buffers) {
color_map[buffer->color()].insert(buffer);
@@ -1057,7 +1060,7 @@
}
Status BufferAssigner::AssignBuffersWithSequentialOrdering(
- const FlatMap<const HloComputation*, FlatSet<const LogicalBuffer*>>&
+ const flat_hash_map<const HloComputation*, FlatSet<const LogicalBuffer*>>&
buffers_to_assign_sequentially,
bool run_whole_module_heap_simulation, BufferAssignment* assignment) {
// Run the sequence of instructions through the heap simulator. The heuristic
@@ -1155,9 +1158,8 @@
const BufferAllocation& allocation, const HeapSimulatorTrace& heap_trace) {
// Create a map from LogicalBuffer::Id to LogicalBuffer* for the logical
// buffers in this allocation.
- tensorflow::gtl::FlatMap<LogicalBuffer::Id, const LogicalBuffer*>
- id_to_buffer;
- tensorflow::gtl::FlatMap<const LogicalBuffer*, int64> buffer_sizes;
+ absl::flat_hash_map<LogicalBuffer::Id, const LogicalBuffer*> id_to_buffer;
+ absl::flat_hash_map<const LogicalBuffer*, int64> buffer_sizes;
for (const auto& pair : allocation.assigned_buffers()) {
const LogicalBuffer* buffer = pair.first;
const BufferAllocation::OffsetSize& offset_size = pair.second;
@@ -1679,7 +1681,7 @@
// First assign buffers for global computatations. Temporary buffers for
// sequential computations are collected in 'buffers_to_assign_sequentially'.
- FlatMap<const HloComputation*, FlatSet<const LogicalBuffer*>>
+ flat_hash_map<const HloComputation*, FlatSet<const LogicalBuffer*>>
buffers_to_assign_sequentially;
for (auto* computation : global_computations) {
TF_RETURN_IF_ERROR(AssignBuffersForComputation(
diff --git a/tensorflow/compiler/xla/service/buffer_assignment.h b/tensorflow/compiler/xla/service/buffer_assignment.h
index 24ba7c1..9ba4061 100644
--- a/tensorflow/compiler/xla/service/buffer_assignment.h
+++ b/tensorflow/compiler/xla/service/buffer_assignment.h
@@ -22,6 +22,7 @@
#include <string>
#include <vector>
+#include "absl/container/flat_hash_map.h"
#include "absl/types/span.h"
#include "tensorflow/compiler/xla/service/buffer_liveness.h"
#include "tensorflow/compiler/xla/service/heap_simulator.h"
@@ -33,7 +34,6 @@
#include "tensorflow/compiler/xla/service/tuple_points_to_analysis.h"
#include "tensorflow/compiler/xla/statusor.h"
#include "tensorflow/compiler/xla/types.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/lib/gtl/flatset.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/macros.h"
@@ -148,7 +148,7 @@
// Access to the logical buffers assigned to this allocation, and their
// associated logical offsets and sizes.
- const tensorflow::gtl::FlatMap<const LogicalBuffer*, OffsetSize>&
+ const absl::flat_hash_map<const LogicalBuffer*, OffsetSize>&
assigned_buffers() const {
return assigned_buffers_;
}
@@ -323,7 +323,7 @@
// Mapping from the set of buffers assigned to this allocation to their
// logical offsets and sizes.
- tensorflow::gtl::FlatMap<const LogicalBuffer*, OffsetSize> assigned_buffers_;
+ absl::flat_hash_map<const LogicalBuffer*, OffsetSize> assigned_buffers_;
int64 fragmentation_bytes_ = 0;
std::vector<HeapSimulatorTrace> heap_traces_;
@@ -500,7 +500,7 @@
int64 temp_allocation_total_size_ = 0;
// Maps Buffers to the index of the BufferAllocation which holds the buffer.
- tensorflow::gtl::FlatMap<const LogicalBuffer*, BufferAllocation::Index>
+ absl::flat_hash_map<const LogicalBuffer*, BufferAllocation::Index>
allocation_index_for_buffer_;
const HloModule* module_;
@@ -557,8 +557,8 @@
const tensorflow::gtl::FlatSet<const LogicalBuffer*>& colocated_buffers,
const tensorflow::gtl::FlatSet<BufferAllocation::Index>&
colocated_allocations,
- tensorflow::gtl::FlatMap<const HloComputation*,
- tensorflow::gtl::FlatSet<const LogicalBuffer*>>*
+ absl::flat_hash_map<const HloComputation*,
+ tensorflow::gtl::FlatSet<const LogicalBuffer*>>*
buffers_to_assign_sequentially,
BufferAssignment* assignment);
@@ -568,9 +568,8 @@
// 'run_whole_module_heap_simulation' is true, the heap simulation will be run
// assuming all global computations are sequentially ordered.
Status AssignBuffersWithSequentialOrdering(
- const tensorflow::gtl::FlatMap<
- const HloComputation*,
- tensorflow::gtl::FlatSet<const LogicalBuffer*>>&
+ const absl::flat_hash_map<const HloComputation*,
+ tensorflow::gtl::FlatSet<const LogicalBuffer*>>&
buffers_to_assign_sequentially,
bool run_whole_module_heap_simulation, BufferAssignment* assignment);
@@ -624,9 +623,9 @@
// Split a set of buffers into several sets, each of which contains buffers
// colored with the same color.
- tensorflow::gtl::FlatMap<LogicalBuffer::Color,
- tensorflow::gtl::FlatSet<const LogicalBuffer*>,
- LogicalBuffer::Color::Hasher>
+ absl::flat_hash_map<LogicalBuffer::Color,
+ tensorflow::gtl::FlatSet<const LogicalBuffer*>,
+ LogicalBuffer::Color::Hasher>
SplitBuffersByColor(
const tensorflow::gtl::FlatSet<const LogicalBuffer*>& buffers);
diff --git a/tensorflow/compiler/xla/service/buffer_liveness.h b/tensorflow/compiler/xla/service/buffer_liveness.h
index cdd3cf4..2911bbc 100644
--- a/tensorflow/compiler/xla/service/buffer_liveness.h
+++ b/tensorflow/compiler/xla/service/buffer_liveness.h
@@ -27,7 +27,6 @@
#include "tensorflow/compiler/xla/statusor.h"
#include "tensorflow/compiler/xla/types.h"
#include "tensorflow/core/lib/core/status.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/lib/gtl/flatset.h"
namespace xla {
diff --git a/tensorflow/compiler/xla/service/call_graph.h b/tensorflow/compiler/xla/service/call_graph.h
index 3af2ab5..0c2e9b9 100644
--- a/tensorflow/compiler/xla/service/call_graph.h
+++ b/tensorflow/compiler/xla/service/call_graph.h
@@ -20,10 +20,10 @@
#include <ostream>
+#include "absl/container/flat_hash_map.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"
-#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/lib/gtl/flatset.h"
namespace xla {
@@ -157,7 +157,7 @@
// The map from instruction to index in callsites_ for looking up the callsite
// (if any) associated with a particular instruction in this computation.
- tensorflow::gtl::FlatMap<const HloInstruction*, int64> callsite_instructions_;
+ absl::flat_hash_map<const HloInstruction*, int64> callsite_instructions_;
// The call sites in other computations which call this computation.
std::vector<CallSite> caller_callsites_;
@@ -267,7 +267,7 @@
// Map from HLO computation to the index of the corresponding call graph node
// in nodes_.
- tensorflow::gtl::FlatMap<const HloComputation*, int64> node_indices_;
+ absl::flat_hash_map<const HloComputation*, int64> node_indices_;
};
} // namespace xla
diff --git a/tensorflow/compiler/xla/service/copy_insertion.cc b/tensorflow/compiler/xla/service/copy_insertion.cc
index b65dfef..7f78412 100644
--- a/tensorflow/compiler/xla/service/copy_insertion.cc
+++ b/tensorflow/compiler/xla/service/copy_insertion.cc
@@ -15,6 +15,7 @@
#include "tensorflow/compiler/xla/service/copy_insertion.h"
+#include "absl/container/flat_hash_map.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_join.h"
#include "tensorflow/compiler/xla/service/hlo_alias_analysis.h"
@@ -31,7 +32,6 @@
#include "tensorflow/compiler/xla/statusor.h"
#include "tensorflow/compiler/xla/types.h"
#include "tensorflow/compiler/xla/util.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/lib/gtl/flatset.h"
#include "tensorflow/core/platform/logging.h"
@@ -432,7 +432,7 @@
// Construct a list for each HLO buffer in the alias analysis. Maintain a
// map from HloValue to the respective list element representing that
// value. The map is used to construct the copy info map below.
- tensorflow::gtl::FlatMap<const HloValue*, ValueNode*> value_to_node;
+ absl::flat_hash_map<const HloValue*, ValueNode*> value_to_node;
for (const HloBuffer& buffer : alias_analysis.buffers()) {
// Verify values contained in the buffer are strictly ordered. This
// should always be the case after adding copies to eliminate
@@ -480,7 +480,7 @@
// respective ValueNode representing that value.
void AddValueList(
absl::Span<const HloValue* const> values,
- tensorflow::gtl::FlatMap<const HloValue*, ValueNode*>* value_to_node) {
+ absl::flat_hash_map<const HloValue*, ValueNode*>* value_to_node) {
ValueNode* tail = nullptr;
ValueNode* head = nullptr;
for (const HloValue* value : values) {
@@ -516,8 +516,7 @@
// respective ValueNode.
void CreateCopyMap(
const HloModule& module,
- const tensorflow::gtl::FlatMap<const HloValue*, ValueNode*>&
- value_to_node) {
+ const absl::flat_hash_map<const HloValue*, ValueNode*>& value_to_node) {
for (HloComputation* computation : module.computations()) {
for (HloInstruction* instruction : computation->instructions()) {
// Add copies with unambiguous source values to the map. Copies with
@@ -916,7 +915,7 @@
ValueNode* src = nullptr;
ValueNode* dest = nullptr;
};
- tensorflow::gtl::FlatMap<const HloInstruction*, CopyNodes> copy_map_;
+ absl::flat_hash_map<const HloInstruction*, CopyNodes> copy_map_;
};
HloModule* module_;
diff --git a/tensorflow/compiler/xla/service/cpu/BUILD b/tensorflow/compiler/xla/service/cpu/BUILD
index b710311..6a83909 100644
--- a/tensorflow/compiler/xla/service/cpu/BUILD
+++ b/tensorflow/compiler/xla/service/cpu/BUILD
@@ -290,6 +290,7 @@
"//tensorflow/compiler/xla/service/llvm_ir:loop_emitter",
"//tensorflow/compiler/xla/service/llvm_ir:tuple_ops",
"//tensorflow/core:lib",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/types:span",
@@ -309,6 +310,7 @@
deps = [
"//tensorflow/compiler/xla:shape_util",
"//tensorflow/core:lib",
+ "@com_google_absl//absl/container:flat_hash_map",
"@llvm//:analysis",
"@llvm//:target",
],
@@ -471,6 +473,7 @@
"//tensorflow/compiler/xla/service/llvm_ir:llvm_util",
"//tensorflow/core:lib",
"//tensorflow/stream_executor",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/synchronization",
"@com_google_absl//absl/types:span",
],
@@ -762,6 +765,7 @@
"//tensorflow/compiler/xla/service:computation_layout",
"//tensorflow/compiler/xla/service:layout_assignment",
"//tensorflow/core:lib",
+ "@com_google_absl//absl/container:flat_hash_map",
],
)
diff --git a/tensorflow/compiler/xla/service/cpu/cpu_layout_assignment.cc b/tensorflow/compiler/xla/service/cpu/cpu_layout_assignment.cc
index bfecbd6..c291bf2 100644
--- a/tensorflow/compiler/xla/service/cpu/cpu_layout_assignment.cc
+++ b/tensorflow/compiler/xla/service/cpu/cpu_layout_assignment.cc
@@ -17,6 +17,7 @@
#include <numeric>
+#include "absl/container/flat_hash_map.h"
#include "tensorflow/compiler/xla/map_util.h"
#include "tensorflow/compiler/xla/service/cpu/dot_op_emitter.h"
#include "tensorflow/compiler/xla/service/cpu/ir_emission_utils.h"
@@ -38,7 +39,7 @@
using absl::optional;
using ShouldMakeOperandColMajorCache =
- tensorflow::gtl::FlatMap<const HloInstruction*, bool>;
+ absl::flat_hash_map<const HloInstruction*, bool>;
} // namespace
static bool ShouldMakeAllUsersColMajor(const HloInstruction* instruction) {
diff --git a/tensorflow/compiler/xla/service/cpu/cpu_runtime.cc b/tensorflow/compiler/xla/service/cpu/cpu_runtime.cc
index 20cf855..a9febe8 100644
--- a/tensorflow/compiler/xla/service/cpu/cpu_runtime.cc
+++ b/tensorflow/compiler/xla/service/cpu/cpu_runtime.cc
@@ -17,6 +17,7 @@
#include <functional>
+#include "absl/container/flat_hash_map.h"
#include "absl/synchronization/mutex.h"
#include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h"
#include "tensorflow/core/platform/dynamic_annotations.h"
@@ -30,8 +31,7 @@
namespace runtime {
XfeedManager* GetXfeedManager(int device_ordinal) {
- static tensorflow::gtl::FlatMap<int, XfeedManager*>* managers =
- new tensorflow::gtl::FlatMap<int, XfeedManager*>();
+ static auto* managers = new absl::flat_hash_map<int, XfeedManager*>();
static absl::Mutex* mutex = new absl::Mutex();
absl::MutexLock lock(mutex);
diff --git a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc
index c3e8020..953a75c 100644
--- a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc
+++ b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc
@@ -24,6 +24,7 @@
#include <utility>
#include <vector>
+#include "absl/container/flat_hash_map.h"
#include "tensorflow/core/lib/math/math_util.h"
#include "tensorflow/core/platform/logging.h"
// IWYU pragma: no_include "llvm/IR/Intrinsics.gen.inc"
@@ -67,7 +68,6 @@
#include "tensorflow/compiler/xla/window_util.h"
#include "tensorflow/core/lib/core/bits.h"
#include "tensorflow/core/lib/core/errors.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/lib/gtl/flatset.h"
namespace xla {
@@ -1398,7 +1398,7 @@
//
// So if we reduce f32[A,B,C,D] on dimensions 1 and 2, this map contains
// [0->0, 3->1].
- gtl::FlatMap<int64, int64> unreduced_dim_map;
+ absl::flat_hash_map<int64, int64> unreduced_dim_map;
gtl::FlatSet<int64> reduced_dims(reduce.dimensions().begin(),
reduce.dimensions().end());
diff --git a/tensorflow/compiler/xla/service/cpu/ir_emitter.h b/tensorflow/compiler/xla/service/cpu/ir_emitter.h
index daafef4..586f27b 100644
--- a/tensorflow/compiler/xla/service/cpu/ir_emitter.h
+++ b/tensorflow/compiler/xla/service/cpu/ir_emitter.h
@@ -23,6 +23,7 @@
#include <unordered_map>
#include <vector>
+#include "absl/container/flat_hash_map.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "llvm/ADT/Triple.h"
@@ -47,7 +48,6 @@
#include "tensorflow/compiler/xla/statusor.h"
#include "tensorflow/compiler/xla/types.h"
#include "tensorflow/compiler/xla/xla_data.pb.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/types.h"
@@ -427,7 +427,7 @@
// Maps the buffer allocation slices for the parameters to the computation
// being compiled to their parameter numbers. Only relevant for thread local
// computations.
- tensorflow::gtl::FlatMap<BufferAllocation::Index, int64>
+ absl::flat_hash_map<BufferAllocation::Index, int64>
computation_parameter_allocations_;
// Maps HLO instructions to their index into the profile counter array.
@@ -567,11 +567,11 @@
}
};
- tensorflow::gtl::FlatMap<const Literal*, llvm::Constant*,
- LiteralPtrHashFunctor, LiteralPtrEqualityFunctor>
+ absl::flat_hash_map<const Literal*, llvm::Constant*, LiteralPtrHashFunctor,
+ LiteralPtrEqualityFunctor>
emitted_literals_;
- tensorflow::gtl::FlatMap<BufferAllocation::Index, llvm::Constant*>
+ absl::flat_hash_map<BufferAllocation::Index, llvm::Constant*>
constant_buffer_to_global_;
std::vector<const HloComputation*> thread_local_computations_;
diff --git a/tensorflow/compiler/xla/service/cpu/target_machine_features.cc b/tensorflow/compiler/xla/service/cpu/target_machine_features.cc
index a0cd8ee..5cdac20 100644
--- a/tensorflow/compiler/xla/service/cpu/target_machine_features.cc
+++ b/tensorflow/compiler/xla/service/cpu/target_machine_features.cc
@@ -14,6 +14,7 @@
==============================================================================*/
#include "tensorflow/compiler/xla/service/cpu/target_machine_features.h"
+#include "tensorflow/core/platform/logging.h"
namespace xla {
namespace cpu {
diff --git a/tensorflow/compiler/xla/service/cpu/target_machine_features.h b/tensorflow/compiler/xla/service/cpu/target_machine_features.h
index 8b00ae9..a383b4a 100644
--- a/tensorflow/compiler/xla/service/cpu/target_machine_features.h
+++ b/tensorflow/compiler/xla/service/cpu/target_machine_features.h
@@ -16,10 +16,10 @@
#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_TARGET_MACHINE_FEATURES_H_
#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_TARGET_MACHINE_FEATURES_H_
+#include "absl/container/flat_hash_map.h"
#include "llvm/Analysis/TargetTransformInfo.h"
#include "llvm/Target/TargetMachine.h"
#include "tensorflow/compiler/xla/primitive_util.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
namespace xla {
namespace cpu {
@@ -97,8 +97,7 @@
// This is mutated from within `GetTargetTransformInfoFor` which is
// semantically a getter (and thus `const`); and is therefore declared
// mutable. Making this mutable is okay because it has cache semantics.
- mutable tensorflow::gtl::FlatMap<const llvm::Function*,
- llvm::TargetTransformInfo>
+ mutable absl::flat_hash_map<const llvm::Function*, llvm::TargetTransformInfo>
target_transform_info_cache_;
llvm::TargetMachine* target_machine_;
};
diff --git a/tensorflow/compiler/xla/service/defuser.cc b/tensorflow/compiler/xla/service/defuser.cc
index d124f74..661539c 100644
--- a/tensorflow/compiler/xla/service/defuser.cc
+++ b/tensorflow/compiler/xla/service/defuser.cc
@@ -22,6 +22,7 @@
#include <utility>
#include <vector>
+#include "absl/container/flat_hash_map.h"
#include "tensorflow/compiler/xla/service/call_graph.h"
#include "tensorflow/compiler/xla/service/hlo_computation.h"
#include "tensorflow/compiler/xla/service/hlo_instruction.h"
@@ -48,7 +49,7 @@
fusion_instruction->fused_instructions_computation();
// A map from fused instruction to its defused clone.
- tensorflow::gtl::FlatMap<const HloInstruction*, HloInstruction*>
+ absl::flat_hash_map<const HloInstruction*, HloInstruction*>
defused_instructions;
// Initialize map to contain the fusion instruction parameters mapping
// to the operands of the fusion instruction.
diff --git a/tensorflow/compiler/xla/service/dfs_hlo_visitor.h b/tensorflow/compiler/xla/service/dfs_hlo_visitor.h
index 5761573..68d01d7 100644
--- a/tensorflow/compiler/xla/service/dfs_hlo_visitor.h
+++ b/tensorflow/compiler/xla/service/dfs_hlo_visitor.h
@@ -27,7 +27,6 @@
#include "tensorflow/compiler/xla/types.h"
#include "tensorflow/compiler/xla/xla_data.pb.h"
#include "tensorflow/core/lib/core/status.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/types.h"
diff --git a/tensorflow/compiler/xla/service/gpu/BUILD b/tensorflow/compiler/xla/service/gpu/BUILD
index 51968d1..e65d3fa 100644
--- a/tensorflow/compiler/xla/service/gpu/BUILD
+++ b/tensorflow/compiler/xla/service/gpu/BUILD
@@ -91,6 +91,7 @@
"//tensorflow/compiler/xla/service:hlo",
"//tensorflow/compiler/xla/service:hlo_reachability",
"//tensorflow/core:lib",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/memory",
],
)
@@ -357,6 +358,7 @@
"//tensorflow/core/platform/default/build_config:cufft_plugin",
"//tensorflow/core/platform/default/build_config:stream_executor_cuda", # build_cleaner: keep
"//tensorflow/stream_executor",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
diff --git a/tensorflow/compiler/xla/service/gpu/gpu_executable.cc b/tensorflow/compiler/xla/service/gpu/gpu_executable.cc
index 31a9f9b..5742632 100644
--- a/tensorflow/compiler/xla/service/gpu/gpu_executable.cc
+++ b/tensorflow/compiler/xla/service/gpu/gpu_executable.cc
@@ -19,6 +19,7 @@
#include <utility>
#include <vector>
+#include "absl/container/flat_hash_map.h"
#include "absl/memory/memory.h"
#include "tensorflow/compiler/xla/map_util.h"
#include "tensorflow/compiler/xla/service/gpu/buffer_allocations.h"
@@ -197,7 +198,7 @@
}
module_spec.AddCudaPtxInMemory(ptx().c_str());
- tensorflow::gtl::FlatMap<int64, se::DeviceMemoryBase> globals;
+ absl::flat_hash_map<int64, se::DeviceMemoryBase> globals;
se::ModuleHandle module_handle;
executor->LoadModule(module_spec, &module_handle);
diff --git a/tensorflow/compiler/xla/service/gpu/gpu_executable.h b/tensorflow/compiler/xla/service/gpu/gpu_executable.h
index 38b0f8f..0e27628 100644
--- a/tensorflow/compiler/xla/service/gpu/gpu_executable.h
+++ b/tensorflow/compiler/xla/service/gpu/gpu_executable.h
@@ -19,6 +19,7 @@
#include <memory>
#include <string>
+#include "absl/container/flat_hash_map.h"
#include "absl/strings/string_view.h"
#include "absl/types/optional.h"
#include "absl/types/span.h"
@@ -35,7 +36,6 @@
#include "tensorflow/compiler/xla/service/tuple_points_to_analysis.h"
#include "tensorflow/compiler/xla/statusor.h"
#include "tensorflow/compiler/xla/types.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/stream_executor_no_cuda.h"
@@ -101,7 +101,7 @@
const PointsToSet& GetRootPointsToSet() const;
using BufferAllocToDeviceMemoryMap =
- tensorflow::gtl::FlatMap<BufferAllocation::Index, se::DeviceMemoryBase>;
+ absl::flat_hash_map<BufferAllocation::Index, se::DeviceMemoryBase>;
// Loads the PTX or CUBIN for this executable into `executor` and resolves the
// globals corresponding to constant buffers. Returns a map mapping buffer
diff --git a/tensorflow/compiler/xla/service/gpu/stream_assignment.h b/tensorflow/compiler/xla/service/gpu/stream_assignment.h
index c2df83a..52d38b6 100644
--- a/tensorflow/compiler/xla/service/gpu/stream_assignment.h
+++ b/tensorflow/compiler/xla/service/gpu/stream_assignment.h
@@ -16,9 +16,9 @@
#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_STREAM_ASSIGNMENT_H_
#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_STREAM_ASSIGNMENT_H_
+#include "absl/container/flat_hash_map.h"
#include "tensorflow/compiler/xla/service/hlo_instruction.h"
#include "tensorflow/compiler/xla/service/hlo_module.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
namespace xla {
namespace gpu {
@@ -34,7 +34,7 @@
private:
int stream_count_ = 1; // At least the main stream.
- tensorflow::gtl::FlatMap<const HloInstruction*, int> hlo_to_stream_number_;
+ absl::flat_hash_map<const HloInstruction*, int> hlo_to_stream_number_;
};
// Assigns GPU streams to instructions in `module`.
diff --git a/tensorflow/compiler/xla/service/heap_simulator.cc b/tensorflow/compiler/xla/service/heap_simulator.cc
index 2bd0425..147776c 100644
--- a/tensorflow/compiler/xla/service/heap_simulator.cc
+++ b/tensorflow/compiler/xla/service/heap_simulator.cc
@@ -18,13 +18,14 @@
#include <algorithm>
#include <vector>
+#include "absl/container/flat_hash_map.h"
#include "absl/memory/memory.h"
#include "tensorflow/compiler/xla/map_util.h"
#include "tensorflow/compiler/xla/util.h"
namespace xla {
-using tensorflow::gtl::FlatMap;
+using absl::flat_hash_map;
using tensorflow::gtl::FlatSet;
/*static*/
@@ -56,7 +57,7 @@
const HloComputation& computation, const HloInstructionSequence& sequence,
const TuplePointsToAnalysis& points_to_analysis,
const LogicalBuffer::SizeFunction& size_function,
- const tensorflow::gtl::FlatMap<const HloComputation*, int64>*
+ const absl::flat_hash_map<const HloComputation*, int64>*
memory_by_computation) {
TF_ASSIGN_OR_RETURN(
HeapSimulator::Result result,
@@ -88,7 +89,7 @@
const HloInstructionSequence& instruction_sequence,
const TuplePointsToAnalysis& points_to_analysis,
const BufferValue::SizeFunction& size_fn, const Options& options,
- const tensorflow::gtl::FlatMap<const HloComputation*, int64>*
+ const absl::flat_hash_map<const HloComputation*, int64>*
memory_by_computation) {
HeapSimulator heap(std::move(algorithm), size_fn, options,
/*schedule=*/nullptr, memory_by_computation);
@@ -115,8 +116,10 @@
// 'used_buffers' is the reverse map - it tracks which buffers were used by an
// instruction, so that we can remove the instructions from a buffer's live
// set after they are visited.
- FlatMap<const BufferValue*, FlatSet<const HloInstruction*>> live_buffers;
- FlatMap<const HloInstruction*, FlatSet<const BufferValue*>> used_buffers;
+ flat_hash_map<const BufferValue*, FlatSet<const HloInstruction*>>
+ live_buffers;
+ flat_hash_map<const HloInstruction*, FlatSet<const BufferValue*>>
+ used_buffers;
auto add_user_to_buffer = [this, &live_buffers, &used_buffers](
const HloInstruction* user,
const BufferValue* buffer) {
@@ -345,7 +348,7 @@
std::unique_ptr<HeapAlgorithm> algorithm,
const BufferValue::SizeFunction& size_fn, const Options& options,
const HloSchedule* schedule,
- const tensorflow::gtl::FlatMap<const HloComputation*, int64>*
+ const absl::flat_hash_map<const HloComputation*, int64>*
memory_by_computation)
: no_fragmentation_stats_(absl::make_unique<NoFragmentationStatsHeap>()),
algorithm_(std::move(algorithm)),
@@ -536,7 +539,7 @@
void NoFragmentationStatsHeap::AccountForSubcomputationMemory(
const HloInstruction* instruction,
- const tensorflow::gtl::FlatMap<const HloComputation*, int64>&
+ const absl::flat_hash_map<const HloComputation*, int64>&
memory_by_computation) {
// We only count the memory usage of the largest subcomputation, instead of
// adding them all, because subcomputations won't execute in parallel.
diff --git a/tensorflow/compiler/xla/service/heap_simulator.h b/tensorflow/compiler/xla/service/heap_simulator.h
index 7d6dcc0..a5bb3f8 100644
--- a/tensorflow/compiler/xla/service/heap_simulator.h
+++ b/tensorflow/compiler/xla/service/heap_simulator.h
@@ -21,6 +21,7 @@
#include <utility>
#include <vector>
+#include "absl/container/flat_hash_map.h"
#include "tensorflow/compiler/xla/service/buffer_value.h"
#include "tensorflow/compiler/xla/service/buffer_value_containers.h"
#include "tensorflow/compiler/xla/service/hlo.pb.h"
@@ -30,7 +31,6 @@
#include "tensorflow/compiler/xla/service/hlo_schedule.h"
#include "tensorflow/compiler/xla/service/tuple_points_to_analysis.h"
#include "tensorflow/compiler/xla/statusor.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/lib/gtl/flatset.h"
namespace xla {
@@ -58,7 +58,7 @@
// Result represents the result of the heap simulation.
struct Result {
// The assignment of buffers to chunks.
- tensorflow::gtl::FlatMap<const BufferValue*, Chunk> chunk_map;
+ absl::flat_hash_map<const BufferValue*, Chunk> chunk_map;
// The total size in bytes of the heap, containing all assigned chunks.
int64 heap_size = 0;
@@ -100,7 +100,7 @@
const HloComputation& computation, const HloInstructionSequence& sequence,
const TuplePointsToAnalysis& points_to_analysis,
const LogicalBuffer::SizeFunction& size_function,
- const tensorflow::gtl::FlatMap<const HloComputation*, int64>*
+ const absl::flat_hash_map<const HloComputation*, int64>*
memory_by_computation = nullptr);
// Run the heap simulation with the given algorithm, assuming the given
@@ -130,7 +130,7 @@
const TuplePointsToAnalysis& points_to_analysis,
const BufferValue::SizeFunction& size_fn,
const Options& options = Options(),
- const tensorflow::gtl::FlatMap<const HloComputation*, int64>*
+ const absl::flat_hash_map<const HloComputation*, int64>*
memory_by_computation = nullptr);
private:
@@ -140,7 +140,7 @@
HeapSimulator(std::unique_ptr<HeapAlgorithm> algorithm,
const BufferValue::SizeFunction& size_fn,
const Options& options, const HloSchedule* schedule = nullptr,
- const tensorflow::gtl::FlatMap<const HloComputation*, int64>*
+ const absl::flat_hash_map<const HloComputation*, int64>*
memory_by_computation = nullptr);
~HeapSimulator();
@@ -172,7 +172,7 @@
// handle subcomputations. It would be good to unify the handling of
// subcomputations, but it's not clear how.
const HloSchedule* schedule_;
- const tensorflow::gtl::FlatMap<const HloComputation*, int64>*
+ const absl::flat_hash_map<const HloComputation*, int64>*
memory_by_computation_;
// In addition to Alloc and Free, the heap simulator exposes a concept of
@@ -193,7 +193,7 @@
const BufferValue* canonical = nullptr;
int64 refcount = 0;
};
- tensorflow::gtl::FlatMap<const BufferValue*, std::shared_ptr<SharedGroup>>
+ absl::flat_hash_map<const BufferValue*, std::shared_ptr<SharedGroup>>
shared_buffers_;
// Hold some sets for error-checking the sequence of Alloc and Free calls.
@@ -235,7 +235,7 @@
// analysis, it's not worth making major changes to HeapSimulator now.
virtual void AccountForSubcomputationMemory(
const HloInstruction* instruction,
- const tensorflow::gtl::FlatMap<const HloComputation*, int64>&
+ const absl::flat_hash_map<const HloComputation*, int64>&
memory_by_computation) {}
// Free de-allocates a previously allocated buffer.
@@ -262,7 +262,7 @@
void AccountForSubcomputationMemory(
const HloInstruction* instruction,
- const tensorflow::gtl::FlatMap<const HloComputation*, int64>&
+ const absl::flat_hash_map<const HloComputation*, int64>&
memory_by_computation) override;
void Free(const BufferValue* buffer, int64 size) override;
@@ -382,8 +382,7 @@
// Free time of the buffer.
int64 end;
};
- tensorflow::gtl::FlatMap<const BufferValue*, BufferInterval>
- buffer_intervals_;
+ absl::flat_hash_map<const BufferValue*, BufferInterval> buffer_intervals_;
};
// A heap algorithm that chooses the best results from other algorithms added to
diff --git a/tensorflow/compiler/xla/service/heap_simulator_test.cc b/tensorflow/compiler/xla/service/heap_simulator_test.cc
index 191fbf8..ea0bced 100644
--- a/tensorflow/compiler/xla/service/heap_simulator_test.cc
+++ b/tensorflow/compiler/xla/service/heap_simulator_test.cc
@@ -19,6 +19,7 @@
#include <utility>
#include <vector>
+#include "absl/container/flat_hash_map.h"
#include "absl/memory/memory.h"
#include "tensorflow/compiler/xla/literal.h"
#include "tensorflow/compiler/xla/service/buffer_value.h"
@@ -31,7 +32,6 @@
#include "tensorflow/compiler/xla/status_macros.h"
#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.h"
#include "tensorflow/core/lib/core/status_test_util.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
namespace xla {
namespace {
@@ -174,7 +174,7 @@
// Construct the module sequence grouped by computation.
HloSchedule schedule(module_.get());
- tensorflow::gtl::FlatMap<const HloInstruction*, int> reverse_position;
+ absl::flat_hash_map<const HloInstruction*, int> reverse_position;
for (int i = 0; i < full_module_sequence.size(); ++i) {
const HloInstruction* instruction = full_module_sequence[i];
schedule.GetOrCreateSequence(instruction->parent())
diff --git a/tensorflow/compiler/xla/service/hlo.proto b/tensorflow/compiler/xla/service/hlo.proto
index b19ec12..caaca16 100644
--- a/tensorflow/compiler/xla/service/hlo.proto
+++ b/tensorflow/compiler/xla/service/hlo.proto
@@ -34,7 +34,7 @@
option cc_enable_arenas = true;
// Serialization of HloInstruction.
-// Next ID: 53
+// Next ID: 54
message HloInstructionProto {
reserved 10;
reserved "parameter_name";
@@ -124,9 +124,13 @@
// The string representation of the infeed configuration.
bytes infeed_config = 27;
- // Name of a global symbol to call, only present for kCustomCall.
+ // Name of a external target (eg, global symbol) to call, only present for
+ // kCustomCall.
string custom_call_target = 28;
+ // Opaque string, only present for kCustomCall.
+ string custom_call_opaque = 53;
+
// Shape of outfeed request.
xla.Shape outfeed_shape = 29;
diff --git a/tensorflow/compiler/xla/service/hlo_alias_analysis.cc b/tensorflow/compiler/xla/service/hlo_alias_analysis.cc
index 0986da6..b6e1f52 100644
--- a/tensorflow/compiler/xla/service/hlo_alias_analysis.cc
+++ b/tensorflow/compiler/xla/service/hlo_alias_analysis.cc
@@ -20,6 +20,7 @@
#include <utility>
#include <vector>
+#include "absl/container/flat_hash_map.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_join.h"
#include "tensorflow/compiler/xla/map_util.h"
@@ -290,13 +291,11 @@
const HloDataflowAnalysis& dataflow_;
// A map containing the set of values contained in each buffer.
- tensorflow::gtl::FlatMap<BufferNumber,
- tensorflow::gtl::FlatSet<const HloValue*>>
+ absl::flat_hash_map<BufferNumber, tensorflow::gtl::FlatSet<const HloValue*>>
buffers_;
// A map indicating which buffer each value is contained in.
- tensorflow::gtl::FlatMap<const HloValue*, BufferNumber>
- value_to_buffer_number_;
+ absl::flat_hash_map<const HloValue*, BufferNumber> value_to_buffer_number_;
// The buffer number of the next buffer to be created.
BufferNumber next_buffer_number_ = 0;
diff --git a/tensorflow/compiler/xla/service/hlo_alias_analysis.h b/tensorflow/compiler/xla/service/hlo_alias_analysis.h
index e345804..372f99f 100644
--- a/tensorflow/compiler/xla/service/hlo_alias_analysis.h
+++ b/tensorflow/compiler/xla/service/hlo_alias_analysis.h
@@ -20,6 +20,7 @@
#include <string>
#include <vector>
+#include "absl/container/flat_hash_map.h"
#include "absl/types/span.h"
#include "tensorflow/compiler/xla/service/hlo_buffer.h"
#include "tensorflow/compiler/xla/service/hlo_dataflow_analysis.h"
@@ -110,7 +111,7 @@
std::unique_ptr<HloDataflowAnalysis> dataflow_analysis_;
// A map indicating which buffer a value is contained in.
- tensorflow::gtl::FlatMap<const HloValue*, HloBuffer*> value_to_buffer_;
+ absl::flat_hash_map<const HloValue*, HloBuffer*> value_to_buffer_;
// A lazily constructed vector containing all HloBuffers sorted by
// HloBuffer::Id.
diff --git a/tensorflow/compiler/xla/service/hlo_clone_context.h b/tensorflow/compiler/xla/service/hlo_clone_context.h
index 658643b..24910ca 100644
--- a/tensorflow/compiler/xla/service/hlo_clone_context.h
+++ b/tensorflow/compiler/xla/service/hlo_clone_context.h
@@ -18,8 +18,8 @@
#include <string>
+#include "absl/container/flat_hash_map.h"
#include "tensorflow/compiler/xla/map_util.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
namespace xla {
@@ -73,12 +73,12 @@
return FindOrDie(computations_, old_computation);
}
- const tensorflow::gtl::FlatMap<const HloInstruction*, HloInstruction*>&
+ const absl::flat_hash_map<const HloInstruction*, HloInstruction*>&
cloned_instructions() const {
return instructions_;
}
- const tensorflow::gtl::FlatMap<const HloComputation*, HloComputation*>&
+ const absl::flat_hash_map<const HloComputation*, HloComputation*>&
cloned_computations() const {
return computations_;
}
@@ -86,10 +86,8 @@
private:
HloModule* module_;
string suffix_;
- tensorflow::gtl::FlatMap<const HloInstruction*, HloInstruction*>
- instructions_;
- tensorflow::gtl::FlatMap<const HloComputation*, HloComputation*>
- computations_;
+ absl::flat_hash_map<const HloInstruction*, HloInstruction*> instructions_;
+ absl::flat_hash_map<const HloComputation*, HloComputation*> computations_;
};
} // namespace xla
diff --git a/tensorflow/compiler/xla/service/hlo_computation.cc b/tensorflow/compiler/xla/service/hlo_computation.cc
index 0e5920a..257dd58 100644
--- a/tensorflow/compiler/xla/service/hlo_computation.cc
+++ b/tensorflow/compiler/xla/service/hlo_computation.cc
@@ -24,6 +24,7 @@
#include <sstream>
#include "absl/algorithm/container.h"
+#include "absl/container/flat_hash_map.h"
#include "absl/memory/memory.h"
#include "absl/strings/numbers.h"
#include "absl/strings/str_cat.h"
@@ -122,30 +123,6 @@
return instructions_.back().get();
}
-namespace {
-
-// Returns the new name for a fusion parameter when we change its number.
-//
-// Fusion parameters are named foo.param_1, bar.param_2, etc. We are
-// renumbering the parameters, so replace the final number in the name with
-// the updated value.
-string RenameFusionParameter(const string& original_name, int64 new_param_no) {
- const string param_underscore = ".param_";
- size_t index = original_name.rfind(param_underscore);
- if (index == string::npos) {
- return original_name;
- }
- string after_param = original_name.substr(index + param_underscore.size());
- int64 numeric_suffix;
- if (absl::SimpleAtoi(after_param, &numeric_suffix)) {
- return StrCat(original_name.substr(0, index + param_underscore.size()),
- new_param_no);
- }
- return original_name;
-}
-
-} // namespace
-
Status HloComputation::RemoveParameter(int64 param_no) {
CHECK_GE(param_no, 0);
CHECK_LT(param_no, param_instructions_.size());
@@ -158,11 +135,9 @@
while (param_no < param_instructions_.size()) {
param_instruction = param_instructions_[param_no];
- string param_name =
- RenameFusionParameter(param_instruction->name(), param_no);
HloInstruction* new_instr =
AddInstructionInternal(HloInstruction::CreateParameter(
- param_no, param_instruction->shape(), param_name));
+ param_no, param_instruction->shape(), StrCat("param_", param_no)));
TF_RETURN_IF_ERROR(param_instruction->ReplaceAllUsesWith(new_instr));
param_instructions_[param_no] = new_instr;
TF_RETURN_IF_ERROR(RemoveInstruction(param_instruction));
@@ -186,11 +161,9 @@
if (removed > 0) {
const int64 param_no = i - removed;
- string param_name =
- RenameFusionParameter(param_instruction->name(), param_no);
- HloInstruction* new_instr =
- AddInstructionInternal(HloInstruction::CreateParameter(
- param_no, param_instruction->shape(), param_name));
+ HloInstruction* new_instr = AddInstructionInternal(
+ HloInstruction::CreateParameter(param_no, param_instruction->shape(),
+ StrCat("param_", param_no)));
TF_RETURN_IF_ERROR(param_instruction->ReplaceAllUsesWith(new_instr));
param_instructions_[param_no] = new_instr;
TF_RETURN_IF_ERROR(RemoveInstruction(param_instruction));
@@ -325,7 +298,7 @@
void HloComputation::ComputeInstructionPostOrder(
const HloComputation::ChannelDependencyMap& channel_dependency_map,
std::vector<HloInstruction*>* post_order, HloInstruction* root,
- tensorflow::gtl::FlatMap<HloInstruction*, VisitState>* visited) const {
+ absl::flat_hash_map<HloInstruction*, VisitState>* visited) const {
std::vector<HloInstruction*> dfs_stack;
dfs_stack.push_back(root);
while (!dfs_stack.empty()) {
@@ -422,7 +395,7 @@
std::vector<HloInstruction*> post_order;
post_order.reserve(instruction_count());
std::vector<HloInstruction*> trace_instructions;
- tensorflow::gtl::FlatMap<HloInstruction*, VisitState> visited;
+ absl::flat_hash_map<HloInstruction*, VisitState> visited;
for (auto& instruction : instructions_) {
if (instruction->opcode() == HloOpcode::kTrace) {
// Trace instructions aren't handled by the DFS visitor. Add trace
@@ -533,9 +506,9 @@
/* static */ StatusOr<std::unique_ptr<HloComputation>>
HloComputation::CreateFromProto(
const HloComputationProto& proto,
- const tensorflow::gtl::FlatMap<int64, HloComputation*>& computation_map) {
- tensorflow::gtl::FlatMap<int64, HloInstruction*> instruction_map;
- tensorflow::gtl::FlatMap<HloInstruction*, int64> to_proto_id;
+ const absl::flat_hash_map<int64, HloComputation*>& computation_map) {
+ absl::flat_hash_map<int64, HloInstruction*> instruction_map;
+ absl::flat_hash_map<HloInstruction*, int64> to_proto_id;
std::vector<std::unique_ptr<HloInstruction>> instructions;
int64 parameter_count = 0;
for (const HloInstructionProto& instruction_proto : proto.instructions()) {
diff --git a/tensorflow/compiler/xla/service/hlo_computation.h b/tensorflow/compiler/xla/service/hlo_computation.h
index 936a53b..af929ac 100644
--- a/tensorflow/compiler/xla/service/hlo_computation.h
+++ b/tensorflow/compiler/xla/service/hlo_computation.h
@@ -25,6 +25,7 @@
#include <utility>
#include <vector>
+#include "absl/container/flat_hash_map.h"
#include "absl/types/span.h"
#include "tensorflow/compiler/xla/iterator_util.h"
#include "tensorflow/compiler/xla/map_util.h"
@@ -40,7 +41,6 @@
#include "tensorflow/compiler/xla/types.h"
#include "tensorflow/compiler/xla/xla_data.pb.h"
#include "tensorflow/core/lib/core/status.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/lib/gtl/flatset.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/types.h"
@@ -188,7 +188,7 @@
// calls.
static StatusOr<std::unique_ptr<HloComputation>> CreateFromProto(
const HloComputationProto& proto,
- const tensorflow::gtl::FlatMap<int64, HloComputation*>& computation_map);
+ const absl::flat_hash_map<int64, HloComputation*>& computation_map);
// Gets the instructions in this computation.
//
@@ -414,14 +414,14 @@
// cross-replica-sum the union of the dependencies for all participating
// instructions.
using ChannelDependencyMap =
- tensorflow::gtl::FlatMap<int64, absl::InlinedVector<HloInstruction*, 1>>;
+ absl::flat_hash_map<int64, absl::InlinedVector<HloInstruction*, 1>>;
ChannelDependencyMap ComputeChannelDependencies() const;
enum VisitState { kVisiting, kVisited };
void ComputeInstructionPostOrder(
const HloComputation::ChannelDependencyMap& channel_dependency_map,
std::vector<HloInstruction*>* post_order, HloInstruction* root,
- tensorflow::gtl::FlatMap<HloInstruction*, VisitState>* visited) const;
+ absl::flat_hash_map<HloInstruction*, VisitState>* visited) const;
string name_;
int64 unique_id_;
@@ -439,7 +439,7 @@
// instruction pointer to location in the list for fast lookup.
using InstructionList = std::list<std::unique_ptr<HloInstruction>>;
InstructionList instructions_;
- tensorflow::gtl::FlatMap<const HloInstruction*, InstructionList::iterator>
+ absl::flat_hash_map<const HloInstruction*, InstructionList::iterator>
instruction_iterators_;
std::vector<HloInstruction*> param_instructions_;
diff --git a/tensorflow/compiler/xla/service/hlo_domain_map.cc b/tensorflow/compiler/xla/service/hlo_domain_map.cc
index 113fd18..159c39d 100644
--- a/tensorflow/compiler/xla/service/hlo_domain_map.cc
+++ b/tensorflow/compiler/xla/service/hlo_domain_map.cc
@@ -17,6 +17,7 @@
#include <algorithm>
+#include "absl/container/flat_hash_map.h"
#include "absl/memory/memory.h"
#include "tensorflow/compiler/xla/map_util.h"
#include "tensorflow/compiler/xla/service/hlo_opcode.h"
@@ -106,8 +107,8 @@
auto equal = [](const DomainMetadata* a, const DomainMetadata* b) {
return a->Matches(*b);
};
- tensorflow::gtl::FlatMap<const DomainMetadata*, int64, decltype(hash),
- decltype(equal)>
+ absl::flat_hash_map<const DomainMetadata*, int64, decltype(hash),
+ decltype(equal)>
domain_metadata(1024, hash, equal);
for (auto& domain : instruction_domains_) {
diff --git a/tensorflow/compiler/xla/service/hlo_domain_map.h b/tensorflow/compiler/xla/service/hlo_domain_map.h
index 56b557d..8584bc0 100644
--- a/tensorflow/compiler/xla/service/hlo_domain_map.h
+++ b/tensorflow/compiler/xla/service/hlo_domain_map.h
@@ -19,13 +19,13 @@
#include <memory>
#include <vector>
+#include "absl/container/flat_hash_map.h"
#include "tensorflow/compiler/xla/service/hlo_computation.h"
#include "tensorflow/compiler/xla/service/hlo_domain_metadata.h"
#include "tensorflow/compiler/xla/service/hlo_instruction.h"
#include "tensorflow/compiler/xla/service/hlo_module.h"
#include "tensorflow/compiler/xla/statusor.h"
#include "tensorflow/core/lib/core/status.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/lib/gtl/flatset.h"
namespace xla {
@@ -77,8 +77,7 @@
private:
// Map used for representing instruction ordering, i.e.
// order_map[a] < order_map[b] means a must be ordered before b.
- using InstructionOrderMap =
- tensorflow::gtl::FlatMap<const HloInstruction*, int64>;
+ using InstructionOrderMap = absl::flat_hash_map<const HloInstruction*, int64>;
HloDomainMap(string domain_kind) : domain_kind_(std::move(domain_kind)) {}
@@ -120,8 +119,8 @@
string domain_kind_;
std::vector<std::unique_ptr<DomainMetadata::Domain>> instruction_domains_;
- tensorflow::gtl::FlatMap<HloInstruction*, int64> instruction_to_domain_;
- tensorflow::gtl::FlatMap<HloInstruction*, int64> domain_metadata_id_;
+ absl::flat_hash_map<HloInstruction*, int64> instruction_to_domain_;
+ absl::flat_hash_map<HloInstruction*, int64> domain_metadata_id_;
};
} // namespace xla
diff --git a/tensorflow/compiler/xla/service/hlo_instruction.cc b/tensorflow/compiler/xla/service/hlo_instruction.cc
index ad58833..5d5c9c7 100644
--- a/tensorflow/compiler/xla/service/hlo_instruction.cc
+++ b/tensorflow/compiler/xla/service/hlo_instruction.cc
@@ -22,6 +22,7 @@
#include <utility>
#include "absl/algorithm/container.h"
+#include "absl/container/flat_hash_map.h"
#include "absl/container/inlined_vector.h"
#include "absl/memory/memory.h"
#include "absl/strings/ascii.h"
@@ -43,7 +44,6 @@
#include "tensorflow/compiler/xla/types.h"
#include "tensorflow/compiler/xla/util.h"
#include "tensorflow/core/lib/core/errors.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/lib/gtl/flatset.h"
#include "tensorflow/core/lib/gtl/map_util.h"
#include "tensorflow/core/platform/human_readable_json.h"
@@ -59,8 +59,8 @@
/* static */
StatusOr<std::unique_ptr<HloInstruction>> HloInstruction::CreateFromProto(
const HloInstructionProto& proto,
- const tensorflow::gtl::FlatMap<int64, HloInstruction*>& instruction_map,
- const tensorflow::gtl::FlatMap<int64, HloComputation*>& computation_map) {
+ const absl::flat_hash_map<int64, HloInstruction*>& instruction_map,
+ const absl::flat_hash_map<int64, HloComputation*>& computation_map) {
TF_RET_CHECK(!proto.opcode().empty());
TF_ASSIGN_OR_RETURN(HloOpcode opcode, StringToHloOpcode(proto.opcode()));
TF_RET_CHECK(proto.has_shape());
@@ -266,7 +266,8 @@
<< "Expect 1 called computation for fusion instruction but sees "
<< proto.called_computation_ids_size();
const int64 fusion_id = proto.called_computation_ids(0);
- auto* fused_computation = FindPtrOrNull(computation_map, fusion_id);
+ auto* fused_computation =
+ tensorflow::gtl::FindPtrOrNull(computation_map, fusion_id);
TF_RET_CHECK(fused_computation != nullptr)
<< "No fusion computation with id " << fusion_id;
instruction = CreateFusion(proto.shape(), fusion_kind, all_operands(),
@@ -379,7 +380,8 @@
break;
case HloOpcode::kCustomCall:
instruction = CreateCustomCall(proto.shape(), all_operands(),
- proto.custom_call_target());
+ proto.custom_call_target(),
+ proto.custom_call_opaque());
if (proto.has_window()) {
static_cast<HloCustomCallInstruction*>(instruction.get())
->set_window(proto.window());
@@ -1108,9 +1110,9 @@
/* static */ std::unique_ptr<HloInstruction> HloInstruction::CreateCustomCall(
const Shape& shape, absl::Span<HloInstruction* const> operands,
- absl::string_view custom_call_target) {
- return absl::make_unique<HloCustomCallInstruction>(shape, operands,
- custom_call_target);
+ absl::string_view custom_call_target, absl::string_view opaque) {
+ return absl::make_unique<HloCustomCallInstruction>(
+ shape, operands, custom_call_target, opaque);
}
/* static */ std::unique_ptr<HloInstruction> HloInstruction::CreateTuple(
@@ -2423,7 +2425,7 @@
static Status PostOrderDFS(HloInstruction* root, Visitor* visitor,
const InternalCompareFunction* operand_order,
bool ignore_control_predecessors) {
- visitor->ReserveVisitStates(root->GetModule()->NumUniqueInstructionIds());
+ visitor->ReserveVisitStates(root->GetModule()->instruction_count());
// dfs_stack holds pairs of <HloInstruction*->unique_id(), HloInstruction*>.
//
@@ -2660,14 +2662,14 @@
// the value of this parameter, which would save stack space but not allow us
// to finish early if we find a reuse.
static UseKind Compute(int64 i, const HloInstruction& hlo) {
- tensorflow::gtl::FlatMap<const HloInstruction*, UseKind> memoization_cache;
+ absl::flat_hash_map<const HloInstruction*, UseKind> memoization_cache;
return ComputeInternal(i, hlo, &memoization_cache);
}
private:
static UseKind ComputeInternal(
int64 i, const HloInstruction& hlo,
- tensorflow::gtl::FlatMap<const HloInstruction*, UseKind>* cache) {
+ absl::flat_hash_map<const HloInstruction*, UseKind>* cache) {
if (auto hlo_param = DynCast<HloParameterInstruction>(&hlo)) {
if (hlo_param->parameter_number() == i) {
return UseKind::kUse;
diff --git a/tensorflow/compiler/xla/service/hlo_instruction.h b/tensorflow/compiler/xla/service/hlo_instruction.h
index d615df0..1bfdc88 100644
--- a/tensorflow/compiler/xla/service/hlo_instruction.h
+++ b/tensorflow/compiler/xla/service/hlo_instruction.h
@@ -32,6 +32,7 @@
#include <unordered_set>
#include <vector>
+#include "absl/container/flat_hash_map.h"
#include "absl/container/inlined_vector.h"
#include "absl/memory/memory.h"
#include "absl/strings/str_cat.h"
@@ -50,7 +51,6 @@
#include "tensorflow/compiler/xla/types.h"
#include "tensorflow/compiler/xla/xla_data.pb.h"
#include "tensorflow/core/lib/core/status.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/lib/gtl/iterator_range.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/macros.h"
@@ -247,7 +247,7 @@
private:
int64 index;
- tensorflow::gtl::FlatMap<string, string> canonical_name_map;
+ absl::flat_hash_map<string, string> canonical_name_map;
};
// HLO instructions are the atomic unit of the high-level compiler's IR.
@@ -350,8 +350,8 @@
// calls.
static StatusOr<std::unique_ptr<HloInstruction>> CreateFromProto(
const HloInstructionProto& proto,
- const tensorflow::gtl::FlatMap<int64, HloInstruction*>& instruction_map,
- const tensorflow::gtl::FlatMap<int64, HloComputation*>& computation_map);
+ const absl::flat_hash_map<int64, HloInstruction*>& instruction_map,
+ const absl::flat_hash_map<int64, HloComputation*>& computation_map);
// Creates a parameter-retrieving instruction.
static std::unique_ptr<HloInstruction> CreateParameter(int64 parameter_number,
@@ -718,10 +718,11 @@
HloComputation* computation);
// Creates a custom call instruction that applies the given custom call target
- // to the given operands. "shape" is the resultant shape.
+ // to the given operands. "opaque" can be an arbitrary string with a
+ // backend-specific interpretation. "shape" is the resultant shape.
static std::unique_ptr<HloInstruction> CreateCustomCall(
const Shape& shape, absl::Span<HloInstruction* const> operands,
- absl::string_view custom_call_target);
+ absl::string_view custom_call_target, absl::string_view opaque = "");
// Creates a tuple instruction with the given elements. This is a convenience
// wrapper around CreateVariadic.
diff --git a/tensorflow/compiler/xla/service/hlo_instructions.cc b/tensorflow/compiler/xla/service/hlo_instructions.cc
index e92882c..1bc168c 100644
--- a/tensorflow/compiler/xla/service/hlo_instructions.cc
+++ b/tensorflow/compiler/xla/service/hlo_instructions.cc
@@ -18,6 +18,7 @@
#include <deque>
#include "absl/algorithm/container.h"
+#include "absl/container/flat_hash_map.h"
#include "absl/memory/memory.h"
#include "absl/strings/escaping.h"
#include "absl/strings/str_cat.h"
@@ -28,7 +29,6 @@
#include "tensorflow/compiler/xla/service/hlo_computation.h"
#include "tensorflow/compiler/xla/service/hlo_module.h"
#include "tensorflow/compiler/xla/window_util.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
namespace xla {
namespace {
@@ -1042,7 +1042,8 @@
const int64 param_no = operand_count();
// Name the parameter after the instruction it represents in the outer
// (non-fusion) computation.
- string param_name = StrCat(new_operand->name(), ".param_", param_no);
+ // string param_name = StrCat(new_operand->name(), ".param_", param_no);
+ string param_name = StrCat("param_", param_no);
HloInstruction* fused_parameter =
fused_instructions_computation()->AddParameter(
HloInstruction::CreateParameter(param_no, new_operand->shape(),
@@ -1098,7 +1099,7 @@
// Note that we add the unfused instructions to this->parent_ computation.
// This is necessary because the unique_id needs for an instruction and
// it's only added when inserting to the computation.
- tensorflow::gtl::FlatMap<HloInstruction*, HloInstruction*> old_to_new;
+ absl::flat_hash_map<HloInstruction*, HloInstruction*> old_to_new;
std::vector<HloInstruction*> unfused_instructions;
auto computation_to_merge =
instruction_to_merge->fused_instructions_computation();
@@ -1391,7 +1392,7 @@
}
Status HloFusionInstruction::DeduplicateFusionOperands() {
- tensorflow::gtl::FlatMap<const HloInstruction*, int> operand_indices;
+ absl::flat_hash_map<const HloInstruction*, int> operand_indices;
std::vector<int> operands_to_remove;
for (int i = 0; i < operand_count(); ++i) {
auto emplace_result = operand_indices.emplace(operand(i), i);
@@ -1830,9 +1831,10 @@
HloCustomCallInstruction::HloCustomCallInstruction(
const Shape& shape, absl::Span<HloInstruction* const> operands,
- absl::string_view custom_call_target)
+ absl::string_view custom_call_target, absl::string_view opaque)
: HloInstruction(HloOpcode::kCustomCall, shape),
custom_call_target_(custom_call_target.begin(), custom_call_target.end()),
+ opaque_(opaque.begin(), opaque.end()),
feature_group_count_(1) {
for (auto operand : operands) {
AppendOperand(operand);
@@ -1849,6 +1851,7 @@
*convolution_dimension_numbers_;
}
proto.set_custom_call_target(custom_call_target_);
+ proto.set_custom_call_opaque(opaque_);
proto.set_feature_group_count(feature_group_count_);
return proto;
}
@@ -1872,6 +1875,11 @@
// an HloComputation.
extra.push_back(
StrCat("custom_call_target=\"", CEscape(custom_call_target_), "\""));
+ // If the opaque string becomes enormous we may want to reconsider printing
+ // this inline and consider other options.
+ if (!opaque_.empty()) {
+ extra.push_back(StrCat("opaque=\"", CEscape(opaque_), "\""));
+ }
return extra;
}
@@ -1897,7 +1905,8 @@
if (feature_group_count_ != casted_other.feature_group_count_) {
return false;
}
- return custom_call_target_ == casted_other.custom_call_target_;
+ return custom_call_target_ == casted_other.custom_call_target_ &&
+ opaque_ == casted_other.opaque_;
}
std::unique_ptr<HloInstruction>
@@ -1905,7 +1914,7 @@
const Shape& shape, absl::Span<HloInstruction* const> new_operands,
HloCloneContext* context) const {
auto cloned = absl::make_unique<HloCustomCallInstruction>(
- shape, new_operands, custom_call_target());
+ shape, new_operands, custom_call_target(), opaque());
if (window_ != nullptr) {
cloned->set_window(*window_);
}
diff --git a/tensorflow/compiler/xla/service/hlo_instructions.h b/tensorflow/compiler/xla/service/hlo_instructions.h
index 2d7bc83..9c22f5d 100644
--- a/tensorflow/compiler/xla/service/hlo_instructions.h
+++ b/tensorflow/compiler/xla/service/hlo_instructions.h
@@ -1070,7 +1070,8 @@
public:
explicit HloCustomCallInstruction(const Shape& shape,
absl::Span<HloInstruction* const> operands,
- absl::string_view custom_call_target);
+ absl::string_view custom_call_target,
+ absl::string_view opaque);
const Window& window() const override {
CHECK(window_ != nullptr);
return *window_;
@@ -1090,6 +1091,7 @@
convolution_dimension_numbers_ =
absl::make_unique<ConvolutionDimensionNumbers>(dnums);
}
+ const string& opaque() const { return opaque_; }
const string& custom_call_target() const { return custom_call_target_; }
void set_feature_group_count(int64 feature_group_count) {
feature_group_count_ = feature_group_count;
@@ -1109,8 +1111,10 @@
std::unique_ptr<HloInstruction> CloneWithNewOperandsImpl(
const Shape& shape, absl::Span<HloInstruction* const> new_operands,
HloCloneContext* context) const override;
- // Name of a global symbol to call, only present for kCustomCall.
+ // Name of a global symbol to call.
string custom_call_target_;
+ // Opaque string interpreted by the backend.
+ string opaque_;
// Describes the window in a windowed operation such as convolution.
std::unique_ptr<Window> window_;
// Describes the dimension numbers used for a convolution.
diff --git a/tensorflow/compiler/xla/service/hlo_memory_scheduler.cc b/tensorflow/compiler/xla/service/hlo_memory_scheduler.cc
index c7ec88d..1c2b286 100644
--- a/tensorflow/compiler/xla/service/hlo_memory_scheduler.cc
+++ b/tensorflow/compiler/xla/service/hlo_memory_scheduler.cc
@@ -20,6 +20,7 @@
#include <utility>
#include <vector>
+#include "absl/container/flat_hash_map.h"
#include "tensorflow/compiler/xla/service/heap_simulator.h"
#include "tensorflow/compiler/xla/service/hlo_computation.h"
#include "tensorflow/compiler/xla/service/tuple_points_to_analysis.h"
@@ -74,7 +75,7 @@
const HloComputation& computation,
const TuplePointsToAnalysis& points_to_analysis,
const LogicalBuffer::SizeFunction& size_function,
- const tensorflow::gtl::FlatMap<const HloComputation*, int64>&
+ const absl::flat_hash_map<const HloComputation*, int64>&
memory_by_computation) {
ListScheduler scheduler(computation, points_to_analysis, size_function,
memory_by_computation);
@@ -99,7 +100,7 @@
ListScheduler(const HloComputation& computation,
const TuplePointsToAnalysis& points_to_analysis,
const LogicalBuffer::SizeFunction& size_function,
- const tensorflow::gtl::FlatMap<const HloComputation*, int64>&
+ const absl::flat_hash_map<const HloComputation*, int64>&
memory_by_computation)
: computation_(computation),
points_to_analysis_(points_to_analysis),
@@ -234,8 +235,7 @@
// Populate the ready list with instructions which have no operands or
// control predecessors.
- tensorflow::gtl::FlatMap<const HloInstruction*, int64>
- unscheduled_pred_count;
+ absl::flat_hash_map<const HloInstruction*, int64> unscheduled_pred_count;
for (auto* instruction : computation_.instructions()) {
// TODO(b/34466113): Replace this and above with successors() or
// predecessors() when these methods are added to HloInstruction.
@@ -251,8 +251,8 @@
std::multimap<Priority, ReadyListEntry> ready_queue;
// Map of ready instructions to their iterators in ready_queue.
- tensorflow::gtl::FlatMap<const HloInstruction*,
- std::multimap<Priority, ReadyListEntry>::iterator>
+ absl::flat_hash_map<const HloInstruction*,
+ std::multimap<Priority, ReadyListEntry>::iterator>
ready_instructions;
auto add_to_ready_queue = [&](HloInstruction* inst) {
@@ -347,12 +347,11 @@
// Computations are analyzed in post-order. When scheduling an instruction
// that includes subcomputations, such as a while loop, we use this map to
// look up the memory needed by subcomputations.
- const tensorflow::gtl::FlatMap<const HloComputation*, int64>&
+ const absl::flat_hash_map<const HloComputation*, int64>&
memory_by_computation_;
// A map containing the LogicalBuffers that each instruction uses.
- tensorflow::gtl::FlatMap<const HloInstruction*,
- std::vector<const LogicalBuffer*>>
+ absl::flat_hash_map<const HloInstruction*, std::vector<const LogicalBuffer*>>
buffer_uses_;
// A map containing the count of unscheduled HLOs which using a particular
@@ -379,7 +378,7 @@
const TuplePointsToAnalysis& points_to_analysis,
const LogicalBuffer::SizeFunction& size_function,
const MemorySchedulerAlgorithm& algorithm,
- const tensorflow::gtl::FlatMap<const HloComputation*, int64>&
+ const absl::flat_hash_map<const HloComputation*, int64>&
memory_by_computation) {
VLOG(2) << "Computation: " << computation.name();
if (algorithm) {
@@ -396,13 +395,13 @@
const HloComputation& computation,
const TuplePointsToAnalysis& points_to_analysis,
const LogicalBuffer::SizeFunction& size_function,
- const tensorflow::gtl::FlatMap<const HloComputation*, int64>&
+ const absl::flat_hash_map<const HloComputation*, int64>&
memory_by_computation) {
// These variables are a hack to prevent overflows.
int64 cumulative_total_size = 0;
- int64 total_hlos = computation.parent()->NumUniqueInstructionIds();
- tensorflow::gtl::FlatMap<const HloInstruction*, int64> extra_users;
- tensorflow::gtl::FlatMap<const HloInstruction*, int64> total_sizes;
+ int64 total_hlos = computation.parent()->instruction_count();
+ absl::flat_hash_map<const HloInstruction*, int64> extra_users;
+ absl::flat_hash_map<const HloInstruction*, int64> total_sizes;
for (const HloInstruction* hlo : computation.MakeInstructionPostOrder()) {
if (ListScheduler::IgnoreInstruction(*hlo)) {
extra_users[hlo] = 0;
@@ -467,7 +466,7 @@
const HloComputation& computation,
const TuplePointsToAnalysis& points_to_analysis,
const LogicalBuffer::SizeFunction& size_function,
- const tensorflow::gtl::FlatMap<const HloComputation*, int64>&
+ const absl::flat_hash_map<const HloComputation*, int64>&
memory_by_computation) {
return ListScheduler::Run(computation, points_to_analysis, size_function,
memory_by_computation);
@@ -477,7 +476,7 @@
const HloComputation& computation,
const TuplePointsToAnalysis& points_to_analysis,
const LogicalBuffer::SizeFunction& size_function,
- const tensorflow::gtl::FlatMap<const HloComputation*, int64>&
+ const absl::flat_hash_map<const HloComputation*, int64>&
memory_by_computation) {
return HloInstructionSequence(computation.MakeInstructionPostOrder());
}
@@ -486,7 +485,7 @@
const HloComputation& computation,
const TuplePointsToAnalysis& points_to_analysis,
const LogicalBuffer::SizeFunction& size_function,
- const tensorflow::gtl::FlatMap<const HloComputation*, int64>&
+ const absl::flat_hash_map<const HloComputation*, int64>&
memory_by_computation) {
// We try a few schedulers and choose whichever returns a lower min-memory,
// not accounting for fragmentation.
@@ -549,7 +548,7 @@
HloSchedule schedule(&module);
TF_ASSIGN_OR_RETURN(std::unique_ptr<TuplePointsToAnalysis> points_to_analysis,
TuplePointsToAnalysis::Run(&module));
- tensorflow::gtl::FlatMap<const HloComputation*, int64> memory_by_computation;
+ absl::flat_hash_map<const HloComputation*, int64> memory_by_computation;
for (const auto* computation : module.MakeComputationPostOrder()) {
if (!computation->IsFusionComputation()) {
TF_ASSIGN_OR_RETURN(HloInstructionSequence computation_sequence,
@@ -577,7 +576,7 @@
CHECK(!computation.IsFusionComputation());
TF_ASSIGN_OR_RETURN(std::unique_ptr<TuplePointsToAnalysis> points_to_analysis,
TuplePointsToAnalysis::Run(computation.parent()));
- tensorflow::gtl::FlatMap<const HloComputation*, int64> empty_map;
+ absl::flat_hash_map<const HloComputation*, int64> empty_map;
return ScheduleComputationHelper(computation, *points_to_analysis,
size_function, nullptr, empty_map);
}
diff --git a/tensorflow/compiler/xla/service/hlo_memory_scheduler.h b/tensorflow/compiler/xla/service/hlo_memory_scheduler.h
index 9964c6f..a4c1d3d 100644
--- a/tensorflow/compiler/xla/service/hlo_memory_scheduler.h
+++ b/tensorflow/compiler/xla/service/hlo_memory_scheduler.h
@@ -18,6 +18,7 @@
#include <vector>
+#include "absl/container/flat_hash_map.h"
#include "tensorflow/compiler/xla/service/hlo_instruction.h"
#include "tensorflow/compiler/xla/service/hlo_module.h"
#include "tensorflow/compiler/xla/service/hlo_ordering.h"
@@ -37,7 +38,7 @@
typedef std::function<StatusOr<HloInstructionSequence>(
const HloComputation&, const TuplePointsToAnalysis&,
const LogicalBuffer::SizeFunction&,
- const tensorflow::gtl::FlatMap<const HloComputation*, int64>&)>
+ const absl::flat_hash_map<const HloComputation*, int64>&)>
MemorySchedulerAlgorithm;
// List scheduler
@@ -45,7 +46,7 @@
const HloComputation& computation,
const TuplePointsToAnalysis& points_to_analysis,
const LogicalBuffer::SizeFunction& size_function,
- const tensorflow::gtl::FlatMap<const HloComputation*, int64>&
+ const absl::flat_hash_map<const HloComputation*, int64>&
memory_by_computation);
// DFS-order scheduler
@@ -53,7 +54,7 @@
const HloComputation& computation,
const TuplePointsToAnalysis& points_to_analysis,
const LogicalBuffer::SizeFunction& size_function,
- const tensorflow::gtl::FlatMap<const HloComputation*, int64>&
+ const absl::flat_hash_map<const HloComputation*, int64>&
memory_by_computation);
// Naive Post Order scheduler
@@ -61,7 +62,7 @@
const HloComputation& computation,
const TuplePointsToAnalysis& points_to_analysis,
const LogicalBuffer::SizeFunction& size_function,
- const tensorflow::gtl::FlatMap<const HloComputation*, int64>&
+ const absl::flat_hash_map<const HloComputation*, int64>&
memory_by_computation);
// The default scheduling algorithm. Runs both the list scheduler
@@ -71,7 +72,7 @@
const HloComputation& computation,
const TuplePointsToAnalysis& points_to_analysis,
const LogicalBuffer::SizeFunction& size_function,
- const tensorflow::gtl::FlatMap<const HloComputation*, int64>&
+ const absl::flat_hash_map<const HloComputation*, int64>&
memory_by_computation);
// Returns an HloSchedule which seeks to minimize the memory required for
diff --git a/tensorflow/compiler/xla/service/hlo_memory_scheduler_test.cc b/tensorflow/compiler/xla/service/hlo_memory_scheduler_test.cc
index 1b9e9bf..5a9fccc 100644
--- a/tensorflow/compiler/xla/service/hlo_memory_scheduler_test.cc
+++ b/tensorflow/compiler/xla/service/hlo_memory_scheduler_test.cc
@@ -19,6 +19,7 @@
#include <string>
#include "absl/algorithm/container.h"
+#include "absl/container/flat_hash_map.h"
#include "tensorflow/compiler/xla/service/heap_simulator.h"
#include "tensorflow/compiler/xla/service/hlo_computation.h"
#include "tensorflow/compiler/xla/service/hlo_dce.h"
@@ -247,7 +248,7 @@
EXPECT_TRUE(ordering.ExecutesBefore(bcast, add));
EXPECT_TRUE(ordering.ExecutesBefore(transpose, add));
- tensorflow::gtl::FlatMap<const HloComputation*, int64> memory_by_computation;
+ absl::flat_hash_map<const HloComputation*, int64> memory_by_computation;
memory_by_computation[cond_computation] = 17;
memory_by_computation[body_computation] = 16;
std::unique_ptr<TuplePointsToAnalysis> points_to_analysis =
@@ -409,7 +410,7 @@
EXPECT_EQ(module->entry_computation()->instruction_count(),
schedule.sequence(module->entry_computation()).size());
- tensorflow::gtl::FlatMap<const HloComputation*, int64> memory_by_computation;
+ absl::flat_hash_map<const HloComputation*, int64> memory_by_computation;
memory_by_computation[cond_computation] = 17;
memory_by_computation[body_computation] = 16;
std::unique_ptr<TuplePointsToAnalysis> points_to_analysis =
diff --git a/tensorflow/compiler/xla/service/hlo_module.cc b/tensorflow/compiler/xla/service/hlo_module.cc
index b3949f3..9359e9a 100644
--- a/tensorflow/compiler/xla/service/hlo_module.cc
+++ b/tensorflow/compiler/xla/service/hlo_module.cc
@@ -23,6 +23,7 @@
#include <utility>
#include "absl/algorithm/container.h"
+#include "absl/container/flat_hash_map.h"
#include "absl/memory/memory.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/compiler/xla/map_util.h"
@@ -285,8 +286,8 @@
<< ShapeUtil::HumanStringWithLayout(expected_program_shape.result())
<< ", actual: " << ShapeUtil::HumanStringWithLayout(result_shape);
- tensorflow::gtl::FlatMap<int64, HloComputation*> computation_map;
- tensorflow::gtl::FlatMap<HloComputation*, int64> to_proto_id;
+ absl::flat_hash_map<int64, HloComputation*> computation_map;
+ absl::flat_hash_map<HloComputation*, int64> to_proto_id;
std::vector<std::unique_ptr<HloComputation>> computations;
HloComputation* entry = nullptr;
for (const HloComputationProto& computation_proto : proto.computations()) {
diff --git a/tensorflow/compiler/xla/service/hlo_module_group_metadata.h b/tensorflow/compiler/xla/service/hlo_module_group_metadata.h
index 278d94c..0311b73 100644
--- a/tensorflow/compiler/xla/service/hlo_module_group_metadata.h
+++ b/tensorflow/compiler/xla/service/hlo_module_group_metadata.h
@@ -22,6 +22,7 @@
#include <unordered_set>
#include <vector>
+#include "absl/container/flat_hash_map.h"
#include "absl/types/optional.h"
#include "tensorflow/compiler/xla/service/hlo_computation.h"
#include "tensorflow/compiler/xla/service/hlo_instruction.h"
@@ -30,7 +31,6 @@
#include "tensorflow/compiler/xla/status.h"
#include "tensorflow/compiler/xla/statusor.h"
#include "tensorflow/core/lib/core/status.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/platform/types.h"
namespace xla {
@@ -250,25 +250,25 @@
std::vector<std::unique_ptr<std::vector<HloInstruction*>>> companion_sets_;
// Map from each companion while instruction to the index into companion_set_.
- tensorflow::gtl::FlatMap<const HloInstruction*, int64> companion_set_index_;
+ absl::flat_hash_map<const HloInstruction*, int64> companion_set_index_;
// Map from computation to the instruction using it (a kWhile, kConditional).
- tensorflow::gtl::FlatMap<const HloComputation*, TrackedInstruction>
+ absl::flat_hash_map<const HloComputation*, TrackedInstruction>
tracked_instructions_;
// Maps tracked instructions (kWhile, kConditional, kCall, ...) to the set of
// communicating instructions within the proper called computation(s).
- tensorflow::gtl::FlatMap<HloInstruction*, std::vector<HloInstruction*>>
+ absl::flat_hash_map<HloInstruction*, std::vector<HloInstruction*>>
tracked_instructions_comms_;
// All channels in the module.
std::vector<Channel> channels_;
// Map from channel ids to the index in channels_.
- tensorflow::gtl::FlatMap<int64, int64> channel_id_map_;
+ absl::flat_hash_map<int64, int64> channel_id_map_;
// Map from all-reduce ids to the all reduce instructions.
- tensorflow::gtl::FlatMap<int64, std::vector<HloInstruction*>> all_reduce_map_;
+ absl::flat_hash_map<int64, std::vector<HloInstruction*>> all_reduce_map_;
// The maximum channel id used in the module group.
int64 max_channel_id_ = -1;
@@ -276,7 +276,7 @@
// The modules that this metadata was built from.
const std::vector<HloModule*>& modules_;
- tensorflow::gtl::FlatMap<HloModule*, std::unique_ptr<TuplePointsToAnalysis>>
+ absl::flat_hash_map<HloModule*, std::unique_ptr<TuplePointsToAnalysis>>
points_to_analyses_;
};
diff --git a/tensorflow/compiler/xla/service/hlo_module_group_util.h b/tensorflow/compiler/xla/service/hlo_module_group_util.h
index 309c230..f21b44b 100644
--- a/tensorflow/compiler/xla/service/hlo_module_group_util.h
+++ b/tensorflow/compiler/xla/service/hlo_module_group_util.h
@@ -20,6 +20,7 @@
#include <memory>
#include <vector>
+#include "absl/container/flat_hash_map.h"
#include "absl/types/span.h"
#include "tensorflow/compiler/xla/service/hlo_computation.h"
#include "tensorflow/compiler/xla/service/hlo_instruction.h"
@@ -28,7 +29,6 @@
#include "tensorflow/compiler/xla/status.h"
#include "tensorflow/compiler/xla/statusor.h"
#include "tensorflow/core/lib/core/status.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
namespace xla {
@@ -87,7 +87,7 @@
// * visit_state: map from each instruction to its visit state.
// * visit_function: function called when each instruction group.
// * root: the root instruction of the traversal.
- using VisitStates = tensorflow::gtl::FlatMap<HloInstruction*, VisitState>;
+ using VisitStates = absl::flat_hash_map<HloInstruction*, VisitState>;
Status VisitTopologicalOrder(VisitStates* visit_state,
const VisitFunction& visit_function,
HloInstruction* root);
diff --git a/tensorflow/compiler/xla/service/hlo_opcode.cc b/tensorflow/compiler/xla/service/hlo_opcode.cc
index 2d4e385..4551a1c 100644
--- a/tensorflow/compiler/xla/service/hlo_opcode.cc
+++ b/tensorflow/compiler/xla/service/hlo_opcode.cc
@@ -14,9 +14,9 @@
==============================================================================*/
#include "tensorflow/compiler/xla/service/hlo_opcode.h"
+#include "absl/container/flat_hash_map.h"
#include "tensorflow/compiler/xla/types.h"
#include "tensorflow/compiler/xla/util.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
namespace xla {
@@ -31,7 +31,7 @@
}
StatusOr<HloOpcode> StringToHloOpcode(const string& opcode_name) {
- static auto* opcode_map = new tensorflow::gtl::FlatMap<string, HloOpcode>({
+ static auto* opcode_map = new absl::flat_hash_map<string, HloOpcode>({
#define STRING_TO_OPCODE_ENTRY(enum_name, opcode_name, ...) \
{opcode_name, HloOpcode::enum_name},
HLO_OPCODE_LIST(STRING_TO_OPCODE_ENTRY)
diff --git a/tensorflow/compiler/xla/service/hlo_ordering.cc b/tensorflow/compiler/xla/service/hlo_ordering.cc
index f1dc08b..23d41d9 100644
--- a/tensorflow/compiler/xla/service/hlo_ordering.cc
+++ b/tensorflow/compiler/xla/service/hlo_ordering.cc
@@ -92,14 +92,18 @@
}
bool HloOrdering::IsDefinedBefore(const HloValue& a, const HloValue& b) const {
- // If 'b' is an entry param then 'a' cannot be defined before 'b' because 'b'
- // is live into the module.
+ // Entry parameter should always be defined before other instructions.
const HloModule* module = b.defining_instruction()->parent()->parent();
if (b.defining_instruction()->parent() == module->entry_computation() &&
b.defining_instruction()->opcode() == HloOpcode::kParameter) {
return false;
}
+ if (a.defining_instruction()->parent() == module->entry_computation() &&
+ a.defining_instruction()->opcode() == HloOpcode::kParameter) {
+ return true;
+ }
+
// Phi values require special handling. Because XLA does not have a phi
// instruction, the definition instruction of the phis values are
// placeholders: either the subcomputation parameter (body or condition) or
@@ -316,7 +320,7 @@
for (auto predecessor : all) {
if (predecessors_.at(computation)
->IsReachable(predecessor, instruction)) {
- pieces.push_back(absl::StrFormat(" %s", predecessor->name()));
+ pieces.push_back(absl::StrFormat(" %s", predecessor->name()));
}
}
}
diff --git a/tensorflow/compiler/xla/service/hlo_ordering.h b/tensorflow/compiler/xla/service/hlo_ordering.h
index b0361c3..6631349 100644
--- a/tensorflow/compiler/xla/service/hlo_ordering.h
+++ b/tensorflow/compiler/xla/service/hlo_ordering.h
@@ -20,6 +20,7 @@
#include <string>
#include <utility>
+#include "absl/container/flat_hash_map.h"
#include "tensorflow/compiler/xla/service/call_graph.h"
#include "tensorflow/compiler/xla/service/hlo.pb.h"
#include "tensorflow/compiler/xla/service/hlo_dataflow_analysis.h"
@@ -28,7 +29,6 @@
#include "tensorflow/compiler/xla/service/hlo_schedule.h"
#include "tensorflow/compiler/xla/service/hlo_value.h"
#include "tensorflow/compiler/xla/types.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
namespace xla {
@@ -120,8 +120,8 @@
// predecessors. An instruction is an element of its own predecessor set.
//
// Subclasses should fill this in to define the desired ordering.
- tensorflow::gtl::FlatMap<const HloComputation*,
- std::unique_ptr<HloReachabilityMap>>
+ absl::flat_hash_map<const HloComputation*,
+ std::unique_ptr<HloReachabilityMap>>
predecessors_;
};
@@ -204,7 +204,7 @@
// this map so more than one instruction may have the same position
// value. This is not a problem because ExecutesBefore also verifies
// instructions are in the same computation.
- tensorflow::gtl::FlatMap<const HloInstruction*, int> order_position_;
+ absl::flat_hash_map<const HloInstruction*, int> order_position_;
};
} // namespace xla
diff --git a/tensorflow/compiler/xla/service/hlo_ordering_test.cc b/tensorflow/compiler/xla/service/hlo_ordering_test.cc
index 00970bc..b045adc 100644
--- a/tensorflow/compiler/xla/service/hlo_ordering_test.cc
+++ b/tensorflow/compiler/xla/service/hlo_ordering_test.cc
@@ -174,6 +174,26 @@
EXPECT_FALSE(ordering.ExecutesBefore(body_param, cond_param));
}
+TEST_F(HloOrderingTest, ParametersDefinedBeforeOthers) {
+ // Entry parameter should always be defined before other instruction.
+ auto module = CreateNewModule();
+ const Shape scalar_shape = ShapeUtil::MakeShape(xla::F32, {});
+ auto builder = HloComputation::Builder(TestName());
+ auto constant = builder.AddInstruction(
+ HloInstruction::CreateConstant(LiteralUtil::CreateR0<float>(1.0)));
+ auto param = builder.AddInstruction(
+ HloInstruction::CreateParameter(0, scalar_shape, "param"));
+ module->AddEntryComputation(builder.Build());
+ TF_ASSERT_OK_AND_ASSIGN(auto dataflow,
+ HloDataflowAnalysis::Run(*module, /*ssa_form=*/true));
+
+ DependencyHloOrdering ordering(module.get());
+ EXPECT_TRUE(ordering.IsDefinedBefore(dataflow->GetValueDefinedAt(param),
+ dataflow->GetValueDefinedAt(constant)));
+ EXPECT_TRUE(!ordering.IsDefinedBefore(dataflow->GetValueDefinedAt(constant),
+ dataflow->GetValueDefinedAt(param)));
+}
+
TEST_F(HloOrderingTest, ValuesInWhileComputations) {
// Tests the ordering of values (defined by dataflow analysis) in the body and
// condition of a while instruction. HLO code:
diff --git a/tensorflow/compiler/xla/service/hlo_parser.cc b/tensorflow/compiler/xla/service/hlo_parser.cc
index 37197b2..25b7074 100644
--- a/tensorflow/compiler/xla/service/hlo_parser.cc
+++ b/tensorflow/compiler/xla/service/hlo_parser.cc
@@ -1266,11 +1266,13 @@
}
case HloOpcode::kCustomCall: {
optional<string> custom_call_target;
+ optional<string> opaque;
optional<Window> window;
optional<ConvolutionDimensionNumbers> dnums;
optional<int64> feature_group_count;
attrs["custom_call_target"] = {/*required=*/true, AttrTy::kString,
&custom_call_target};
+ attrs["opaque"] = {/*required=*/false, AttrTy::kString, &opaque};
attrs["window"] = {/*required=*/false, AttrTy::kWindow, &window};
attrs["dim_labels"] = {/*required=*/false,
AttrTy::kConvolutionDimensionNumbers, &dnums};
@@ -1279,8 +1281,9 @@
if (!ParseOperands(&operands) || !ParseAttributes(attrs)) {
return false;
}
- instruction = builder->AddInstruction(HloInstruction::CreateCustomCall(
- shape, operands, *custom_call_target));
+ instruction = builder->AddInstruction(
+ HloInstruction::CreateCustomCall(shape, operands, *custom_call_target,
+ opaque.has_value() ? *opaque : ""));
if (window.has_value()) {
instruction->set_window(*window);
}
diff --git a/tensorflow/compiler/xla/service/hlo_parser_test.cc b/tensorflow/compiler/xla/service/hlo_parser_test.cc
index cca50fa..96db96b 100644
--- a/tensorflow/compiler/xla/service/hlo_parser_test.cc
+++ b/tensorflow/compiler/xla/service/hlo_parser_test.cc
@@ -1004,6 +1004,18 @@
)"
},
+// CustomCall with opaque value.
+{
+"CustomCallWithOpaque",
+R"(HloModule custom_call
+
+ENTRY CustomCall {
+ constant = f32[1]{0} constant({12345})
+ ROOT custom-call = f32[1,2,3]{0,2,1} custom-call(constant), custom_call_target="foo\"bar", opaque="this string is opaque"
+}
+
+)"
+},
// Variables with non-default names
{
"NonDefaultNames",
diff --git a/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc b/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc
index 8c2f928..59fd01c 100644
--- a/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc
+++ b/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc
@@ -17,6 +17,7 @@
#include <functional>
+#include "absl/container/flat_hash_map.h"
#include "absl/strings/str_format.h"
#include "absl/strings/str_join.h"
#include "tensorflow/compiler/xla/service/hlo_graph_dumper.h"
@@ -98,7 +99,7 @@
if (!proto_dump_path.empty()) {
static tensorflow::mutex mu(tensorflow::LINKER_INITIALIZED);
static auto* const module_id_to_pass_number =
- new tensorflow::gtl::FlatMap<int64, int64>();
+ new absl::flat_hash_map<int64, int64>();
tensorflow::mutex_lock lock(mu);
const int64 pass_number = (*module_id_to_pass_number)[module.unique_id()]++;
diff --git a/tensorflow/compiler/xla/service/hlo_reachability.h b/tensorflow/compiler/xla/service/hlo_reachability.h
index b66a2aa..5a5f01f 100644
--- a/tensorflow/compiler/xla/service/hlo_reachability.h
+++ b/tensorflow/compiler/xla/service/hlo_reachability.h
@@ -19,11 +19,11 @@
#include <list>
#include <vector>
+#include "absl/container/flat_hash_map.h"
#include "absl/types/span.h"
#include "tensorflow/compiler/xla/map_util.h"
#include "tensorflow/compiler/xla/types.h"
#include "tensorflow/core/lib/core/status.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/platform/types.h"
namespace xla {
@@ -154,7 +154,7 @@
// Dense assignment from HloInstruction* to number. These numbers index
// into the bit_vectors_ vector and into the bits within a BitVector.
- tensorflow::gtl::FlatMap<const HloInstruction*, int> indices_;
+ absl::flat_hash_map<const HloInstruction*, int> indices_;
// Bitvectors holding the reachability to each instruction. The bit vector for
// instruction X includes ones for each instruction which X is reachable from.
diff --git a/tensorflow/compiler/xla/service/hlo_rematerialization.cc b/tensorflow/compiler/xla/service/hlo_rematerialization.cc
index a438671..abdd9a9 100644
--- a/tensorflow/compiler/xla/service/hlo_rematerialization.cc
+++ b/tensorflow/compiler/xla/service/hlo_rematerialization.cc
@@ -20,6 +20,7 @@
#include <set>
#include <string>
+#include "absl/container/flat_hash_map.h"
#include "absl/container/inlined_vector.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_format.h"
@@ -75,7 +76,7 @@
// cache before, and eventually calling the IsRematerializable() API.
bool CanBeRematerialized(
const HloInstruction* instruction,
- tensorflow::gtl::FlatMap<const HloInstruction*, bool>* remat_able) {
+ absl::flat_hash_map<const HloInstruction*, bool>* remat_able) {
auto it = remat_able->find(instruction);
if (it != remat_able->end()) {
return it->second;
@@ -268,7 +269,7 @@
Item* first_;
// Item for each instruction.
- tensorflow::gtl::FlatMap<const HloInstruction*, Item*> item_map_;
+ absl::flat_hash_map<const HloInstruction*, Item*> item_map_;
};
// Return the items which use the given LogicalBuffer. Sets
@@ -503,7 +504,7 @@
PointsToSet::BufferSet live_out_set =
points_to_analysis.GetPointsToSet(computation_->root_instruction())
.CreateFlattenedSet();
- tensorflow::gtl::FlatMap<const LogicalBuffer*, BufferId>
+ absl::flat_hash_map<const LogicalBuffer*, BufferId>
logical_buffer_to_buffer_id;
for (auto* item = instruction_list_.first(); item != nullptr;
@@ -854,7 +855,7 @@
Item* PickRematerializationCandidate(
const MemoryUsageTracker& memory_tracker,
const InstructionList& instruction_list, int64 memory_limit_bytes,
- tensorflow::gtl::FlatMap<const HloInstruction*, bool>* remat_able) {
+ absl::flat_hash_map<const HloInstruction*, bool>* remat_able) {
Item* best_item = nullptr;
int64 best_cost = 0;
@@ -983,7 +984,7 @@
tensorflow::gtl::FlatSet<const HloInstruction*> remat_move_instructions;
// The map from instructions to their rematerializable status.
- tensorflow::gtl::FlatMap<const HloInstruction*, bool> remat_able;
+ absl::flat_hash_map<const HloInstruction*, bool> remat_able;
// The peak memory of the computation at any point in the instruction
// sequence.
diff --git a/tensorflow/compiler/xla/service/hlo_rematerialization.h b/tensorflow/compiler/xla/service/hlo_rematerialization.h
index 7330d73..5a02e3a 100644
--- a/tensorflow/compiler/xla/service/hlo_rematerialization.h
+++ b/tensorflow/compiler/xla/service/hlo_rematerialization.h
@@ -15,6 +15,7 @@
#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_HLO_REMATERIALIZATION_H_
#define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_REMATERIALIZATION_H_
+#include "absl/container/flat_hash_map.h"
#include "tensorflow/compiler/xla/service/buffer_liveness.h"
#include "tensorflow/compiler/xla/service/call_graph.h"
#include "tensorflow/compiler/xla/service/hlo_computation.h"
@@ -115,8 +116,7 @@
// computations called from sequential context
// (CallContext::kSequential). These values are updated as rematerialization
// occurs.
- tensorflow::gtl::FlatMap<const HloComputation*, int64>
- computation_peak_memory_;
+ absl::flat_hash_map<const HloComputation*, int64> computation_peak_memory_;
std::unique_ptr<TuplePointsToAnalysis> points_to_analysis_;
diff --git a/tensorflow/compiler/xla/service/hlo_schedule.cc b/tensorflow/compiler/xla/service/hlo_schedule.cc
index 3fc5dbe..7c5c98f 100644
--- a/tensorflow/compiler/xla/service/hlo_schedule.cc
+++ b/tensorflow/compiler/xla/service/hlo_schedule.cc
@@ -18,6 +18,7 @@
#include <queue>
#include <vector>
+#include "absl/container/flat_hash_map.h"
#include "absl/strings/str_format.h"
#include "absl/strings/str_join.h"
#include "tensorflow/compiler/xla/map_util.h"
@@ -30,7 +31,7 @@
/* static */ StatusOr<HloSchedule> HloSchedule::CreateFromProto(
const HloModule* module, const HloScheduleProto& proto) {
- tensorflow::gtl::FlatMap<int64, const HloComputation*> id_to_computation;
+ absl::flat_hash_map<int64, const HloComputation*> id_to_computation;
for (const HloComputation* computation : module->computations()) {
id_to_computation[computation->unique_id()] = computation;
}
@@ -44,7 +45,7 @@
<< "No computation exists in HLO module with id " << computation_id;
const HloComputation* computation = comp_it->second;
- tensorflow::gtl::FlatMap<int64, const HloInstruction*> id_to_instruction;
+ absl::flat_hash_map<int64, const HloInstruction*> id_to_instruction;
for (const HloInstruction* instruction : computation->instructions()) {
id_to_instruction[instruction->unique_id()] = instruction;
}
@@ -112,7 +113,7 @@
const HloComputation* computation) {
// Map from unique ID to HloInstruction pointer for instructions in the
// computation.
- tensorflow::gtl::FlatMap<int, const HloInstruction*> id_to_instruction;
+ absl::flat_hash_map<int, const HloInstruction*> id_to_instruction;
for (const HloInstruction* instruction : computation->instructions()) {
InsertOrDie(&id_to_instruction, instruction->unique_id(), instruction);
}
@@ -126,15 +127,13 @@
// Map from HloInstruction X to newly added instructions (instruction is in
// computation, but not in schedule) which use X. If an instruction is not in
// the map, then it has no users which are newly added instructions.
- tensorflow::gtl::FlatMap<const HloInstruction*,
- std::vector<const HloInstruction*>>
+ absl::flat_hash_map<const HloInstruction*, std::vector<const HloInstruction*>>
new_instruction_uses;
// For each newly added instruction, this is the count of the instruction's
// operands that have not yet been scheduled. When this value reaches zero,
// then the instruction may be placed in the schedule.
- tensorflow::gtl::FlatMap<const HloInstruction*, int>
- unscheduled_operand_count;
+ absl::flat_hash_map<const HloInstruction*, int> unscheduled_operand_count;
// Create a worklist of newly added instructions which are ready to be added
// to the schedule. Initialize worklist with those that have zero operands.
@@ -217,9 +216,9 @@
}
for (auto it = sequences_.begin(); it != sequences_.end();) {
if (nonfusion_computations_ids.count(it->first) == 0) {
- it = sequences_.erase(it);
+ sequences_.erase(it++);
} else {
- it++;
+ ++it;
}
}
}
@@ -254,7 +253,7 @@
// For each computation verify the set of instructions is the same and that
// each dependency and control edge is honored.
for (const HloComputation* computation : nonfusion_computations) {
- tensorflow::gtl::FlatMap<const HloInstruction*, int> instruction_position;
+ absl::flat_hash_map<const HloInstruction*, int> instruction_position;
int pos = 0;
for (const HloInstruction* instruction :
sequence(computation).instructions()) {
diff --git a/tensorflow/compiler/xla/service/hlo_schedule.h b/tensorflow/compiler/xla/service/hlo_schedule.h
index 270fe60..0a71410 100644
--- a/tensorflow/compiler/xla/service/hlo_schedule.h
+++ b/tensorflow/compiler/xla/service/hlo_schedule.h
@@ -18,6 +18,7 @@
#include <vector>
+#include "absl/container/flat_hash_map.h"
#include "absl/types/span.h"
#include "tensorflow/compiler/xla/service/hlo_computation.h"
#include "tensorflow/compiler/xla/service/hlo_instruction.h"
@@ -103,8 +104,7 @@
// Returns a map from HloComputation unique ID to instruction sequence. The
// map contains all sequences in the schedule.
- const tensorflow::gtl::FlatMap<int64, HloInstructionSequence>& sequences()
- const {
+ const absl::flat_hash_map<int64, HloInstructionSequence>& sequences() const {
return sequences_;
}
@@ -148,7 +148,7 @@
// A map from computation unique ID to instruction sequence. Unique IDs are
// used rather than HloComputation pointers because HLO pointers are not
// unique across HLO transformations because pointers may be recycled.
- tensorflow::gtl::FlatMap<int64, HloInstructionSequence> sequences_;
+ absl::flat_hash_map<int64, HloInstructionSequence> sequences_;
};
std::ostream& operator<<(std::ostream& out, const HloSchedule& schedule);
diff --git a/tensorflow/compiler/xla/service/hlo_verifier.cc b/tensorflow/compiler/xla/service/hlo_verifier.cc
index 6eb6658..a772782 100644
--- a/tensorflow/compiler/xla/service/hlo_verifier.cc
+++ b/tensorflow/compiler/xla/service/hlo_verifier.cc
@@ -15,6 +15,7 @@
#include <set>
+#include "absl/container/flat_hash_map.h"
#include "absl/strings/str_join.h"
#include "tensorflow/compiler/xla/service/hlo_casting_utils.h"
#include "tensorflow/compiler/xla/service/hlo_instructions.h"
@@ -23,7 +24,6 @@
#include "tensorflow/compiler/xla/status_macros.h"
#include "tensorflow/compiler/xla/util.h"
#include "tensorflow/core/lib/core/errors.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
namespace xla {
@@ -993,7 +993,7 @@
// Checks various invariants of send and recv instructions.
Status VerifySendsAndRecvs(const HloModule& module) {
- tensorflow::gtl::FlatMap<int64, const HloInstruction*> host_channels;
+ absl::flat_hash_map<int64, const HloInstruction*> host_channels;
// Host send/recv instructions must have their own unique channel.
auto check_unique_host_channel = [&](const HloInstruction* instruction) {
const HloSendRecvInstruction* sendrecv =
@@ -1061,7 +1061,7 @@
TF_RETURN_IF_ERROR(VerifyHloStructure(module));
TF_RETURN_IF_ERROR(VerifySendsAndRecvs(*module));
- tensorflow::gtl::FlatMap<string, const HloInstruction*> instructions;
+ absl::flat_hash_map<string, const HloInstruction*> instructions;
for (auto* computation : module->computations()) {
for (const auto& instruction : computation->instructions()) {
diff --git a/tensorflow/compiler/xla/service/indexed_array_analysis.cc b/tensorflow/compiler/xla/service/indexed_array_analysis.cc
index 06f0e1e..7ee7892 100644
--- a/tensorflow/compiler/xla/service/indexed_array_analysis.cc
+++ b/tensorflow/compiler/xla/service/indexed_array_analysis.cc
@@ -16,6 +16,7 @@
#include "tensorflow/compiler/xla/service/indexed_array_analysis.h"
#include "absl/algorithm/container.h"
+#include "absl/container/flat_hash_map.h"
#include "absl/container/inlined_vector.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_join.h"
@@ -95,7 +96,7 @@
absl::InlinedVector<const HloInstruction*, 4> stack;
enum DfsState { kDiscovered, kVisited };
- gtl::FlatMap<const HloInstruction*, DfsState> dfs_state_map;
+ absl::flat_hash_map<const HloInstruction*, DfsState> dfs_state_map;
stack.push_back(root);
InsertOrDie(&dfs_state_map, root, kDiscovered);
diff --git a/tensorflow/compiler/xla/service/indexed_array_analysis.h b/tensorflow/compiler/xla/service/indexed_array_analysis.h
index 3e238f9..e5aa67f 100644
--- a/tensorflow/compiler/xla/service/indexed_array_analysis.h
+++ b/tensorflow/compiler/xla/service/indexed_array_analysis.h
@@ -18,10 +18,10 @@
#include <type_traits>
+#include "absl/container/flat_hash_map.h"
#include "tensorflow/compiler/xla/service/hlo_instruction.h"
#include "tensorflow/compiler/xla/service/hlo_module.h"
#include "tensorflow/compiler/xla/service/hlo_pass_interface.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/util/ptr_util.h"
namespace xla {
@@ -360,7 +360,7 @@
std::vector<std::unique_ptr<Array>> owned_tensors_;
std::vector<Literal> owned_literals_;
- tensorflow::gtl::FlatMap<const HloInstruction*, Array*> cache_;
+ absl::flat_hash_map<const HloInstruction*, Array*> cache_;
};
// A pass that prints all non-trivial results returned by IndexedArrayAnalysis.
diff --git a/tensorflow/compiler/xla/service/instruction_fusion.cc b/tensorflow/compiler/xla/service/instruction_fusion.cc
index 3fdc2ce..5a99c40 100644
--- a/tensorflow/compiler/xla/service/instruction_fusion.cc
+++ b/tensorflow/compiler/xla/service/instruction_fusion.cc
@@ -22,11 +22,11 @@
#include <vector>
#include "absl/algorithm/container.h"
+#include "absl/container/flat_hash_map.h"
#include "absl/memory/memory.h"
#include "tensorflow/compiler/xla/map_util.h"
#include "tensorflow/compiler/xla/service/hlo_opcode.h"
#include "tensorflow/core/lib/core/errors.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/platform/logging.h"
namespace xla {
@@ -188,13 +188,20 @@
bool InstructionFusion::CanFuseOnAllPaths(
HloInstruction* producer, HloInstruction* consumer,
- const HloInstructionSet& do_not_duplicate) {
+ const HloInstructionSet& do_not_fuse,
+ absl::flat_hash_map<std::pair<HloInstruction*, HloInstruction*>, bool>*
+ result_cache) {
if (consumer == producer) {
return true;
}
if (!consumer->IsFusible()) {
return false;
}
+ auto cache_it = result_cache->find(std::make_pair(producer, consumer));
+ if (cache_it != result_cache->end()) {
+ return cache_it->second;
+ }
+ bool result = true;
for (int64 i = 0, e = consumer->operand_count(); i < e; ++i) {
auto* consumer_operand = consumer->mutable_operand(i);
// If the operand is not on a path to the producer, it doesn't matter
@@ -202,20 +209,23 @@
if (!reachability_->IsReachable(producer, consumer_operand)) {
continue;
}
- if (do_not_duplicate.count(consumer_operand) > 0 ||
- !ShouldFuse(consumer, i)) {
- return false;
+ if (do_not_fuse.count(consumer_operand) > 0 || !ShouldFuse(consumer, i)) {
+ result = false;
+ break;
}
// The producer is reachable from consumer_operand which means we need
// to be able to fuse consumer_operand into consumer in order for
// producer to be fusible into consumer on all paths.
// Perform the recursive step: make sure producer can be fused into
// consumer_operand on all paths.
- if (!CanFuseOnAllPaths(producer, consumer_operand, do_not_duplicate)) {
- return false;
+ if (!CanFuseOnAllPaths(producer, consumer_operand, do_not_fuse,
+ result_cache)) {
+ result = false;
+ break;
}
}
- return true;
+ result_cache->emplace(std::make_pair(producer, consumer), result);
+ return result;
}
InstructionFusion::HloInstructionSet
@@ -231,6 +241,8 @@
// fusing operations that require duplication later depending on
// is_expensive_().
HloInstructionSet do_not_duplicate;
+ absl::flat_hash_map<std::pair<HloInstruction*, HloInstruction*>, bool>
+ can_fuse_on_all_paths_result_cache;
for (HloInstruction* consumer : post_order) {
for (HloInstruction* producer : consumer->operands()) {
if (do_not_duplicate.count(producer) > 0) {
@@ -286,7 +298,8 @@
// A will be not allowed to be fused into B, as it cannot be fused via
// all paths.
if (producer->IsFusible() &&
- CanFuseOnAllPaths(producer, consumer, do_not_duplicate)) {
+ CanFuseOnAllPaths(producer, consumer, do_not_duplicate,
+ &can_fuse_on_all_paths_result_cache)) {
continue;
}
do_not_duplicate.insert(producer);
@@ -417,7 +430,7 @@
private:
std::vector<HloInstruction*> post_order_;
- tensorflow::gtl::FlatMap<HloInstruction*, int> post_order_index_;
+ absl::flat_hash_map<HloInstruction*, int> post_order_index_;
};
} // namespace
diff --git a/tensorflow/compiler/xla/service/instruction_fusion.h b/tensorflow/compiler/xla/service/instruction_fusion.h
index 7e1196f..da2032f 100644
--- a/tensorflow/compiler/xla/service/instruction_fusion.h
+++ b/tensorflow/compiler/xla/service/instruction_fusion.h
@@ -1,3 +1,4 @@
+#include "absl/container/flat_hash_map.h"
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
@@ -151,8 +152,15 @@
// Whether or not we can fuse producer into consumer on all paths
// from the producer to the consumer where nodes are HLOs and edges are uses.
- bool CanFuseOnAllPaths(HloInstruction* producer, HloInstruction* consumer,
- const HloInstructionSet& do_not_fuse);
+ //
+ // A map from <producer, consumer> to a bool is required as the result cache
+ // to store and query the results of calls to this function, in order to avoid
+ // repeated computations.
+ bool CanFuseOnAllPaths(
+ HloInstruction* producer, HloInstruction* consumer,
+ const HloInstructionSet& do_not_fuse,
+ absl::flat_hash_map<std::pair<HloInstruction*, HloInstruction*>, bool>*
+ result_cache);
// Computes the set of nodes that we do not want to fuse into any of their
// consumers based on a global analysis of the HLO graph.
diff --git a/tensorflow/compiler/xla/service/layout_assignment.h b/tensorflow/compiler/xla/service/layout_assignment.h
index e29c199..1591256 100644
--- a/tensorflow/compiler/xla/service/layout_assignment.h
+++ b/tensorflow/compiler/xla/service/layout_assignment.h
@@ -25,6 +25,7 @@
#include <utility>
#include <vector>
+#include "absl/container/flat_hash_map.h"
#include "tensorflow/compiler/xla/service/computation_layout.h"
#include "tensorflow/compiler/xla/service/hlo_computation.h"
#include "tensorflow/compiler/xla/service/hlo_instruction.h"
@@ -38,7 +39,6 @@
#include "tensorflow/compiler/xla/types.h"
#include "tensorflow/compiler/xla/xla_data.pb.h"
#include "tensorflow/core/lib/core/status.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/lib/gtl/flatset.h"
#include "tensorflow/core/platform/types.h"
@@ -228,8 +228,8 @@
// Array-shaped buffers which have not yet been constrained.
std::set<LogicalBuffer::Id> unconstrained_buffer_ids_;
- mutable tensorflow::gtl::FlatMap<const HloInstruction*,
- std::unique_ptr<PointsToSet::BufferSet>>
+ mutable absl::flat_hash_map<const HloInstruction*,
+ std::unique_ptr<PointsToSet::BufferSet>>
buffer_sets_cache_;
HloComputation* computation_;
diff --git a/tensorflow/compiler/xla/service/llvm_ir/BUILD b/tensorflow/compiler/xla/service/llvm_ir/BUILD
index 540bbb7..3934d2e 100644
--- a/tensorflow/compiler/xla/service/llvm_ir/BUILD
+++ b/tensorflow/compiler/xla/service/llvm_ir/BUILD
@@ -38,6 +38,7 @@
"//tensorflow/compiler/xla/service:hlo",
"//tensorflow/compiler/xla/service:logical_buffer",
"//tensorflow/core:lib",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/strings",
"@llvm//:core",
],
diff --git a/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.h b/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.h
index 8d9fa99..88cde2d 100644
--- a/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.h
+++ b/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.h
@@ -16,13 +16,13 @@
#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_LLVM_IR_ALIAS_ANALYSIS_H_
#define TENSORFLOW_COMPILER_XLA_SERVICE_LLVM_IR_ALIAS_ANALYSIS_H_
+#include "absl/container/flat_hash_map.h"
#include "absl/strings/str_cat.h"
#include "llvm/IR/Module.h"
#include "tensorflow/compiler/xla/service/buffer_assignment.h"
#include "tensorflow/compiler/xla/service/hlo_instruction.h"
#include "tensorflow/compiler/xla/service/llvm_ir/ir_array.h"
#include "tensorflow/compiler/xla/types.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/lib/gtl/flatset.h"
namespace xla {
@@ -77,14 +77,14 @@
// A map from a buffer slice to metadata corresponding to its alias.scope
// metadata. The index kParameterAliasSet is used to hold aliasing
// information for parameters.
- tensorflow::gtl::FlatMap<BufferAllocation::Slice, llvm::MDNode*,
- BufferAllocation::Slice::Hasher>
+ absl::flat_hash_map<BufferAllocation::Slice, llvm::MDNode*,
+ BufferAllocation::Slice::Hasher>
alias_scope_metadata_;
// A map from a buffer slice to metadata corresponding to its noalias
// metadata.
- tensorflow::gtl::FlatMap<BufferAllocation::Slice, llvm::MDNode*,
- BufferAllocation::Slice::Hasher>
+ absl::flat_hash_map<BufferAllocation::Slice, llvm::MDNode*,
+ BufferAllocation::Slice::Hasher>
noalias_metadata_;
};
diff --git a/tensorflow/compiler/xla/service/logical_buffer_analysis.cc b/tensorflow/compiler/xla/service/logical_buffer_analysis.cc
index eaa0959..ec52a24 100644
--- a/tensorflow/compiler/xla/service/logical_buffer_analysis.cc
+++ b/tensorflow/compiler/xla/service/logical_buffer_analysis.cc
@@ -54,7 +54,7 @@
// so reserve 10% more than the number of instructions to avoid frequent
// resizes.
logical_buffers_.clear();
- logical_buffers_.reserve((module_->NumUniqueInstructionIds() * 11) / 10);
+ logical_buffers_.reserve((module_->instruction_count() * 11) / 10);
// We filter out fusion computations, and get to them through fusion
// instructions. This is because it's possible to have orphaned (unreachable)
diff --git a/tensorflow/compiler/xla/service/multi_output_fusion.cc b/tensorflow/compiler/xla/service/multi_output_fusion.cc
index b9ec31c..95b1c20 100644
--- a/tensorflow/compiler/xla/service/multi_output_fusion.cc
+++ b/tensorflow/compiler/xla/service/multi_output_fusion.cc
@@ -18,7 +18,7 @@
#include "tensorflow/compiler/xla/service/hlo_instruction.h"
#include "tensorflow/compiler/xla/service/hlo_opcode.h"
#include "tensorflow/compiler/xla/shape_util.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
+#include "tensorflow/core/lib/gtl/flatset.h"
#include "tensorflow/core/platform/types.h"
namespace xla {
diff --git a/tensorflow/compiler/xla/service/multi_output_fusion.h b/tensorflow/compiler/xla/service/multi_output_fusion.h
index 0344626..9508ab2 100644
--- a/tensorflow/compiler/xla/service/multi_output_fusion.h
+++ b/tensorflow/compiler/xla/service/multi_output_fusion.h
@@ -19,6 +19,7 @@
#include <queue>
#include <vector>
+#include "absl/container/flat_hash_map.h"
#include "absl/strings/string_view.h"
#include "tensorflow/compiler/xla/service/hlo_module.h"
#include "tensorflow/compiler/xla/service/hlo_pass_interface.h"
@@ -126,7 +127,7 @@
std::vector<FusionCandidate> candidates_;
// A map that maps an instruction to the index_.
- tensorflow::gtl::FlatMap<HloInstruction*, int> candidates_index_;
+ absl::flat_hash_map<HloInstruction*, int> candidates_index_;
// The reachability map of current computation.
std::unique_ptr<HloReachabilityMap> reachability_;
diff --git a/tensorflow/compiler/xla/service/name_uniquer.h b/tensorflow/compiler/xla/service/name_uniquer.h
index 6dd89c2..1ac60f1 100644
--- a/tensorflow/compiler/xla/service/name_uniquer.h
+++ b/tensorflow/compiler/xla/service/name_uniquer.h
@@ -18,9 +18,9 @@
#include <string>
+#include "absl/container/flat_hash_map.h"
#include "absl/strings/string_view.h"
#include "tensorflow/compiler/xla/types.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/lib/gtl/flatset.h"
#include "tensorflow/core/platform/macros.h"
@@ -78,7 +78,7 @@
// Map from name prefix to the generator data structure which tracks used
// identifiers and generates new ones.
- tensorflow::gtl::FlatMap<string, SequentialIdGenerator> generated_names_;
+ absl::flat_hash_map<string, SequentialIdGenerator> generated_names_;
TF_DISALLOW_COPY_AND_ASSIGN(NameUniquer);
};
diff --git a/tensorflow/compiler/xla/service/reduce_precision_insertion.h b/tensorflow/compiler/xla/service/reduce_precision_insertion.h
index 4bb2242..0b4e82e 100644
--- a/tensorflow/compiler/xla/service/reduce_precision_insertion.h
+++ b/tensorflow/compiler/xla/service/reduce_precision_insertion.h
@@ -22,7 +22,6 @@
#include "tensorflow/compiler/xla/service/hlo_module.h"
#include "tensorflow/compiler/xla/service/hlo_pass_interface.h"
#include "tensorflow/compiler/xla/service/hlo_pass_pipeline.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
namespace xla {
diff --git a/tensorflow/compiler/xla/service/shape_inference.cc b/tensorflow/compiler/xla/service/shape_inference.cc
index 7194b2c..6ccea9d 100644
--- a/tensorflow/compiler/xla/service/shape_inference.cc
+++ b/tensorflow/compiler/xla/service/shape_inference.cc
@@ -2380,7 +2380,9 @@
!std::is_permutation(dimensions.begin(), dimensions.end(),
indices.begin())) {
return InvalidArgument(
- "Transpose dimensions not a permutation of the operand dimensions.");
+ "Transpose dimensions [%s] are not a permutation of the operand "
+ "dimensions (operand shape is %s).",
+ StrJoin(dimensions, ","), ShapeUtil::HumanString(operand));
}
// Permute(dimensions,input) computes output[dimensions[i]]=input[i]. However,
diff --git a/tensorflow/compiler/xla/service/tuple_points_to_analysis.h b/tensorflow/compiler/xla/service/tuple_points_to_analysis.h
index a9e8a51..78392d3 100644
--- a/tensorflow/compiler/xla/service/tuple_points_to_analysis.h
+++ b/tensorflow/compiler/xla/service/tuple_points_to_analysis.h
@@ -36,7 +36,6 @@
#include "tensorflow/compiler/xla/xla_data.pb.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/gtl/compactptrset.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/lib/gtl/flatset.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/types.h"
diff --git a/tensorflow/compiler/xla/service/while_loop_constant_sinking.cc b/tensorflow/compiler/xla/service/while_loop_constant_sinking.cc
index 5614582..067cfcc 100644
--- a/tensorflow/compiler/xla/service/while_loop_constant_sinking.cc
+++ b/tensorflow/compiler/xla/service/while_loop_constant_sinking.cc
@@ -18,7 +18,6 @@
#include "absl/container/inlined_vector.h"
#include "tensorflow/compiler/xla/service/while_util.h"
#include "tensorflow/compiler/xla/util.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
namespace xla {
diff --git a/tensorflow/compiler/xla/service/while_loop_invariant_code_motion.cc b/tensorflow/compiler/xla/service/while_loop_invariant_code_motion.cc
index e8fe33e..2590473 100644
--- a/tensorflow/compiler/xla/service/while_loop_invariant_code_motion.cc
+++ b/tensorflow/compiler/xla/service/while_loop_invariant_code_motion.cc
@@ -15,17 +15,17 @@
#include "tensorflow/compiler/xla/service/while_loop_invariant_code_motion.h"
#include "absl/algorithm/container.h"
+#include "absl/container/flat_hash_map.h"
#include "absl/container/inlined_vector.h"
#include "tensorflow/compiler/xla/service/tuple_util.h"
#include "tensorflow/compiler/xla/service/while_util.h"
#include "tensorflow/compiler/xla/util.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/lib/gtl/flatset.h"
namespace xla {
+using absl::flat_hash_map;
using absl::InlinedVector;
-using tensorflow::gtl::FlatMap;
using tensorflow::gtl::FlatSet;
// Copies `to_hoist` to the computation containing `while_instr`, hoisting its
@@ -34,7 +34,7 @@
// function hoists the operands in `unhoisted_invariant_instructions` and moves
// them into `hoisted_instructions`.
static void CreateLoopInvariantCopy(
- FlatMap<HloInstruction*, HloInstruction*>* hoisted_instructions,
+ flat_hash_map<HloInstruction*, HloInstruction*>* hoisted_instructions,
FlatSet<HloInstruction*>* unhoisted_invariant_instructions,
HloInstruction* while_instr, HloInstruction* to_hoist) {
HloComputation* parent_of_while = while_instr->parent();
@@ -147,7 +147,7 @@
// Maps instructions in the while body to instructions hoisted outside the
// while that compute the same value.
- FlatMap<HloInstruction*, HloInstruction*> hoisted_instructions;
+ flat_hash_map<HloInstruction*, HloInstruction*> hoisted_instructions;
// Contains instructions that can be legally hoisted, but were deemed to be
// unprofitable to be hoisted alone by NotWorthHoistingIndividually. When we
diff --git a/tensorflow/compiler/xla/service/while_loop_simplifier.cc b/tensorflow/compiler/xla/service/while_loop_simplifier.cc
index 9a74f22..07de849 100644
--- a/tensorflow/compiler/xla/service/while_loop_simplifier.cc
+++ b/tensorflow/compiler/xla/service/while_loop_simplifier.cc
@@ -14,12 +14,12 @@
==============================================================================*/
#include "tensorflow/compiler/xla/service/while_loop_simplifier.h"
+#include "absl/container/flat_hash_map.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_join.h"
#include "absl/types/optional.h"
#include "tensorflow/compiler/xla/service/call_inliner.h"
#include "tensorflow/compiler/xla/service/while_loop_analysis.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
namespace xla {
@@ -181,7 +181,7 @@
used_tuple_indices.end());
std::sort(new_to_old_tuple_idx.begin(), new_to_old_tuple_idx.end());
- tensorflow::gtl::FlatMap<int64, int64> old_to_new_tuple_idx;
+ absl::flat_hash_map<int64, int64> old_to_new_tuple_idx;
for (int64 new_idx = 0; new_idx < new_to_old_tuple_idx.size(); ++new_idx) {
int64 old_idx = new_to_old_tuple_idx[new_idx];
old_to_new_tuple_idx[old_idx] = new_idx;
@@ -405,7 +405,7 @@
// build a map from the tuple element index to the constant value. Limit this
// to scalar constant values because propagating array constants can regress
// performance by forcing us to copy constants.
- tensorflow::gtl::FlatMap<int, const HloInstruction*> index_to_constant;
+ absl::flat_hash_map<int, const HloInstruction*> index_to_constant;
for (int i = 0; i < root_operands.size(); i++) {
HloInstruction* instr = root_operands[i];
if (instr->opcode() == HloOpcode::kGetTupleElement &&
diff --git a/tensorflow/compiler/xla/tests/BUILD b/tensorflow/compiler/xla/tests/BUILD
index f474ecb1..06b6330 100644
--- a/tensorflow/compiler/xla/tests/BUILD
+++ b/tensorflow/compiler/xla/tests/BUILD
@@ -422,6 +422,7 @@
"//tensorflow/core:regexp_internal",
"//tensorflow/core:test",
"@com_google_absl//absl/algorithm:container",
+ "@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/strings",
],
)
diff --git a/tensorflow/compiler/xla/tests/xla_hlo_profile_test.cc b/tensorflow/compiler/xla/tests/xla_hlo_profile_test.cc
index db5a824..a6e70eb 100644
--- a/tensorflow/compiler/xla/tests/xla_hlo_profile_test.cc
+++ b/tensorflow/compiler/xla/tests/xla_hlo_profile_test.cc
@@ -17,6 +17,7 @@
#include <vector>
#include "absl/algorithm/container.h"
+#include "absl/container/flat_hash_map.h"
#include "absl/strings/match.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_split.h"
@@ -32,7 +33,6 @@
#include "tensorflow/compiler/xla/tests/test_macros.h"
#include "tensorflow/compiler/xla/tests/test_utils.h"
#include "tensorflow/core/lib/core/status_test_util.h"
-#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/platform/regexp.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/types.h"
@@ -83,7 +83,7 @@
Status ParseOneProfileOutputLine(
const string& line, bool expect_hlo,
- gtl::FlatMap<string, ParsedProfileOutputLine>* parsed_results,
+ absl::flat_hash_map<string, ParsedProfileOutputLine>* parsed_results,
absl::Span<const absl::string_view> opcodes_to_ignore = {}) {
string separator = "[^:]*:: +";
string match_percentage = R"(\d+\.\d*% +\d+Σ)";
@@ -208,7 +208,7 @@
std::vector<string> profile_output_lines =
absl::StrSplit(profile_output, '\n');
- gtl::FlatMap<string, ParsedProfileOutputLine> parsed_profile_lines;
+ absl::flat_hash_map<string, ParsedProfileOutputLine> parsed_profile_lines;
TF_ASSERT_OK(ParseOneProfileOutputLine(
profile_output_lines[1], /*expect_hlo=*/false, &parsed_profile_lines));
@@ -314,7 +314,7 @@
ASSERT_NE(while_body_profile_end, profile_output_lines.end());
- gtl::FlatMap<string, ParsedProfileOutputLine> parsed_profile_lines;
+ absl::flat_hash_map<string, ParsedProfileOutputLine> parsed_profile_lines;
for (auto while_body_profile_i = while_body_profile_start + 1;
while_body_profile_i != while_body_profile_end; while_body_profile_i++) {
diff --git a/tensorflow/compiler/xrt/ops/xrt_execute_op.cc b/tensorflow/compiler/xrt/ops/xrt_execute_op.cc
index fda4c31..40ec1b0 100644
--- a/tensorflow/compiler/xrt/ops/xrt_execute_op.cc
+++ b/tensorflow/compiler/xrt/ops/xrt_execute_op.cc
@@ -21,7 +21,7 @@
namespace tensorflow {
REGISTER_OP("XRTExecute")
- .Attr("Ninputs: int")
+ .Attr("Ninputs: int >= 0")
.Input("computation_handle: int64")
.Input("execution_config: string")
.Input("input_handles: Ninputs * int64")
diff --git a/tensorflow/compiler/xrt/tests/raw_api_test.cc b/tensorflow/compiler/xrt/tests/raw_api_test.cc
index 2952feb..f590fbf 100644
--- a/tensorflow/compiler/xrt/tests/raw_api_test.cc
+++ b/tensorflow/compiler/xrt/tests/raw_api_test.cc
@@ -108,6 +108,14 @@
return equal;
}
+xla::XlaComputation OnePlusTwo() {
+ xla::XlaBuilder builder("OnePlusTwo");
+ auto c0 = xla::ConstantR0(&builder, 1.0f);
+ auto c1 = xla::ConstantR0(&builder, 2.0f);
+ xla::Add(c0, c1);
+ return builder.Build().ValueOrDie();
+}
+
xla::XlaComputation AddAndScale() {
xla::XlaBuilder builder("AddAndScale");
auto p0 = xla::Parameter(&builder, 0,
@@ -346,6 +354,39 @@
EXPECT_TRUE(CompareLiteralToLiteralProto(expected, response));
}
+TEST(RawApiTest, CompileAndExecuteZeroArg) {
+ xrt::XLAComputation c;
+ auto config = c.mutable_config();
+ auto shapes = config->mutable_program_shape();
+ *shapes->mutable_result() = xla::ShapeUtil::MakeShape(xla::F32, {});
+
+ xrt::XRTExecutionConfig e;
+ e.set_release_input_handles(true);
+ e.set_release_compilation_handle(true);
+ StoreComputationSnapshot(OnePlusTwo(), c.mutable_hlo_snapshot());
+
+ Scope root = Scope::NewRootScope().WithDevice(DeviceFromFlag());
+ auto e_config =
+ ops::Const(root.WithDevice("/device:CPU:0"), e.SerializeAsString());
+ auto computation =
+ ops::Const(root.WithDevice("/device:CPU:0"), c.SerializeAsString());
+ auto c_handle = ops::XRTCompile(root, computation);
+ auto result = ops::XRTExecute(root, c_handle, e_config,
+ std::initializer_list<Input>({}));
+ auto read_back = ops::XRTReadLiteralAndRelease(root, result);
+ TF_ASSERT_OK(root.status());
+
+ ClientSession session(root);
+ std::vector<Tensor> outputs;
+ TF_EXPECT_OK(session.Run({read_back}, &outputs));
+
+ xla::LiteralProto response;
+ EXPECT_TRUE(response.ParseFromString(outputs[0].scalar<string>()()));
+
+ auto expected = xla::LiteralUtil::CreateR0<float>(3.0f);
+ EXPECT_TRUE(CompareLiteralToLiteralProto(expected, response));
+}
+
TEST(RawApiTest, CompileAndExecuteReturnTuple) {
xrt::XLAAllocation p0;
p0.set_device_ordinal(0);
diff --git a/tensorflow/contrib/BUILD b/tensorflow/contrib/BUILD
index ae5ca32..fbe0573 100644
--- a/tensorflow/contrib/BUILD
+++ b/tensorflow/contrib/BUILD
@@ -112,29 +112,22 @@
"//tensorflow/python:util",
"//tensorflow/python/estimator:estimator_py",
] + if_mpi(["//tensorflow/contrib/mpi_collectives:mpi_collectives_py"]) + select({
- "//tensorflow:with_kafka_support_windows_override": [],
- "//tensorflow:with_kafka_support": [
- "//tensorflow/contrib/kafka",
- ],
- "//conditions:default": [],
- }) + select({
- "//tensorflow:with_aws_support_windows_override": [],
- "//tensorflow:with_aws_support": [
- "//tensorflow/contrib/kinesis",
- ],
- "//conditions:default": [],
- }) + if_not_windows_cuda([
- "//tensorflow/contrib/fused_conv:fused_conv_py", # unresolved symbols, need to export more symbols
- ]) + if_not_windows([
- ]) + select({
"//tensorflow:linux_s390x": [],
"//tensorflow:windows": [],
"//conditions:default": [
"//tensorflow/contrib/bigtable",
"//tensorflow/contrib/cloud:cloud_py",
+ "//tensorflow/contrib/fused_conv:fused_conv_py", # unresolved symbols, need to export more symbols
+ "//tensorflow/contrib/kafka",
+ "//tensorflow/contrib/kinesis",
"//tensorflow/contrib/tensorrt:init_py",
"//tensorflow/contrib/ffmpeg:ffmpeg_ops_py",
],
+ }) + select({
+ "//tensorflow:with_ignite_support": [
+ "//tensorflow/contrib/ignite",
+ ],
+ "//conditions:default": [],
}),
)
@@ -144,7 +137,6 @@
deps = [
"//tensorflow/contrib/boosted_trees:boosted_trees_kernels",
"//tensorflow/contrib/coder:all_kernels",
- "//tensorflow/contrib/data/kernels:dataset_kernels",
"//tensorflow/contrib/factorization/kernels:all_kernels",
"//tensorflow/contrib/hadoop:dataset_kernels",
"//tensorflow/contrib/input_pipeline:input_pipeline_ops_kernels",
@@ -159,20 +151,14 @@
] + if_mpi(["//tensorflow/contrib/mpi_collectives:mpi_collectives_py"]) + if_cuda([
"//tensorflow/contrib/nccl:nccl_kernels",
]) + select({
- "//tensorflow:with_kafka_support_windows_override": [],
- "//tensorflow:with_kafka_support": [
+ "//tensorflow:linux_s390x": [],
+ "//tensorflow:windows": [],
+ "//conditions:default": [
"//tensorflow/contrib/kafka:dataset_kernels",
- ],
- "//conditions:default": [],
- }) + select({
- "//tensorflow:with_aws_support_windows_override": [],
- "//tensorflow:with_aws_support": [
"//tensorflow/contrib/kinesis:dataset_kernels",
+ "//tensorflow/contrib/tensorrt:trt_engine_op_kernel",
],
- "//conditions:default": [],
- }) + if_not_windows([
- "//tensorflow/contrib/tensorrt:trt_engine_op_kernel",
- ]),
+ }),
)
cc_library(
@@ -181,8 +167,6 @@
deps = [
"//tensorflow/contrib/boosted_trees:boosted_trees_ops_op_lib",
"//tensorflow/contrib/coder:all_ops",
- "//tensorflow/contrib/data:dataset_ops_op_lib",
- "//tensorflow/contrib/data:indexed_dataset_ops_op_lib",
"//tensorflow/contrib/factorization:all_ops",
"//tensorflow/contrib/framework:all_ops",
"//tensorflow/contrib/hadoop:dataset_ops_op_lib",
@@ -198,18 +182,17 @@
"//tensorflow/contrib/text:all_ops",
"//tensorflow/contrib/tpu:all_ops",
] + select({
- "//tensorflow:with_kafka_support_windows_override": [],
- "//tensorflow:with_kafka_support": [
+ "//tensorflow:linux_s390x": [],
+ "//tensorflow:windows": [],
+ "//conditions:default": [
"//tensorflow/contrib/kafka:dataset_ops_op_lib",
- ],
- "//conditions:default": [],
- }) + select({
- "//tensorflow:with_aws_support_windows_override": [],
- "//tensorflow:with_aws_support": [
"//tensorflow/contrib/kinesis:dataset_ops_op_lib",
+ "//tensorflow/contrib/tensorrt:trt_engine_op_op_lib",
+ ],
+ }) + select({
+ "//tensorflow:with_ignite_support": [
+ "//tensorflow/contrib/ignite:dataset_ops_op_lib",
],
"//conditions:default": [],
- }) + if_not_windows([
- "//tensorflow/contrib/tensorrt:trt_engine_op_op_lib",
- ]),
+ }),
)
diff --git a/tensorflow/contrib/bigtable/README.md b/tensorflow/contrib/bigtable/README.md
index f33eaf7..2c44abe 100644
--- a/tensorflow/contrib/bigtable/README.md
+++ b/tensorflow/contrib/bigtable/README.md
@@ -203,7 +203,7 @@
start = tf.string_join(['training_data_', start_idx_str])
end = tf.string_join(['training_data_', end_idx_str])
return table.scan_range(start_idx, end_idx, columns=columns)
-ds = ds.apply(tf.contrib.data.parallel_interleave(
+ds = ds.apply(tf.data.experimental.parallel_interleave(
interleave_fn, cycle_length=NUM_PARALLEL_READS, prefetch_input_elements=1))
```
@@ -249,7 +249,7 @@
- ...
- fake-data-23498103
"""
- counter_dataset = tf.contrib.data.Counter()
+ counter_dataset = tf.data.experimental.Counter()
width = 8
row_key_prefix = 'fake-data-'
ds = counter_dataset.map(lambda index: tf.as_string(index,
diff --git a/tensorflow/contrib/bigtable/kernels/bigtable_lookup_dataset_op.cc b/tensorflow/contrib/bigtable/kernels/bigtable_lookup_dataset_op.cc
index 11f530e..2c63171 100644
--- a/tensorflow/contrib/bigtable/kernels/bigtable_lookup_dataset_op.cc
+++ b/tensorflow/contrib/bigtable/kernels/bigtable_lookup_dataset_op.cc
@@ -28,6 +28,7 @@
DatasetBase** output) override {
BigtableTableResource* table;
OP_REQUIRES_OK(ctx, LookupResource(ctx, HandleFromInput(ctx, 1), &table));
+ core::ScopedUnref scoped_unref(table);
std::vector<string> column_families;
std::vector<string> columns;
diff --git a/tensorflow/contrib/bigtable/kernels/bigtable_prefix_key_dataset_op.cc b/tensorflow/contrib/bigtable/kernels/bigtable_prefix_key_dataset_op.cc
index 5cab729..92a3658 100644
--- a/tensorflow/contrib/bigtable/kernels/bigtable_prefix_key_dataset_op.cc
+++ b/tensorflow/contrib/bigtable/kernels/bigtable_prefix_key_dataset_op.cc
@@ -31,6 +31,7 @@
BigtableTableResource* resource;
OP_REQUIRES_OK(ctx,
LookupResource(ctx, HandleFromInput(ctx, 0), &resource));
+ core::ScopedUnref scoped_unref(resource);
*output = new Dataset(ctx, resource, std::move(prefix));
}
diff --git a/tensorflow/contrib/bigtable/kernels/bigtable_range_key_dataset_op.cc b/tensorflow/contrib/bigtable/kernels/bigtable_range_key_dataset_op.cc
index 4dc4647..bd8805a 100644
--- a/tensorflow/contrib/bigtable/kernels/bigtable_range_key_dataset_op.cc
+++ b/tensorflow/contrib/bigtable/kernels/bigtable_range_key_dataset_op.cc
@@ -34,6 +34,7 @@
BigtableTableResource* resource;
OP_REQUIRES_OK(ctx,
LookupResource(ctx, HandleFromInput(ctx, 0), &resource));
+ core::ScopedUnref scoped_unref(resource);
*output =
new Dataset(ctx, resource, std::move(start_key), std::move(end_key));
diff --git a/tensorflow/contrib/bigtable/kernels/bigtable_sample_key_pairs_dataset_op.cc b/tensorflow/contrib/bigtable/kernels/bigtable_sample_key_pairs_dataset_op.cc
index 736775b..01608dc 100644
--- a/tensorflow/contrib/bigtable/kernels/bigtable_sample_key_pairs_dataset_op.cc
+++ b/tensorflow/contrib/bigtable/kernels/bigtable_sample_key_pairs_dataset_op.cc
@@ -38,6 +38,7 @@
BigtableTableResource* resource;
OP_REQUIRES_OK(ctx,
LookupResource(ctx, HandleFromInput(ctx, 0), &resource));
+ core::ScopedUnref scoped_unref(resource);
OP_REQUIRES(ctx, prefix.empty() || start_key.empty(),
errors::InvalidArgument(
diff --git a/tensorflow/contrib/bigtable/kernels/bigtable_sample_keys_dataset_op.cc b/tensorflow/contrib/bigtable/kernels/bigtable_sample_keys_dataset_op.cc
index 208b7b3..9b60e0a 100644
--- a/tensorflow/contrib/bigtable/kernels/bigtable_sample_keys_dataset_op.cc
+++ b/tensorflow/contrib/bigtable/kernels/bigtable_sample_keys_dataset_op.cc
@@ -28,6 +28,7 @@
BigtableTableResource* resource;
OP_REQUIRES_OK(ctx,
LookupResource(ctx, HandleFromInput(ctx, 0), &resource));
+ core::ScopedUnref scoped_unref(resource);
*output = new Dataset(ctx, resource);
}
diff --git a/tensorflow/contrib/bigtable/kernels/bigtable_scan_dataset_op.cc b/tensorflow/contrib/bigtable/kernels/bigtable_scan_dataset_op.cc
index 9407855..688289a 100644
--- a/tensorflow/contrib/bigtable/kernels/bigtable_scan_dataset_op.cc
+++ b/tensorflow/contrib/bigtable/kernels/bigtable_scan_dataset_op.cc
@@ -67,6 +67,7 @@
BigtableTableResource* resource;
OP_REQUIRES_OK(ctx,
LookupResource(ctx, HandleFromInput(ctx, 0), &resource));
+ core::ScopedUnref scoped_unref(resource);
const uint64 num_outputs = columns.size() + 1;
std::vector<PartialTensorShape> output_shapes;
diff --git a/tensorflow/contrib/bigtable/python/ops/bigtable_api.py b/tensorflow/contrib/bigtable/python/ops/bigtable_api.py
index cf56822..7c87b0d 100644
--- a/tensorflow/contrib/bigtable/python/ops/bigtable_api.py
+++ b/tensorflow/contrib/bigtable/python/ops/bigtable_api.py
@@ -31,8 +31,8 @@
from six import string_types
from tensorflow.contrib.bigtable.ops import gen_bigtable_ops
-from tensorflow.contrib.data.python.ops import interleave_ops
from tensorflow.contrib.util import loader
+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.framework import dtypes
@@ -228,7 +228,7 @@
"""Retrieves a sampling of row keys from the Bigtable table.
This dataset is most often used in conjunction with
- `tf.contrib.data.parallel_interleave` to construct a set of ranges for
+ `tf.data.experimental.parallel_interleave` to construct a set of ranges for
scanning in parallel.
Returns:
diff --git a/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch.py b/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch.py
index c7eb249..8531e97 100644
--- a/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch.py
+++ b/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch.py
@@ -402,13 +402,13 @@
self._feature_columns = feature_columns
self._learner_config_serialized = learner_config.SerializeToString()
self._num_quantiles = num_quantiles
- self._max_tree_depth = variables.Variable(
+ self._max_tree_depth = variables.VariableV1(
initial_value=self._learner_config.constraints.max_tree_depth)
- self._attempted_trees = variables.Variable(
+ self._attempted_trees = variables.VariableV1(
initial_value=array_ops.zeros([], dtypes.int64),
trainable=False,
name="attempted_trees")
- self._finalized_trees = variables.Variable(
+ self._finalized_trees = variables.VariableV1(
initial_value=array_ops.zeros([], dtypes.int64),
trainable=False,
name="finalized_trees")
@@ -770,28 +770,28 @@
fc_name_idx += 1
# Create ensemble stats variables.
- num_layer_examples = variables.Variable(
+ num_layer_examples = variables.VariableV1(
initial_value=array_ops.zeros([], dtypes.int64),
name="num_layer_examples",
trainable=False)
- num_layer_steps = variables.Variable(
+ num_layer_steps = variables.VariableV1(
initial_value=array_ops.zeros([], dtypes.int64),
name="num_layer_steps",
trainable=False)
- num_layers = variables.Variable(
+ num_layers = variables.VariableV1(
initial_value=array_ops.zeros([], dtypes.int64),
name="num_layers",
trainable=False)
- active_tree = variables.Variable(
+ active_tree = variables.VariableV1(
initial_value=array_ops.zeros([], dtypes.int64),
name="active_tree",
trainable=False)
- active_layer = variables.Variable(
+ active_layer = variables.VariableV1(
initial_value=array_ops.zeros([], dtypes.int64),
name="active_layer",
trainable=False)
# Variable that becomes false once bias centering is done.
- continue_centering = variables.Variable(
+ continue_centering = variables.VariableV1(
initial_value=self._center_bias,
name="continue_centering",
trainable=False)
diff --git a/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch_test.py b/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch_test.py
index 9d9941f..6d20a2e 100644
--- a/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch_test.py
+++ b/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch_test.py
@@ -239,7 +239,7 @@
predictions = array_ops.constant(
[[0.0], [1.0], [0.0], [2.0]], dtype=dtypes.float32)
partition_ids = array_ops.zeros([4], dtypes.int32)
- ensemble_stamp = variables.Variable(
+ ensemble_stamp = variables.VariableV1(
initial_value=0,
name="ensemble_stamp",
trainable=False,
@@ -503,7 +503,7 @@
predictions = array_ops.constant(
[[0.0], [1.0], [0.0], [2.0]], dtype=dtypes.float32)
partition_ids = array_ops.zeros([4], dtypes.int32)
- ensemble_stamp = variables.Variable(
+ ensemble_stamp = variables.VariableV1(
initial_value=0,
name="ensemble_stamp",
trainable=False,
@@ -607,7 +607,7 @@
predictions = array_ops.constant(
[[0.0], [1.0], [0.0], [2.0]], dtype=dtypes.float32)
partition_ids = array_ops.zeros([4], dtypes.int32)
- ensemble_stamp = variables.Variable(
+ ensemble_stamp = variables.VariableV1(
initial_value=0,
name="ensemble_stamp",
trainable=False,
@@ -711,7 +711,7 @@
predictions = array_ops.constant(
[[0.0], [1.0], [0.0], [2.0]], dtype=dtypes.float32)
partition_ids = array_ops.zeros([4], dtypes.int32)
- ensemble_stamp = variables.Variable(
+ ensemble_stamp = variables.VariableV1(
initial_value=0,
name="ensemble_stamp",
trainable=False,
@@ -783,7 +783,7 @@
predictions = array_ops.constant(
[[0.0], [1.0], [0.0], [2.0]], dtype=dtypes.float32)
partition_ids = array_ops.zeros([4], dtypes.int32)
- ensemble_stamp = variables.Variable(
+ ensemble_stamp = variables.VariableV1(
initial_value=0,
name="ensemble_stamp",
trainable=False,
@@ -847,7 +847,7 @@
predictions = array_ops.constant(
[[0.0], [1.0], [0.0], [2.0]], dtype=dtypes.float32)
partition_ids = array_ops.zeros([4], dtypes.int32)
- ensemble_stamp = variables.Variable(
+ ensemble_stamp = variables.VariableV1(
initial_value=0,
name="ensemble_stamp",
trainable=False,
@@ -1090,7 +1090,7 @@
weights = array_ops.ones([batch_size, 1], dtypes.float32)
partition_ids = array_ops.zeros([batch_size], dtypes.int32)
- ensemble_stamp = variables.Variable(
+ ensemble_stamp = variables.VariableV1(
initial_value=0,
name="ensemble_stamp",
trainable=False,
@@ -1194,7 +1194,7 @@
weights = array_ops.ones([batch_size, 1], dtypes.float32)
partition_ids = array_ops.zeros([batch_size], dtypes.int32)
- ensemble_stamp = variables.Variable(
+ ensemble_stamp = variables.VariableV1(
initial_value=0,
name="ensemble_stamp",
trainable=False,
@@ -1299,7 +1299,7 @@
weights = array_ops.ones([batch_size, 1], dtypes.float32)
partition_ids = array_ops.zeros([batch_size], dtypes.int32)
- ensemble_stamp = variables.Variable(
+ ensemble_stamp = variables.VariableV1(
initial_value=0,
name="ensemble_stamp",
trainable=False,
@@ -1405,7 +1405,7 @@
predictions = array_ops.constant(
[[0.0], [1.0], [0.0], [2.0]], dtype=dtypes.float32)
partition_ids = array_ops.zeros([4], dtypes.int32)
- ensemble_stamp = variables.Variable(
+ ensemble_stamp = variables.VariableV1(
initial_value=0,
name="ensemble_stamp",
trainable=False,
@@ -1524,7 +1524,7 @@
predictions = array_ops.constant(
[[0.0], [1.0], [0.0], [2.0]], dtype=dtypes.float32)
partition_ids = array_ops.zeros([4], dtypes.int32)
- ensemble_stamp = variables.Variable(
+ ensemble_stamp = variables.VariableV1(
initial_value=0,
name="ensemble_stamp",
trainable=False,
@@ -1656,7 +1656,7 @@
predictions = array_ops.constant(
[[0.0], [1.0], [0.0], [2.0]], dtype=dtypes.float32)
partition_ids = array_ops.zeros([4], dtypes.int32)
- ensemble_stamp = variables.Variable(
+ ensemble_stamp = variables.VariableV1(
initial_value=0,
name="ensemble_stamp",
trainable=False,
diff --git a/tensorflow/contrib/cmake/CMakeLists.txt b/tensorflow/contrib/cmake/CMakeLists.txt
index c6d6f04..f675c13 100644
--- a/tensorflow/contrib/cmake/CMakeLists.txt
+++ b/tensorflow/contrib/cmake/CMakeLists.txt
@@ -30,7 +30,6 @@
option(tensorflow_ENABLE_GRPC_SUPPORT "Enable gRPC support" ON)
option(tensorflow_ENABLE_HDFS_SUPPORT "Enable HDFS support" OFF)
-option(tensorflow_ENABLE_JEMALLOC_SUPPORT "Enable jemalloc support" OFF)
option(tensorflow_BUILD_CC_EXAMPLE "Build the C++ tutorial example" ON)
option(tensorflow_BUILD_PYTHON_BINDINGS "Build the Python bindings" ON)
option(tensorflow_BUILD_ALL_KERNELS "Build all OpKernels" ON)
@@ -218,10 +217,6 @@
endif()
endif()
-if (tensorflow_ENABLE_JEMALLOC_SUPPORT)
- add_definitions(-DTENSORFLOW_USE_JEMALLOC -DJEMALLOC_EXPORT=)
-endif()
-
# External dependencies
include(zlib)
include(gif)
@@ -329,12 +324,6 @@
list(APPEND tensorflow_EXTERNAL_DEPENDENCIES boringssl)
endif()
endif()
-if(tensorflow_ENABLE_JEMALLOC_SUPPORT)
- include(jemalloc)
- list(APPEND tensorflow_EXTERNAL_LIBRARIES ${jemalloc_STATIC_LIBRARIES})
- list(APPEND tensorflow_EXTERNAL_DEPENDENCIES jemalloc)
- include_directories(${jemalloc_INCLUDE_DIRS})
-endif()
if(tensorflow_ENABLE_SNAPPY_SUPPORT)
include(snappy)
list(APPEND tensorflow_EXTERNAL_LIBRARIES ${snappy_STATIC_LIBRARIES})
diff --git a/tensorflow/contrib/cmake/external/jemalloc.cmake b/tensorflow/contrib/cmake/external/jemalloc.cmake
deleted file mode 100644
index afadcc0..0000000
--- a/tensorflow/contrib/cmake/external/jemalloc.cmake
+++ /dev/null
@@ -1,50 +0,0 @@
-# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-include (ExternalProject)
-
-set(jemalloc_INCLUDE_DIRS ${CMAKE_CURRENT_BINARY_DIR}/jemalloc/src/jemalloc/include)
-set(jemalloc_URL https://mirror.bazel.build/github.com/jemalloc/jemalloc-cmake/archive/jemalloc-cmake.4.3.1.tar.gz)
-set(jemalloc_HASH SHA256=f9be9a05fe906deb5c1c8ca818071a7d2e27d66fd87f5ba9a7bf3750bcedeaf0)
-set(jemalloc_BUILD ${CMAKE_CURRENT_BINARY_DIR}/jemalloc/src/jemalloc)
-
-if (WIN32)
- set(jemalloc_INCLUDE_DIRS
- ${jemalloc_INCLUDE_DIRS}
- ${CMAKE_CURRENT_BINARY_DIR}/jemalloc/src/jemalloc/include/msvc_compat
- )
- if(${CMAKE_GENERATOR} MATCHES "Visual Studio.*")
- set(jemalloc_STATIC_LIBRARIES ${jemalloc_BUILD}/Release/jemalloc.lib)
- else()
- set(jemalloc_STATIC_LIBRARIES ${jemalloc_BUILD}/jemalloc.lib)
- endif()
-else()
- set(jemalloc_STATIC_LIBRARIES ${jemalloc_BUILD}/Release/jemalloc.a)
-endif()
-
-ExternalProject_Add(jemalloc
- PREFIX jemalloc
- URL ${jemalloc_URL}
- URL_HASH ${jemalloc_HASH}
- DOWNLOAD_DIR "${DOWNLOAD_LOCATION}"
- BUILD_IN_SOURCE 1
- BUILD_BYPRODUCTS ${jemalloc_STATIC_LIBRARIES}
- BUILD_COMMAND ${CMAKE_COMMAND} --build . --config Release --target jemalloc
- INSTALL_COMMAND ${CMAKE_COMMAND} -E echo "Skipping install step."
- CMAKE_CACHE_ARGS
- -DCMAKE_BUILD_TYPE:STRING=Release
- -DCMAKE_VERBOSE_MAKEFILE:BOOL=OFF
- -Dwith-jemalloc-prefix:STRING=jemalloc_
- -Dwithout-export:BOOL=ON
-)
diff --git a/tensorflow/contrib/cmake/python_modules.txt b/tensorflow/contrib/cmake/python_modules.txt
index c0763f4..6e726701 100644
--- a/tensorflow/contrib/cmake/python_modules.txt
+++ b/tensorflow/contrib/cmake/python_modules.txt
@@ -132,10 +132,8 @@
tensorflow/contrib/cudnn_rnn/python/layers
tensorflow/contrib/cudnn_rnn/python/ops
tensorflow/contrib/data
-tensorflow/contrib/data/kernels
tensorflow/contrib/data/python
tensorflow/contrib/data/python/kernel_tests
-tensorflow/contrib/data/python/kernel_tests/serialization
tensorflow/contrib/data/python/ops
tensorflow/contrib/decision_trees
tensorflow/contrib/decision_trees/proto
@@ -207,6 +205,8 @@
tensorflow/contrib/integrate/python/ops
tensorflow/contrib/kafka/python
tensorflow/contrib/kafka/python/ops
+tensorflow/contrib/ignite/python
+tensorflow/contrib/ignite/python/ops
tensorflow/contrib/keras
tensorflow/contrib/keras/api
tensorflow/contrib/keras/api/keras
diff --git a/tensorflow/contrib/compiler/xla.py b/tensorflow/contrib/compiler/xla.py
index 1e30525..873b035 100644
--- a/tensorflow/contrib/compiler/xla.py
+++ b/tensorflow/contrib/compiler/xla.py
@@ -293,7 +293,8 @@
saved_use_resource = vscope.use_resource
vscope.set_use_resource(True)
- outputs = computation(*computation_inputs)
+ with _disable_summary_context():
+ outputs = computation(*computation_inputs)
# Restore variable scope after computation.
vscope.set_use_resource(saved_use_resource)
@@ -371,13 +372,13 @@
Yields:
None.
"""
- origional_skip_summary_func = summary_op_util.skip_summary
+ original_skip_summary_func = summary_op_util.skip_summary
summary_op_util.skip_summary = lambda: True
try:
yield
finally:
- summary_op_util.skip_summary = origional_skip_summary_func
+ summary_op_util.skip_summary = original_skip_summary_func
class _CapturedObject(object):
@@ -436,8 +437,7 @@
if mode == model_fn_lib.ModeKeys.TRAIN:
train_step, captured_scaffold_fn = self._make_train_step(
features, labels, params)
- with _disable_summary_context():
- (loss,) = compile(train_step)
+ (loss,) = compile(train_step)
return model_fn_lib.EstimatorSpec(
mode=mode,
loss=loss,
@@ -446,8 +446,7 @@
elif mode == model_fn_lib.ModeKeys.EVAL:
eval_step, captured_eval_metric_fn, captured_scaffold_fn = (
self._make_eval_step(features, labels, params))
- with _disable_summary_context():
- outputs = compile(eval_step)
+ outputs = compile(eval_step)
loss = outputs[0]
# Calculate eval_metric_ops if eval_metric_fn is set and captured.
diff --git a/tensorflow/contrib/constrained_optimization/python/external_regret_optimizer.py b/tensorflow/contrib/constrained_optimization/python/external_regret_optimizer.py
index d1af15f..67f8ac2 100644
--- a/tensorflow/contrib/constrained_optimization/python/external_regret_optimizer.py
+++ b/tensorflow/contrib/constrained_optimization/python/external_regret_optimizer.py
@@ -102,9 +102,9 @@
0.0,
(radius - standard_ops.reduce_sum(multipliers)) / standard_ops.maximum(
1.0, standard_ops.reduce_sum(inactive)))
- multipliers += scale * inactive
+ multipliers = multipliers + (scale * inactive)
new_inactive = standard_ops.cast(multipliers > 0, multipliers.dtype)
- multipliers *= new_inactive
+ multipliers = multipliers * new_inactive
return (iteration, multipliers, new_inactive, inactive)
iteration = standard_ops.constant(0)
diff --git a/tensorflow/contrib/constrained_optimization/python/swap_regret_optimizer.py b/tensorflow/contrib/constrained_optimization/python/swap_regret_optimizer.py
index 2c673d9..a6cb1f6 100644
--- a/tensorflow/contrib/constrained_optimization/python/swap_regret_optimizer.py
+++ b/tensorflow/contrib/constrained_optimization/python/swap_regret_optimizer.py
@@ -175,9 +175,9 @@
scale = (1.0 - standard_ops.reduce_sum(
matrix, axis=0, keepdims=True)) / standard_ops.maximum(
1.0, standard_ops.reduce_sum(inactive, axis=0, keepdims=True))
- matrix += scale * inactive
+ matrix = matrix + (scale * inactive)
new_inactive = standard_ops.cast(matrix > 0, matrix.dtype)
- matrix *= new_inactive
+ matrix = matrix * new_inactive
return (iteration, matrix, new_inactive, inactive)
iteration = standard_ops.constant(0)
@@ -210,8 +210,9 @@
# For numerical reasons, make sure that the largest matrix element is zero
# before exponentiating.
- log_matrix -= standard_ops.reduce_max(log_matrix, axis=0, keepdims=True)
- log_matrix -= standard_ops.log(
+ log_matrix = log_matrix - standard_ops.reduce_max(
+ log_matrix, axis=0, keepdims=True)
+ log_matrix = log_matrix - standard_ops.log(
standard_ops.reduce_sum(
standard_ops.exp(log_matrix), axis=0, keepdims=True))
return log_matrix
diff --git a/tensorflow/contrib/copy_graph/python/util/copy_elements.py b/tensorflow/contrib/copy_graph/python/util/copy_elements.py
index 6c9ab6a..9c5871d 100644
--- a/tensorflow/contrib/copy_graph/python/util/copy_elements.py
+++ b/tensorflow/contrib/copy_graph/python/util/copy_elements.py
@@ -31,7 +31,7 @@
from __future__ import print_function
from copy import deepcopy
-from tensorflow.python.ops.variables import Variable
+from tensorflow.python.ops.variables import VariableV1
from tensorflow.python.client.session import Session
from tensorflow.python.framework import ops
@@ -55,7 +55,7 @@
TypeError: If `org_instance` is not a `Variable`.
"""
- if not isinstance(org_instance, Variable):
+ if not isinstance(org_instance, VariableV1):
raise TypeError(str(org_instance) + ' is not a Variable')
#The name of the new variable
@@ -88,7 +88,7 @@
#Initialize the new variable
with to_graph.as_default():
- new_var = Variable(
+ new_var = VariableV1(
init_value,
trainable,
name=new_name,
diff --git a/tensorflow/contrib/copy_graph/python/util/copy_test.py b/tensorflow/contrib/copy_graph/python/util/copy_test.py
index 05744be..ba97c78 100644
--- a/tensorflow/contrib/copy_graph/python/util/copy_test.py
+++ b/tensorflow/contrib/copy_graph/python/util/copy_test.py
@@ -36,7 +36,7 @@
with graph1.as_default():
#Define a Variable in graph1
- some_var = variables.Variable(2)
+ some_var = variables.VariableV1(2)
#Initialize session
sess1 = session_lib.Session()
#Initialize the Variable
@@ -72,7 +72,7 @@
with graph1.as_default():
#Initialize a basic expression y = ax + b
x = array_ops.placeholder("float")
- a = variables.Variable(3.0)
+ a = variables.VariableV1(3.0)
b = constant_op.constant(4.0)
ax = math_ops.multiply(x, a)
y = math_ops.add(ax, b)
diff --git a/tensorflow/contrib/crf/python/ops/crf.py b/tensorflow/contrib/crf/python/ops/crf.py
index 2a91dcb..43bb431 100644
--- a/tensorflow/contrib/crf/python/ops/crf.py
+++ b/tensorflow/contrib/crf/python/ops/crf.py
@@ -56,7 +56,6 @@
from tensorflow.python.framework import dtypes
from tensorflow.python.layers import utils
from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import gen_array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import rnn
@@ -214,10 +213,11 @@
log_norm)
return log_norm
- max_seq_len = array_ops.shape(inputs)[1]
- return control_flow_ops.cond(pred=math_ops.equal(max_seq_len, 1),
- true_fn=_single_seq_fn,
- false_fn=_multi_seq_fn)
+ return utils.smart_cond(
+ pred=math_ops.equal(inputs.shape[1].value or
+ array_ops.shape(inputs)[1], 1),
+ true_fn=_single_seq_fn,
+ false_fn=_multi_seq_fn)
def crf_log_likelihood(inputs,
diff --git a/tensorflow/contrib/data/BUILD b/tensorflow/contrib/data/BUILD
index 9f71061..38f1c65 100644
--- a/tensorflow/contrib/data/BUILD
+++ b/tensorflow/contrib/data/BUILD
@@ -4,17 +4,6 @@
exports_files(["LICENSE"])
-load(
- "//tensorflow:tensorflow.bzl",
- "tf_custom_op_library",
- "tf_gen_op_libs",
- "if_not_windows",
-)
-load(
- "//tensorflow/core:platform/default/build_config_root.bzl",
- "if_static",
-)
-
py_library(
name = "data",
srcs = ["__init__.py"],
@@ -25,30 +14,3 @@
"//tensorflow/python:util",
],
)
-
-cc_library(
- name = "lib_proto_parsing_for_dataset_ops",
- deps = if_not_windows(["//tensorflow/core:lib_proto_parsing"]),
-)
-
-tf_custom_op_library(
- name = "_dataset_ops.so",
- srcs = [
- "ops/dataset_ops.cc",
- "ops/indexed_dataset_ops.cc",
- ],
- deps = [
- "//tensorflow/contrib/data/kernels:dataset_kernels",
- "//tensorflow/contrib/data/kernels:indexed_dataset",
- ] + if_static(
- extra_deps = [":lib_proto_parsing_for_dataset_ops"],
- otherwise = [],
- ),
-)
-
-tf_gen_op_libs(
- op_lib_names = [
- "dataset_ops",
- "indexed_dataset_ops",
- ],
-)
diff --git a/tensorflow/contrib/data/README.md b/tensorflow/contrib/data/README.md
index 848782e..90be7a6 100644
--- a/tensorflow/contrib/data/README.md
+++ b/tensorflow/contrib/data/README.md
@@ -1,10 +1,12 @@
`tf.contrib.data` API
=====================
-NOTE: The `tf.contrib.data` module has been deprecated. Use `tf.data` instead.
-We are continuing to support existing code using the `tf.contrib.data` APIs in
-the current version of TensorFlow, but will eventually remove support. The
-`tf.data` APIs are subject to backwards compatibility guarantees.
+NOTE: The `tf.contrib.data` module has been deprecated. Use `tf.data` instead,
+or `tf.data.experimental` for the experimental transformations previously hosted
+in this module. We are continuing to support existing code using the
+`tf.contrib.data` APIs in the current version of TensorFlow, but will eventually
+remove support. The non-experimental `tf.data` APIs are subject to backwards
+compatibility guarantees.
Porting your code to `tf.data`
------------------------------
@@ -25,13 +27,13 @@
changes is as follows:
* `dataset.dense_to_sparse_batch(...)` is now
- `dataset.apply(tf.contrib.data.dense_to_sparse_batch(...)`.
+ `dataset.apply(tf.data.experimental.dense_to_sparse_batch(...)`.
* `dataset.enumerate(...)` is now
- `dataset.apply(tf.contrib.data.enumerate_dataset(...))`.
+ `dataset.apply(tf.data.experimental.enumerate_dataset(...))`.
* `dataset.group_by_window(...)` is now
- `dataset.apply(tf.contrib.data.group_by_window(...))`.
+ `dataset.apply(tf.data.experimental.group_by_window(...))`.
* `dataset.ignore_errors()` is now
- `dataset.apply(tf.contrib.data.ignore_errors())`.
+ `dataset.apply(tf.data.experimental.ignore_errors())`.
* `dataset.unbatch()` is now `dataset.apply(tf.contrib.data.unbatch())`.
The `Dataset.make_dataset_resource()` and `Iterator.dispose_op()` methods have
diff --git a/tensorflow/contrib/data/__init__.py b/tensorflow/contrib/data/__init__.py
index 3cb5127..c3d3e98 100644
--- a/tensorflow/contrib/data/__init__.py
+++ b/tensorflow/contrib/data/__init__.py
@@ -96,10 +96,6 @@
from tensorflow.contrib.data.python.ops.interleave_ops import sloppy_interleave
from tensorflow.contrib.data.python.ops.iterator_ops import CheckpointInputPipelineHook
from tensorflow.contrib.data.python.ops.iterator_ops import make_saveable_from_iterator
-
-# Optimization constant that can be used to enable auto-tuning.
-from tensorflow.contrib.data.python.ops.optimization import AUTOTUNE
-
from tensorflow.contrib.data.python.ops.parsing_ops import parse_example_dataset
from tensorflow.contrib.data.python.ops.prefetching_ops import copy_to_device
from tensorflow.contrib.data.python.ops.prefetching_ops import prefetch_to_device
@@ -114,11 +110,12 @@
from tensorflow.contrib.data.python.ops.scan_ops import scan
from tensorflow.contrib.data.python.ops.shuffle_ops import shuffle_and_repeat
from tensorflow.contrib.data.python.ops.sliding import sliding_window_batch
-from tensorflow.contrib.data.python.ops.stats_ops import latency_stats
-from tensorflow.contrib.data.python.ops.stats_ops import set_stats_aggregator
-from tensorflow.contrib.data.python.ops.stats_ops import StatsAggregator
from tensorflow.contrib.data.python.ops.unique import unique
from tensorflow.contrib.data.python.ops.writers import TFRecordWriter
+
+# Optimization constant that can be used to enable auto-tuning.
+from tensorflow.python.data.experimental.ops.optimization import AUTOTUNE
+
from tensorflow.python.data.ops.iterator_ops import get_next_as_optional
from tensorflow.python.data.ops.optional_ops import Optional
# pylint: enable=unused-import
diff --git a/tensorflow/contrib/data/kernels/BUILD b/tensorflow/contrib/data/kernels/BUILD
deleted file mode 100644
index ec6cb37..0000000
--- a/tensorflow/contrib/data/kernels/BUILD
+++ /dev/null
@@ -1,139 +0,0 @@
-# Description:
-# Contains kernels for datasets and iterators.
-package(default_visibility = ["//tensorflow:internal"])
-
-licenses(["notice"]) # Apache 2.0
-
-exports_files(["LICENSE"])
-
-cc_library(
- name = "indexed_dataset_headers",
- hdrs = ["indexed_dataset.h"],
- deps = [
- "//tensorflow/core:framework_headers_lib",
- "//third_party/eigen3",
- "@protobuf_archive//:protobuf_headers",
- ],
-)
-
-cc_library(
- name = "indexed_dataset",
- srcs = [
- "identity_indexed_dataset.cc",
- "indexed_dataset.cc",
- ],
- deps = [
- ":indexed_dataset_headers",
- "//tensorflow/core:framework_headers_lib",
- "//third_party/eigen3",
- "@protobuf_archive//:protobuf_headers",
- ],
- alwayslink = 1,
-)
-
-cc_library(
- name = "prefetching_kernels",
- srcs = ["prefetching_kernels.cc"],
- deps = [
- "//tensorflow/core:core_cpu_headers_lib",
- "//tensorflow/core:framework_headers_lib",
- "//third_party/eigen3",
- "@protobuf_archive//:protobuf_headers",
- ],
- alwayslink = 1,
-)
-
-cc_library(
- name = "directed_interleave_dataset_op",
- srcs = ["directed_interleave_dataset_op.cc"],
- deps = [
- "//tensorflow/core:framework_headers_lib",
- "//third_party/eigen3",
- "@protobuf_archive//:protobuf_headers",
- ],
- alwayslink = 1,
-)
-
-cc_library(
- name = "csv_dataset_op",
- srcs = ["csv_dataset_op.cc"],
- deps = [
- "//tensorflow/core:framework_headers_lib",
- "//third_party/eigen3",
- "@protobuf_archive//:protobuf_headers",
- ],
- alwayslink = 1,
-)
-
-cc_library(
- name = "ignore_errors_dataset_op",
- srcs = ["ignore_errors_dataset_op.cc"],
- deps = [
- "//tensorflow/core:framework_headers_lib",
- "//third_party/eigen3",
- "@protobuf_archive//:protobuf_headers",
- ],
- alwayslink = 1,
-)
-
-cc_library(
- name = "lmdb_dataset_op",
- srcs = ["lmdb_dataset_op.cc"],
- deps = [
- "//tensorflow/core:framework_headers_lib",
- "//third_party/eigen3",
- "@lmdb",
- "@protobuf_archive//:protobuf_headers",
- ],
-)
-
-cc_library(
- name = "threadpool_dataset_op",
- srcs = ["threadpool_dataset_op.cc"],
- deps = [
- "//tensorflow/core:framework_headers_lib",
- "//third_party/eigen3",
- "@protobuf_archive//:protobuf_headers",
- ],
- alwayslink = 1,
-)
-
-cc_library(
- name = "unique_dataset_op",
- srcs = ["unique_dataset_op.cc"],
- deps = [
- "//tensorflow/core:framework_headers_lib",
- "//third_party/eigen3",
- "@protobuf_archive//:protobuf_headers",
- ],
- alwayslink = 1,
-)
-
-cc_library(
- name = "assert_next_dataset_op",
- srcs = ["assert_next_dataset_op.cc"],
- deps = [
- "//tensorflow/core:framework_headers_lib",
- "//third_party/eigen3",
- "@protobuf_archive//:protobuf_headers",
- ],
- alwayslink = 1,
-)
-
-cc_library(
- name = "dataset_kernels",
- deps = [
- ":assert_next_dataset_op",
- ":csv_dataset_op",
- ":directed_interleave_dataset_op",
- ":ignore_errors_dataset_op",
- ":indexed_dataset",
- ":lmdb_dataset_op",
- ":prefetching_kernels",
- ":threadpool_dataset_op",
- ":unique_dataset_op",
- "//tensorflow/core:framework_headers_lib",
- "//third_party/eigen3",
- "@protobuf_archive//:protobuf_headers",
- ],
-)
diff --git a/tensorflow/contrib/data/ops/indexed_dataset_ops.cc b/tensorflow/contrib/data/ops/indexed_dataset_ops.cc
deleted file mode 100644
index cd9b7c6..0000000
--- a/tensorflow/contrib/data/ops/indexed_dataset_ops.cc
+++ /dev/null
@@ -1,80 +0,0 @@
-/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
- http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-==============================================================================*/
-#include "tensorflow/core/framework/common_shape_fns.h"
-#include "tensorflow/core/framework/op.h"
-
-namespace tensorflow {
-
-REGISTER_OP("IdentityIndexedDataset")
- .Input("size: uint64")
- .Output("handle: variant")
- .SetIsStateful()
- .SetShapeFn(
- shape_inference::ScalarShape); // TODO(saeta): check input shapes.
-
-///////////////////////////////////////////////////////////////////////////////
-// IndexedDataset Internals
-///////////////////////////////////////////////////////////////////////////////
-
-// Creates the handle.
-REGISTER_OP("MaterializedIndexDatasetHandle")
- .Output("handle: resource")
- .Attr("container: string")
- .Attr("shared_name: string")
- .Attr("output_types: list(type) >= 1")
- .Attr("output_shapes: list(shape) >= 1")
- .SetShapeFn(shape_inference::ScalarShape);
-
-// Actually materialize the materialize handle.
-REGISTER_OP("IndexedDatasetMaterialize")
- .Input("dataset: variant")
- .Input("materialized: resource")
- .SetShapeFn(shape_inference::NoOutputs);
-
-namespace {
-
-Status GetShapeFn(shape_inference::InferenceContext* c) {
- shape_inference::ShapeHandle unused;
- TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &unused));
- std::vector<PartialTensorShape> output_shapes;
- TF_RETURN_IF_ERROR(c->GetAttr("output_shapes", &output_shapes));
- if (output_shapes.size() != c->num_outputs()) {
- return errors::InvalidArgument(
- "`output_shapes` must be the same length as `output_types` (",
- output_shapes.size(), " vs. ", c->num_outputs());
- }
- for (size_t i = 0; i < output_shapes.size(); ++i) {
- shape_inference::ShapeHandle output_shape_handle;
- TF_RETURN_IF_ERROR(c->MakeShapeFromPartialTensorShape(
- output_shapes[i], &output_shape_handle));
- c->set_output(static_cast<int>(i), output_shape_handle);
- }
- return Status::OK();
-}
-
-} // namespace
-
-REGISTER_OP("IndexedDatasetGet")
- .Input("materialized: resource")
- .Input("index: uint64")
- .Output("components: output_types")
- .Attr("output_types: list(type) >= 1")
- .Attr("output_shapes: list(shape) >= 1")
- .SetShapeFn(GetShapeFn)
- .Doc(R"doc(
-Gets the element at `index` from `materialized` IndexedDataset.
-)doc");
-
-} // namespace tensorflow
diff --git a/tensorflow/contrib/data/python/kernel_tests/BUILD b/tensorflow/contrib/data/python/kernel_tests/BUILD
index ce52c99..42f538b 100644
--- a/tensorflow/contrib/data/python/kernel_tests/BUILD
+++ b/tensorflow/contrib/data/python/kernel_tests/BUILD
@@ -8,189 +8,22 @@
load("//tensorflow:tensorflow.bzl", "py_test")
py_test(
- name = "batch_dataset_op_test",
- size = "medium",
- srcs = ["batch_dataset_op_test.py"],
+ name = "assert_element_shape_test",
+ srcs = ["assert_element_shape_test.py"],
srcs_version = "PY2AND3",
- tags = [
- "no_oss", # (b/79552534)
- "no_pip",
- ],
- deps = [
- "//tensorflow/contrib/data/python/ops:batching",
- "//tensorflow/python:array_ops",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:constant_op",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:errors",
- "//tensorflow/python:framework_ops",
- "//tensorflow/python:math_ops",
- "//tensorflow/python:script_ops",
- "//tensorflow/python:session",
- "//tensorflow/python:sparse_tensor",
- "//tensorflow/python:string_ops",
- "//tensorflow/python:tensor_shape",
- "//tensorflow/python:util",
- "//tensorflow/python/data/ops:dataset_ops",
- "//third_party/py/numpy",
- "@absl_py//absl/testing:parameterized",
- ],
-)
-
-py_test(
- name = "bucketing_test",
- size = "medium",
- srcs = ["bucketing_test.py"],
- srcs_version = "PY2AND3",
- deps = [
- "//tensorflow/contrib/data/python/ops:grouping",
- "//tensorflow/python:array_ops",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:constant_op",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:errors",
- "//tensorflow/python:framework_ops",
- "//tensorflow/python:math_ops",
- "//tensorflow/python:sparse_tensor",
- "//tensorflow/python:string_ops",
- "//tensorflow/python:tensor_shape",
- "//tensorflow/python/data/ops:dataset_ops",
- "//third_party/py/numpy",
- ],
-)
-
-py_test(
- name = "csv_dataset_op_test",
- size = "medium",
- srcs = ["csv_dataset_op_test.py"],
- srcs_version = "PY2AND3",
- tags = ["no_pip"],
- deps = [
- "//tensorflow/contrib/data/python/ops:error_ops",
- "//tensorflow/contrib/data/python/ops:readers",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:constant_op",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:errors",
- "//tensorflow/python:framework_test_lib",
- "//tensorflow/python:parsing_ops",
- "//tensorflow/python:platform",
- "//tensorflow/python:platform_test",
- "//tensorflow/python:session",
- "//tensorflow/python/data/ops:readers",
- "//tensorflow/python/eager:context",
- "//third_party/py/numpy",
- ],
-)
-
-py_test(
- name = "dataset_constructor_op_test",
- size = "medium",
- srcs = ["dataset_constructor_op_test.py"],
- srcs_version = "PY2AND3",
- tags = [
- "manual",
- "nomac", # b/62040583
- ],
deps = [
"//tensorflow/contrib/data/python/ops:batching",
"//tensorflow/python:array_ops",
"//tensorflow/python:client_testlib",
"//tensorflow/python:dtypes",
- "//tensorflow/python/data/ops:dataset_ops",
- "//tensorflow/python/data/util:nest",
- ],
-)
-
-py_test(
- name = "directed_interleave_dataset_test",
- size = "medium",
- srcs = ["directed_interleave_dataset_test.py"],
- srcs_version = "PY2AND3",
- deps = [
- "//tensorflow/contrib/data/python/ops:interleave_ops",
- "//tensorflow/python:client_testlib",
"//tensorflow/python:errors",
- "//tensorflow/python:random_seed",
- "//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"],
- deps = [
- "//tensorflow/contrib/data/python/ops:get_single_element",
- "//tensorflow/contrib/data/python/ops:grouping",
- "//tensorflow/python:array_ops",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:constant_op",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:errors",
- "//tensorflow/python:sparse_tensor",
- "//tensorflow/python/data/ops:dataset_ops",
- "@absl_py//absl/testing:parameterized",
- ],
-)
-
-py_test(
- name = "indexed_dataset_ops_test",
- srcs = ["indexed_dataset_ops_test.py"],
- deps = [
- "//tensorflow/contrib/data/python/ops:contrib_op_loader",
- "//tensorflow/contrib/data/python/ops:gen_dataset_ops",
- "//tensorflow/contrib/data/python/ops:indexed_dataset_ops",
- "//tensorflow/python:array_ops",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:dtypes",
- "//tensorflow/python/data/ops:dataset_ops",
- "//third_party/py/numpy",
- ],
-)
-
-py_test(
- name = "interleave_dataset_op_test",
- size = "medium",
- srcs = ["interleave_dataset_op_test.py"],
- srcs_version = "PY2AND3",
- tags = [
- "no_oss",
- "no_pip",
- "notap",
- ],
- deps = [
- "//tensorflow/contrib/data/python/ops:interleave_ops",
- "//tensorflow/python:array_ops",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:errors",
- "//tensorflow/python:math_ops",
- "//tensorflow/python:script_ops",
- "//tensorflow/python:sparse_ops",
- "//tensorflow/python:sparse_tensor",
- "//tensorflow/python/data/ops:dataset_ops",
- "@six_archive//:six",
- ],
-)
-
-py_test(
- name = "iterator_ops_test",
- size = "small",
- srcs = ["iterator_ops_test.py"],
- srcs_version = "PY2AND3",
- tags = ["no_pip"],
- deps = [
- "//tensorflow/contrib/data/python/ops:iterator_ops",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:constant_op",
- "//tensorflow/python:dtypes",
"//tensorflow/python:framework_ops",
- "//tensorflow/python:training",
- "//tensorflow/python:variables",
+ "//tensorflow/python:script_ops",
+ "//tensorflow/python:string_ops",
+ "//tensorflow/python:tensor_shape",
+ "//tensorflow/python/data/kernel_tests:test_base",
"//tensorflow/python/data/ops:dataset_ops",
- "//tensorflow/python/estimator:estimator_py",
+ "//third_party/py/numpy",
],
)
@@ -215,247 +48,24 @@
"//tensorflow/python:platform",
"//tensorflow/python:platform_test",
"//tensorflow/python:session",
+ "//tensorflow/python/data/kernel_tests:test_base",
"//third_party/py/numpy",
],
)
py_test(
- name = "map_dataset_op_test",
- size = "medium",
- srcs = ["map_dataset_op_test.py"],
- srcs_version = "PY2AND3",
- tags = [
- "no_pip",
- "noasan", # times out
- "optonly",
- ],
+ name = "reduce_dataset_test",
+ size = "small",
+ srcs = ["reduce_dataset_test.py"],
deps = [
- "//tensorflow/contrib/data/python/ops:batching",
- "//tensorflow/contrib/data/python/ops:error_ops",
- "//tensorflow/contrib/data/python/ops:optimization",
+ "//tensorflow/contrib/data/python/ops:get_single_element",
+ "//tensorflow/contrib/data/python/ops:grouping",
"//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:dtypes",
+ "//tensorflow/python/data/kernel_tests:test_base",
"//tensorflow/python/data/ops:dataset_ops",
- "//third_party/py/numpy",
- ],
-)
-
-py_test(
- name = "filter_dataset_op_test",
- size = "medium",
- srcs = ["filter_dataset_op_test.py"],
- srcs_version = "PY2AND3",
- deps = [
- "//tensorflow/contrib/data/python/ops:optimization",
- "//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/ops:dataset_ops",
- "//third_party/py/numpy",
- ],
-)
-
-py_test(
- name = "map_defun_op_test",
- size = "small",
- srcs = ["map_defun_op_test.py"],
- srcs_version = "PY2AND3",
- tags = ["no_pip"],
- deps = [
- "//tensorflow/contrib/data/python/ops:map_defun",
- "//tensorflow/python:array_ops",
- "//tensorflow/python:check_ops",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:constant_op",
- "//tensorflow/python:data_flow_ops",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:framework_ops",
- "//tensorflow/python:function",
- "//tensorflow/python:functional_ops",
- "//tensorflow/python:math_ops",
- "//tensorflow/python:session",
- ],
-)
-
-py_test(
- name = "parsing_ops_test",
- size = "small",
- srcs = ["parsing_ops_test.py"],
- srcs_version = "PY2AND3",
- deps = [
- "//tensorflow/contrib/data/python/ops:parsing_ops",
- "//tensorflow/core:protos_all_py",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:errors",
- "//tensorflow/python:framework_ops",
- "//tensorflow/python:parsing_ops",
- "//tensorflow/python:platform",
- "//tensorflow/python:sparse_tensor",
- "//tensorflow/python/data/ops:dataset_ops",
- "//tensorflow/python/data/util:nest",
- "//third_party/py/numpy",
- ],
-)
-
-cuda_py_test(
- name = "prefetching_ops_test",
- size = "small",
- srcs = ["prefetching_ops_test.py"],
- additional_deps = [
- "//tensorflow/contrib/data/python/ops:prefetching_ops",
- "//tensorflow/core:protos_all_py",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:constant_op",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:framework_ops",
- "//tensorflow/python:framework_test_lib",
- "//tensorflow/python:function",
- "//tensorflow/python:resource_variable_ops",
- "//tensorflow/python/compat:compat",
- "//tensorflow/python/data/ops:dataset_ops",
- "//tensorflow/python/data/ops:iterator_ops",
- ],
- tags = ["no_windows_gpu"],
-)
-
-py_test(
- name = "range_dataset_op_test",
- size = "small",
- srcs = ["range_dataset_op_test.py"],
- srcs_version = "PY2AND3",
- deps = [
- "//tensorflow/contrib/data/python/ops:counter",
- "//tensorflow/contrib/data/python/ops:enumerate_ops",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:constant_op",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:errors",
- "//tensorflow/python:tensor_shape",
- "//tensorflow/python/data/ops:dataset_ops",
- ],
-)
-
-py_library(
- name = "reader_dataset_ops_test_base",
- testonly = 1,
- srcs = [
- "reader_dataset_ops_test_base.py",
- ],
- srcs_version = "PY2AND3",
- visibility = [
- "//tensorflow/contrib/data/python/kernel_tests:__pkg__",
- "//tensorflow/contrib/data/python/kernel_tests/serialization:__pkg__",
- ],
- deps = [
- "//tensorflow/contrib/data/python/ops:readers",
- "//tensorflow/core:protos_all_py",
- "//tensorflow/python:array_ops",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:constant_op",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:lib",
- "//tensorflow/python:parsing_ops",
- "//tensorflow/python:util",
- "//tensorflow/python/data/ops:iterator_ops",
- "//tensorflow/python/data/ops:readers",
- ],
-)
-
-py_test(
- name = "reader_dataset_ops_test",
- size = "medium",
- srcs = ["reader_dataset_ops_test.py"],
- srcs_version = "PY2AND3",
- tags = ["no_pip"],
- deps = [
- ":reader_dataset_ops_test_base",
- "//tensorflow/contrib/data/python/ops:readers",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:constant_op",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:errors",
- "//tensorflow/python:framework_ops",
- "//tensorflow/python:parsing_ops",
- "//tensorflow/python:string_ops",
- "//tensorflow/python/data/ops:readers",
- "//tensorflow/python/data/util:nest",
- "//third_party/py/numpy",
- ],
-)
-
-py_test(
- name = "resample_test",
- size = "medium",
- srcs = ["resample_test.py"],
- shard_count = 2,
- srcs_version = "PY2AND3",
- tags = [
- "noasan",
- "optonly",
- ],
- deps = [
- "//tensorflow/contrib/data/python/ops:resampling",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:errors",
- "//tensorflow/python:math_ops",
- "//tensorflow/python:random_ops",
- "//tensorflow/python:string_ops",
- "//tensorflow/python:util",
- "//tensorflow/python/data/ops:dataset_ops",
- "//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
- "@six_archive//:six",
- ],
-)
-
-py_test(
- name = "scan_dataset_op_test",
- size = "small",
- srcs = ["scan_dataset_op_test.py"],
- srcs_version = "PY2AND3",
- tags = ["no_pip"],
- deps = [
- "//tensorflow/contrib/data/python/ops:scan_ops",
- "//tensorflow/python:array_ops",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:constant_op",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:errors",
- "//tensorflow/python:framework_test_lib",
- "//tensorflow/python:sparse_tensor",
- "//tensorflow/python/data/ops:dataset_ops",
- "//tensorflow/python/eager:context",
- "//third_party/py/numpy",
- ],
-)
-
-py_test(
- name = "shuffle_dataset_op_test",
- size = "medium",
- srcs = ["shuffle_dataset_op_test.py"],
- srcs_version = "PY2AND3",
- tags = [
- "no_pip",
- "optonly",
- ],
- deps = [
- "//tensorflow/contrib/data/python/ops:shuffle_ops",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:errors",
- "//tensorflow/python:framework_ops",
- "//tensorflow/python/data/ops:dataset_ops",
- "//third_party/py/numpy",
],
)
@@ -471,151 +81,9 @@
"//tensorflow/python:errors",
"//tensorflow/python:math_ops",
"//tensorflow/python:sparse_tensor",
+ "//tensorflow/python/data/kernel_tests:test_base",
"//tensorflow/python/data/ops:dataset_ops",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
-
-py_library(
- name = "sql_dataset_op_test_base",
- srcs = ["sql_dataset_op_test_base.py"],
- srcs_version = "PY2AND3",
- visibility = [
- "//tensorflow/contrib/data/python/kernel_tests:__pkg__",
- "//tensorflow/contrib/data/python/kernel_tests/serialization:__pkg__",
- ],
- deps = [
- "//tensorflow/contrib/data/python/ops:readers",
- "//tensorflow/python:array_ops",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:dtypes",
- "@org_sqlite//:python",
- ],
-)
-
-py_test(
- name = "sql_dataset_op_test",
- size = "small",
- srcs = ["sql_dataset_op_test.py"],
- srcs_version = "PY2AND3",
- tags = ["no_pip"],
- deps = [
- ":sql_dataset_op_test_base",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:errors",
- ],
-)
-
-py_test(
- name = "stats_dataset_ops_test",
- size = "medium",
- srcs = ["stats_dataset_ops_test.py"],
- srcs_version = "PY2AND3",
- tags = ["no_pip"],
- deps = [
- ":reader_dataset_ops_test_base",
- ":stats_dataset_test_base",
- "//tensorflow/contrib/data/python/ops:stats_ops",
- "//tensorflow/python:array_ops",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:errors",
- "//tensorflow/python:framework_ops",
- "//tensorflow/python/data/ops:dataset_ops",
- "//third_party/py/numpy",
- ],
-)
-
-py_library(
- name = "stats_dataset_test_base",
- srcs = ["stats_dataset_test_base.py"],
- srcs_version = "PY2AND3",
- deps = [
- "//tensorflow/core:protos_all_py",
- "//tensorflow/python:client_testlib",
- ],
-)
-
-py_test(
- name = "threadpool_dataset_ops_test",
- size = "small",
- srcs = ["threadpool_dataset_ops_test.py"],
- srcs_version = "PY2AND3",
- tags = ["no_pip"],
- deps = [
- "//tensorflow/contrib/data/python/ops:threadpool",
- "//tensorflow/contrib/data/python/ops:unique",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:errors",
- "//tensorflow/python:script_ops",
- "//tensorflow/python/data/ops:dataset_ops",
- "//third_party/py/numpy",
- "@absl_py//absl/testing:parameterized",
- ],
-)
-
-py_test(
- name = "unique_dataset_op_test",
- size = "small",
- srcs = ["unique_dataset_op_test.py"],
- srcs_version = "PY2AND3",
- tags = ["no_pip"],
- deps = [
- "//tensorflow/contrib/data/python/ops:unique",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:errors",
- "//tensorflow/python:util",
- "//tensorflow/python/data/ops:dataset_ops",
- ],
-)
-
-py_test(
- name = "window_dataset_op_test",
- size = "medium",
- srcs = ["window_dataset_op_test.py"],
- srcs_version = "PY2AND3",
- tags = [
- "no_pip",
- ],
- deps = [
- "//tensorflow/contrib/data/python/ops:batching",
- "//tensorflow/contrib/data/python/ops:grouping",
- "//tensorflow/python:array_ops",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:math_ops",
- "//tensorflow/python:sparse_tensor",
- "//tensorflow/python/data/ops:dataset_ops",
- "//third_party/py/numpy",
- "@absl_py//absl/testing:parameterized",
- ],
-)
-
-py_test(
- name = "writer_ops_test",
- size = "small",
- srcs = ["writer_ops_test.py"],
- deps = [
- "//tensorflow/contrib/data/python/ops:writers",
- "//tensorflow/python:array_ops",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:lib",
- "//tensorflow/python:util",
- "//tensorflow/python/data/ops:dataset_ops",
- "//tensorflow/python/data/ops:readers",
- ],
-)
-
-py_library(
- name = "test_utils",
- srcs = ["test_utils.py"],
- deps = [
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:errors",
- "//tensorflow/python/data/util:nest",
- ],
-)
diff --git a/tensorflow/contrib/data/python/kernel_tests/assert_element_shape_test.py b/tensorflow/contrib/data/python/kernel_tests/assert_element_shape_test.py
new file mode 100644
index 0000000..0456463
--- /dev/null
+++ b/tensorflow/contrib/data/python/kernel_tests/assert_element_shape_test.py
@@ -0,0 +1,226 @@
+# 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 the experimental input pipeline ops."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from tensorflow.contrib.data.python.ops import batching
+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 tensor_shape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import script_ops
+from tensorflow.python.platform import test
+
+
+class AssertElementShapeTest(test_base.DatasetTestBase):
+
+ def test_assert_element_shape(self):
+
+ def create_dataset(_):
+ return (array_ops.ones(2, dtype=dtypes.float32),
+ array_ops.zeros((3, 4), dtype=dtypes.int32))
+
+ dataset = dataset_ops.Dataset.range(5).map(create_dataset)
+ expected_shapes = (tensor_shape.TensorShape(2),
+ tensor_shape.TensorShape((3, 4)))
+ self.assertEqual(expected_shapes, dataset.output_shapes)
+
+ result = dataset.apply(batching.assert_element_shape(expected_shapes))
+ self.assertEqual(expected_shapes, result.output_shapes)
+
+ iterator = result.make_initializable_iterator()
+ init_op = iterator.initializer
+ get_next = iterator.get_next()
+ with self.cached_session() as sess:
+ sess.run(init_op)
+ for _ in range(5):
+ sess.run(get_next)
+ with self.assertRaises(errors.OutOfRangeError):
+ sess.run(get_next)
+
+ def test_assert_wrong_element_shape(self):
+
+ def create_dataset(_):
+ return (array_ops.ones(2, dtype=dtypes.float32),
+ array_ops.zeros((3, 4), dtype=dtypes.int32))
+
+ dataset = dataset_ops.Dataset.range(3).map(create_dataset)
+ wrong_shapes = (tensor_shape.TensorShape(2),
+ tensor_shape.TensorShape((3, 10)))
+ with self.assertRaises(ValueError):
+ dataset.apply(batching.assert_element_shape(wrong_shapes))
+
+ def test_assert_element_shape_on_unknown_shape_dataset(self):
+
+ def create_unknown_shape_dataset(x):
+ return script_ops.py_func(
+ lambda _: ( # pylint: disable=g-long-lambda
+ np.ones(2, dtype=np.float32),
+ np.zeros((3, 4), dtype=np.int32)),
+ [x],
+ [dtypes.float32, dtypes.int32])
+
+ dataset = dataset_ops.Dataset.range(5).map(create_unknown_shape_dataset)
+ unknown_shapes = (tensor_shape.TensorShape(None),
+ tensor_shape.TensorShape(None))
+ self.assertEqual(unknown_shapes, dataset.output_shapes)
+
+ expected_shapes = (tensor_shape.TensorShape(2),
+ tensor_shape.TensorShape((3, 4)))
+ result = dataset.apply(batching.assert_element_shape(expected_shapes))
+ self.assertEqual(expected_shapes, result.output_shapes)
+
+ iterator = result.make_initializable_iterator()
+ init_op = iterator.initializer
+ get_next = iterator.get_next()
+ with self.cached_session() as sess:
+ sess.run(init_op)
+ for _ in range(5):
+ sess.run(get_next)
+ with self.assertRaises(errors.OutOfRangeError):
+ sess.run(get_next)
+
+ def test_assert_wrong_element_shape_on_unknown_shape_dataset(self):
+
+ def create_unknown_shape_dataset(x):
+ return script_ops.py_func(
+ lambda _: ( # pylint: disable=g-long-lambda
+ np.ones(2, dtype=np.float32),
+ np.zeros((3, 4), dtype=np.int32)),
+ [x],
+ [dtypes.float32, dtypes.int32])
+
+ dataset = dataset_ops.Dataset.range(3).map(create_unknown_shape_dataset)
+ unknown_shapes = (tensor_shape.TensorShape(None),
+ tensor_shape.TensorShape(None))
+ self.assertEqual(unknown_shapes, dataset.output_shapes)
+
+ wrong_shapes = (tensor_shape.TensorShape(2),
+ tensor_shape.TensorShape((3, 10)))
+ iterator = (
+ dataset.apply(batching.assert_element_shape(wrong_shapes))
+ .make_initializable_iterator())
+ init_op = iterator.initializer
+ get_next = iterator.get_next()
+ with self.cached_session() as sess:
+ sess.run(init_op)
+ with self.assertRaises(errors.InvalidArgumentError):
+ sess.run(get_next)
+
+ def test_assert_partial_element_shape(self):
+
+ def create_dataset(_):
+ return (array_ops.ones(2, dtype=dtypes.float32),
+ array_ops.zeros((3, 4), dtype=dtypes.int32))
+
+ dataset = dataset_ops.Dataset.range(5).map(create_dataset)
+ partial_expected_shape = (
+ tensor_shape.TensorShape(None), # Unknown shape
+ tensor_shape.TensorShape((None, 4))) # Partial shape
+ result = dataset.apply(
+ batching.assert_element_shape(partial_expected_shape))
+ # Partial shapes are merged with actual shapes:
+ actual_shapes = (tensor_shape.TensorShape(2),
+ tensor_shape.TensorShape((3, 4)))
+ self.assertEqual(actual_shapes, result.output_shapes)
+
+ iterator = result.make_initializable_iterator()
+ init_op = iterator.initializer
+ get_next = iterator.get_next()
+ with self.cached_session() as sess:
+ sess.run(init_op)
+ for _ in range(5):
+ sess.run(get_next)
+ with self.assertRaises(errors.OutOfRangeError):
+ sess.run(get_next)
+
+ def test_assert_wrong_partial_element_shape(self):
+
+ def create_dataset(_):
+ return (array_ops.ones(2, dtype=dtypes.float32),
+ array_ops.zeros((3, 4), dtype=dtypes.int32))
+
+ dataset = dataset_ops.Dataset.range(3).map(create_dataset)
+ wrong_shapes = (tensor_shape.TensorShape(2),
+ tensor_shape.TensorShape((None, 10)))
+ with self.assertRaises(ValueError):
+ dataset.apply(batching.assert_element_shape(wrong_shapes))
+
+ def test_assert_partial_element_shape_on_unknown_shape_dataset(self):
+
+ def create_unknown_shape_dataset(x):
+ return script_ops.py_func(
+ lambda _: ( # pylint: disable=g-long-lambda
+ np.ones(2, dtype=np.float32),
+ np.zeros((3, 4), dtype=np.int32)),
+ [x],
+ [dtypes.float32, dtypes.int32])
+
+ dataset = dataset_ops.Dataset.range(5).map(create_unknown_shape_dataset)
+ unknown_shapes = (tensor_shape.TensorShape(None),
+ tensor_shape.TensorShape(None))
+ self.assertEqual(unknown_shapes, dataset.output_shapes)
+
+ expected_shapes = (tensor_shape.TensorShape(2),
+ tensor_shape.TensorShape((None, 4)))
+ result = dataset.apply(batching.assert_element_shape(expected_shapes))
+ self.assertEqual(expected_shapes, result.output_shapes)
+
+ iterator = result.make_initializable_iterator()
+ init_op = iterator.initializer
+ get_next = iterator.get_next()
+ with self.cached_session() as sess:
+ sess.run(init_op)
+ for _ in range(5):
+ sess.run(get_next)
+ with self.assertRaises(errors.OutOfRangeError):
+ sess.run(get_next)
+
+ def test_assert_wrong_partial_element_shape_on_unknown_shape_dataset(self):
+
+ def create_unknown_shape_dataset(x):
+ return script_ops.py_func(
+ lambda _: ( # pylint: disable=g-long-lambda
+ np.ones(2, dtype=np.float32),
+ np.zeros((3, 4), dtype=np.int32)),
+ [x],
+ [dtypes.float32, dtypes.int32])
+
+ dataset = dataset_ops.Dataset.range(3).map(create_unknown_shape_dataset)
+ unknown_shapes = (tensor_shape.TensorShape(None),
+ tensor_shape.TensorShape(None))
+ self.assertEqual(unknown_shapes, dataset.output_shapes)
+
+ wrong_shapes = (tensor_shape.TensorShape(2),
+ tensor_shape.TensorShape((None, 10)))
+ iterator = (
+ dataset.apply(batching.assert_element_shape(wrong_shapes))
+ .make_initializable_iterator())
+ init_op = iterator.initializer
+ get_next = iterator.get_next()
+ with self.cached_session() as sess:
+ sess.run(init_op)
+ with self.assertRaises(errors.InvalidArgumentError):
+ sess.run(get_next)
+
+
+if __name__ == "__main__":
+ test.main()
diff --git a/tensorflow/contrib/data/python/kernel_tests/lmdb_dataset_op_test.py b/tensorflow/contrib/data/python/kernel_tests/lmdb_dataset_op_test.py
index 1cc5ddc..d2a7227 100644
--- a/tensorflow/contrib/data/python/kernel_tests/lmdb_dataset_op_test.py
+++ b/tensorflow/contrib/data/python/kernel_tests/lmdb_dataset_op_test.py
@@ -22,6 +22,7 @@
import shutil
from tensorflow.contrib.data.python.ops import readers
+from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
@@ -31,7 +32,7 @@
prefix_path = "tensorflow/core/lib"
-class LMDBDatasetTest(test.TestCase):
+class LMDBDatasetTest(test_base.DatasetTestBase):
def setUp(self):
super(LMDBDatasetTest, self).setUp()
diff --git a/tensorflow/contrib/data/python/kernel_tests/reduce_dataset_test.py b/tensorflow/contrib/data/python/kernel_tests/reduce_dataset_test.py
new file mode 100644
index 0000000..e7281d5
--- /dev/null
+++ b/tensorflow/contrib/data/python/kernel_tests/reduce_dataset_test.py
@@ -0,0 +1,62 @@
+# 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 the experimental input pipeline 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.contrib.data.python.ops import get_single_element
+from tensorflow.contrib.data.python.ops import grouping
+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.ops import array_ops
+from tensorflow.python.platform import test
+
+
+class ReduceDatasetTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+ @parameterized.named_parameters(
+ ("SumZero", 0),
+ ("SumOne", 1),
+ ("SumFive", 5),
+ ("SumTen", 10),
+ )
+ def testReduceDataset(self, stop):
+ def init_fn(_):
+ return np.int64(0)
+
+ def reduce_fn(state, value):
+ return state + value
+
+ def finalize_fn(state):
+ return state
+
+ sum_reducer = grouping.Reducer(init_fn, reduce_fn, finalize_fn)
+
+ stop_t = array_ops.placeholder(dtypes.int64, shape=[])
+ dataset = dataset_ops.Dataset.range(stop_t)
+ element = get_single_element.reduce_dataset(dataset, sum_reducer)
+
+ with self.cached_session() as sess:
+ value = sess.run(element, feed_dict={stop_t: stop})
+ self.assertEqual(stop * (stop - 1) / 2, value)
+
+
+if __name__ == "__main__":
+ test.main()
diff --git a/tensorflow/contrib/data/python/kernel_tests/slide_dataset_op_test.py b/tensorflow/contrib/data/python/kernel_tests/slide_dataset_op_test.py
index 90d18dc..c5a7862 100644
--- a/tensorflow/contrib/data/python/kernel_tests/slide_dataset_op_test.py
+++ b/tensorflow/contrib/data/python/kernel_tests/slide_dataset_op_test.py
@@ -21,6 +21,7 @@
import numpy as np
from tensorflow.contrib.data.python.ops import sliding
+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
@@ -30,7 +31,7 @@
from tensorflow.python.platform import test
-class SlideDatasetTest(test.TestCase, parameterized.TestCase):
+class SlideDatasetTest(test_base.DatasetTestBase, parameterized.TestCase):
@parameterized.named_parameters(
("1", 20, 14, 7, 1),
@@ -197,11 +198,6 @@
sliding.sliding_window_batch(
window_size=1, stride=1, window_shift=1, window_stride=1))
- def assertSparseValuesEqual(self, a, b):
- self.assertAllEqual(a.indices, b.indices)
- self.assertAllEqual(a.values, b.values)
- self.assertAllEqual(a.dense_shape, b.dense_shape)
-
def testSlideSparse(self):
def _sparse(i):
diff --git a/tensorflow/contrib/data/python/kernel_tests/test_utils.py b/tensorflow/contrib/data/python/kernel_tests/test_utils.py
deleted file mode 100644
index 4c3353f..0000000
--- a/tensorflow/contrib/data/python/kernel_tests/test_utils.py
+++ /dev/null
@@ -1,73 +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.
-# ==============================================================================
-"""Test utilities for tf.data functionality."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import re
-
-from tensorflow.python.data.util import nest
-from tensorflow.python.framework import errors
-from tensorflow.python.platform import test
-
-
-class DatasetTestBase(test.TestCase):
- """Base class for dataset tests."""
-
- def _assert_datasets_equal(self, dataset1, dataset2):
- # TODO(rachelim): support sparse tensor outputs
- next1 = dataset1.make_one_shot_iterator().get_next()
- next2 = dataset2.make_one_shot_iterator().get_next()
- with self.cached_session() as sess:
- while True:
- try:
- op1 = sess.run(next1)
- except errors.OutOfRangeError:
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(next2)
- break
- op2 = sess.run(next2)
-
- op1 = nest.flatten(op1)
- op2 = nest.flatten(op2)
- assert len(op1) == len(op2)
- for i in range(len(op1)):
- self.assertAllEqual(op1[i], op2[i])
-
- def _assert_datasets_raise_same_error(self,
- dataset1,
- dataset2,
- exception_class,
- replacements=None):
- # We are defining next1 and next2 in the same line so that we get identical
- # file:line_number in the error messages
- # pylint: disable=line-too-long
- next1, next2 = dataset1.make_one_shot_iterator().get_next(), dataset2.make_one_shot_iterator().get_next()
- # pylint: enable=line-too-long
- with self.cached_session() as sess:
- try:
- sess.run(next1)
- raise ValueError(
- "Expected dataset to raise an error of type %s, but it did not." %
- repr(exception_class))
- except exception_class as e:
- expected_message = e.message
- for old, new, count in replacements:
- expected_message = expected_message.replace(old, new, count)
- # Check that the first segment of the error messages are the same.
- with self.assertRaisesRegexp(exception_class,
- re.escape(expected_message)):
- sess.run(next2)
diff --git a/tensorflow/contrib/data/python/kernel_tests/window_dataset_op_test.py b/tensorflow/contrib/data/python/kernel_tests/window_dataset_op_test.py
deleted file mode 100644
index 8b7b3ac..0000000
--- a/tensorflow/contrib/data/python/kernel_tests/window_dataset_op_test.py
+++ /dev/null
@@ -1,526 +0,0 @@
-# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-"""Tests for the experimental input pipeline 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.contrib.data.python.ops import batching
-from tensorflow.contrib.data.python.ops import grouping
-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.ops import array_ops
-from tensorflow.python.ops import sparse_ops
-from tensorflow.python.platform import test
-
-
-class WindowDatasetTest(test.TestCase, parameterized.TestCase):
-
- def _structuredDataset(self, structure, shape, dtype):
- if structure is None:
- return dataset_ops.Dataset.from_tensors(
- array_ops.zeros(shape, dtype=dtype))
- else:
- return dataset_ops.Dataset.zip(
- tuple([
- self._structuredDataset(substructure, shape, dtype)
- for substructure in structure
- ]))
-
- def _structuredElement(self, structure, shape, dtype):
- if structure is None:
- return array_ops.zeros(shape, dtype=dtype)
- else:
- return tuple([
- self._structuredElement(substructure, shape, dtype)
- for substructure in structure
- ])
-
- def _assertEqual(self, xs, ys):
- self.assertEqual(type(xs), type(ys))
- if isinstance(xs, tuple) and isinstance(ys, tuple):
- self.assertEqual(len(xs), len(ys))
- for x, y in zip(xs, ys):
- self._assertEqual(x, y)
- elif isinstance(xs, np.ndarray) and isinstance(ys, np.ndarray):
- self.assertAllEqual(xs, ys)
- else:
- self.assertEqual(xs, ys)
-
- @parameterized.named_parameters(
- ("1", None, np.int32([]), dtypes.bool),
- ("2", None, np.int32([]), dtypes.int32),
- ("3", None, np.int32([]), dtypes.float32),
- ("4", None, np.int32([]), dtypes.string),
- ("5", None, np.int32([2]), dtypes.int32),
- ("6", None, np.int32([2, 2]), dtypes.int32),
- ("7", (None, None, None), np.int32([]), dtypes.int32),
- ("8", (None, (None, None)), np.int32([]), dtypes.int32),
- )
- def testWindowDatasetFlatMap(self, structure, shape, dtype):
- """Tests windowing by chaining it with flat map.
-
- Args:
- structure: the input structure
- shape: the input shape
- dtype: the input data type
- """
-
- def fn(*args):
- if len(args) == 1 and not isinstance(args[0], tuple):
- return args[0]
- return dataset_ops.Dataset.zip(
- tuple([fn(*arg) if isinstance(arg, tuple) else arg for arg in args]))
-
- dataset = self._structuredDataset(structure, shape, dtype).repeat(5).apply(
- grouping.window_dataset(5)).flat_map(fn)
- get_next = dataset.make_one_shot_iterator().get_next()
- with self.cached_session() as sess:
- expected = sess.run(self._structuredElement(structure, shape, dtype))
- for _ in range(5):
- actual = sess.run(get_next)
- self._assertEqual(expected, actual)
-
- @parameterized.named_parameters(
- ("1", None, np.int32([]), dtypes.bool),
- ("2", None, np.int32([]), dtypes.int32),
- ("3", None, np.int32([]), dtypes.float32),
- ("4", None, np.int32([]), dtypes.string),
- ("5", None, np.int32([2]), dtypes.int32),
- ("6", None, np.int32([2, 2]), dtypes.int32),
- ("7", (None, None, None), np.int32([]), dtypes.int32),
- ("8", (None, (None, None)), np.int32([]), dtypes.int32),
- )
- def testWindowDatasetBatchDense(self, structure, shape, dtype):
- """Tests batching of dense tensor windows.
-
- Args:
- structure: the input structure
- shape: the input shape
- dtype: the input data type
- """
-
- def fn(*args):
- if len(args) == 1 and not isinstance(args[0], tuple):
- return batching.batch_window(args[0])
-
- return tuple([
- fn(*arg) if isinstance(arg, tuple) else batching.batch_window(arg)
- for arg in args
- ])
-
- dataset = self._structuredDataset(structure, shape, dtype).repeat(5).apply(
- grouping.window_dataset(5)).apply(grouping._map_x_dataset(fn))
- get_next = dataset.make_one_shot_iterator().get_next()
- with self.cached_session() as sess:
- expected = sess.run(
- self._structuredElement(structure, np.concatenate(
- ([5], shape), axis=0), dtype))
- actual = sess.run(get_next)
- self._assertEqual(expected, actual)
-
- @parameterized.named_parameters(
- ("1", np.int32([])),
- ("2", np.int32([1])),
- ("3", np.int32([1, 2, 3])),
- )
- def testWindowDatasetBatchDenseDynamicShape(self, shape):
- """Tests batching of dynamically shaped dense tensor windows.
-
- Args:
- shape: the input shape
- """
-
- shape_t = array_ops.placeholder(dtypes.int32)
- dataset = dataset_ops.Dataset.from_tensors(
- array_ops.zeros(shape_t)).repeat(5).apply(
- grouping.window_dataset(5)).apply(
- grouping._map_x_dataset(batching.batch_window))
- iterator = dataset.make_initializable_iterator()
- init_op = iterator.initializer
- get_next = iterator.get_next()
- with self.cached_session() as sess:
- sess.run(init_op, {shape_t: shape})
- expected = sess.run(
- self._structuredElement(None, np.concatenate(([5], shape), axis=0),
- dtypes.int32))
- actual = sess.run(get_next)
- self._assertEqual(expected, actual)
-
- def _make_dense_to_sparse_fn(self, is_scalar):
-
- def dense_to_sparse_scalar(tensor):
- indices = [[]]
- values = array_ops.expand_dims(tensor, 0)
- shape = []
- return sparse_tensor.SparseTensorValue(indices, values, shape)
-
- def dense_to_sparse_non_scalar(tensor):
- indices = array_ops.where(array_ops.ones_like(tensor, dtype=dtypes.bool))
- values = array_ops.gather_nd(tensor, indices)
- shape = array_ops.shape(tensor, out_type=dtypes.int64)
- return sparse_tensor.SparseTensorValue(indices, values, shape)
-
- if is_scalar:
- return dense_to_sparse_scalar
- return dense_to_sparse_non_scalar
-
- def _structuredSparseDataset(self, structure, shape, dtype):
- dense_to_sparse = self._make_dense_to_sparse_fn(len(shape) == 0) # pylint: disable=g-explicit-length-test
- if structure is None:
- return dataset_ops.Dataset.from_tensors(
- dense_to_sparse(array_ops.zeros(shape, dtype=dtype)))
- else:
- return dataset_ops.Dataset.zip(
- tuple([
- self._structuredSparseDataset(substructure, shape, dtype)
- for substructure in structure
- ]))
-
- def _structuredSparseElement(self, structure, shape, dtype):
- dense_to_sparse = self._make_dense_to_sparse_fn(len(shape) == 0) # pylint: disable=g-explicit-length-test
- if structure is None:
- return dense_to_sparse(array_ops.zeros(shape, dtype=dtype))
- else:
- return tuple([
- self._structuredSparseElement(substructure, shape, dtype)
- for substructure in structure
- ])
-
- @parameterized.named_parameters(
- ("1", None, np.int32([]), dtypes.bool),
- ("2", None, np.int32([]), dtypes.int32),
- ("3", None, np.int32([]), dtypes.float32),
- ("4", None, np.int32([]), dtypes.string),
- ("5", None, np.int32([2]), dtypes.int32),
- ("6", None, np.int32([2, 2]), dtypes.int32),
- ("7", (None, None, None), np.int32([]), dtypes.int32),
- ("8", (None, (None, None)), np.int32([]), dtypes.int32),
- )
- def testWindowDatasetBatchSparse(self, structure, shape, dtype):
- """Tests batching of sparse tensor windows.
-
- Args:
- structure: the input structure
- shape: the input shape
- dtype: the input data type
- """
-
- def fn(*args):
- if len(args) == 1 and not isinstance(args[0], tuple):
- return batching.batch_window(args[0])
-
- return tuple([
- fn(*arg) if isinstance(arg, tuple) else batching.batch_window(arg)
- for arg in args
- ])
-
- dataset = self._structuredSparseDataset(
- structure, shape, dtype).repeat(5).apply(
- grouping.window_dataset(5)).apply(grouping._map_x_dataset(fn))
- get_next = dataset.make_one_shot_iterator().get_next()
- with self.cached_session() as sess:
- expected = sess.run(
- self._structuredSparseElement(structure,
- np.concatenate(([5], shape), axis=0),
- dtype))
- actual = sess.run(get_next)
- self._assertEqual(expected, actual)
-
- @parameterized.named_parameters(
- ("1", np.int32([])),
- ("2", np.int32([1])),
- ("3", np.int32([1, 2, 3])),
- )
- def testWindowDatasetBatchSparseDynamicShape(self, shape):
- """Tests batching of dynamically shaped sparse tensor windows.
-
- Args:
- shape: the input shape
- """
-
- shape_t = array_ops.placeholder(dtypes.int32)
- dataset = dataset_ops.Dataset.from_tensors(array_ops.zeros(shape_t)).map(
- self._make_dense_to_sparse_fn(len(shape) == 0)).repeat(5).apply( # pylint: disable=g-explicit-length-test
- grouping.window_dataset(5)).apply(
- grouping._map_x_dataset(batching.batch_window))
- iterator = dataset.make_initializable_iterator()
- init_op = iterator.initializer
- get_next = iterator.get_next()
- with self.cached_session() as sess:
- sess.run(init_op, {shape_t: shape})
- expected = sess.run(
- self._structuredSparseElement(None,
- np.concatenate(([5], shape), axis=0),
- dtypes.int32))
- actual = sess.run(get_next)
- self._assertEqual(expected, actual)
-
- def _structuredRaggedDataset(self, structure, shapes, dtype):
-
- if structure is None:
- return dataset_ops.Dataset.from_tensor_slices(shapes).map(
- lambda shape: array_ops.zeros(shape, dtype=dtype))
- else:
- return dataset_ops.Dataset.zip(
- tuple([
- self._structuredRaggedDataset(substructure, shapes, dtype)
- for substructure in structure
- ]))
-
- @parameterized.named_parameters(
- ("1", None, np.int32([[1], [2], [3]]), dtypes.bool, [-1]),
- ("2", None, np.int32([[1], [2], [3]]), dtypes.int32, [-1]),
- ("3", None, np.int32([[1], [2], [3]]), dtypes.float32, [-1]),
- ("4", None, np.int32([[1], [2], [3]]), dtypes.string, [-1]),
- ("5", None, np.int32([[1, 3], [2, 2], [3, 1]]), dtypes.int32, [-1, -1]),
- ("6", None, np.int32([[3, 1, 3], [1, 3, 1]]), dtypes.int32, [-1, -1, -1]),
- ("7", (None, None, None), np.int32([[1], [2], [3]]), dtypes.int32, [-1]),
- ("8", (None,
- (None, None)), np.int32([[1], [2], [3]]), dtypes.int32, [-1]),
- ("9", None, np.int32([[1], [2], [3]]), dtypes.int32, [-1]),
- ("10", None, np.int32([[1], [2], [3]]), dtypes.int32, np.int32([10])),
- )
- def testWindowDatasetPaddedBatchDense(self, structure, shapes, dtype,
- padded_shape):
- """Tests padded batching of dense tensor windows.
-
- Args:
- structure: the input structure
- shapes: the input shapes
- dtype: the input data type
- padded_shape: the shape to pad the output to
- """
-
- def fn(*args):
- if len(args) == 1 and not isinstance(args[0], tuple):
- return batching.padded_batch_window(args[0], padded_shape)
-
- return tuple([
- fn(*arg) if isinstance(arg, tuple) else batching.padded_batch_window(
- arg, padded_shape) for arg in args
- ])
-
- dataset = self._structuredRaggedDataset(structure, shapes, dtype).apply(
- grouping.window_dataset(len(shapes))).apply(
- grouping._map_x_dataset(fn))
- get_next = dataset.make_one_shot_iterator().get_next()
- with self.cached_session() as sess:
- expected_shape = np.maximum(np.amax(shapes, axis=0), padded_shape)
- expected = sess.run(
- self._structuredElement(
- structure,
- np.concatenate((np.int32([len(shapes)]), expected_shape)), dtype))
- actual = sess.run(get_next)
- self._assertEqual(expected, actual)
-
- @parameterized.named_parameters(
- ("1", np.int32([[1], [2], [3]]), [-1]),
- ("2", np.int32([[1, 3], [2, 2], [3, 1]]), [-1, -1]),
- ("3", np.int32([[3, 1, 3], [1, 3, 1]]), [-1, -1, -1]),
- )
- def testWindowDatasetPaddedBatchDenseDynamicShape(self, shapes, padded_shape):
- """Tests padded batching of dynamically shaped dense tensor windows.
-
- Args:
- shapes: the input shapes
- padded_shape: the shape to pad the output to
- """
-
- shapes_t = array_ops.placeholder(dtypes.int32)
- dataset = dataset_ops.Dataset.from_tensor_slices(shapes_t).map(
- lambda shape: array_ops.zeros(shape, dtype=dtypes.int32)).apply(
- grouping.window_dataset(len(shapes))).apply(
- grouping._map_x_dataset(
- lambda x: batching.padded_batch_window(x, padded_shape)))
- iterator = dataset.make_initializable_iterator()
- init_op = iterator.initializer
- get_next = iterator.get_next()
- with self.cached_session() as sess:
- sess.run(init_op, {shapes_t: shapes})
- expected_shape = np.maximum(np.amax(shapes, axis=0), padded_shape)
- expected = sess.run(
- self._structuredElement(
- None, np.concatenate((np.int32([len(shapes)]), expected_shape)),
- dtypes.int32))
- actual = sess.run(get_next)
- self._assertEqual(expected, actual)
-
- @parameterized.named_parameters(
- ("1", np.int32([[1]]), np.int32([0])),
- ("2", np.int32([[10], [20]]), np.int32([15])),
- )
- def testWindowDatasetPaddedBatchDenseInvalid(self, shapes, padded_shape):
- """Tests invalid padded batching of dense tensor windows.
-
- Args:
- shapes: the input shapes
- padded_shape: the shape to pad the output to
- """
-
- dataset = dataset_ops.Dataset.from_tensor_slices(shapes).map(
- lambda shape: array_ops.zeros(shape, dtype=dtypes.int32)).apply(
- grouping.window_dataset(len(shapes))).apply(
- grouping._map_x_dataset(
- lambda x: batching.padded_batch_window(x, padded_shape)))
- get_next = dataset.make_one_shot_iterator().get_next()
- with self.cached_session() as sess:
- with self.assertRaises(errors.InvalidArgumentError):
- sess.run(get_next)
-
- def _structuredRaggedSparseDataset(self, structure, shapes, dtype):
-
- def map_fn(shape):
- dense_to_sparse = self._make_dense_to_sparse_fn(False)
- return dense_to_sparse(array_ops.zeros(shape, dtype=dtype))
-
- if structure is None:
- return dataset_ops.Dataset.from_tensor_slices(shapes).map(map_fn)
- else:
- return dataset_ops.Dataset.zip(
- tuple([
- self._structuredRaggedSparseDataset(substructure, shapes, dtype)
- for substructure in structure
- ]))
-
- def _structuredRaggedSparseElement(self, structure, shapes, dtype,
- padded_shape):
- if structure is None:
- dense_shape = np.maximum(np.amax(shapes, axis=0), padded_shape)
- values = []
- for shape in shapes:
- dense_to_sparse = self._make_dense_to_sparse_fn(len(shape) == 0) # pylint: disable=g-explicit-length-test
- sparse = dense_to_sparse(array_ops.zeros(shape, dtype=dtype))
- padded_sparse = sparse_tensor.SparseTensor(sparse.indices,
- sparse.values, dense_shape)
- reshaped_sparse = sparse_ops.sparse_reshape(
- padded_sparse,
- array_ops.concat([np.array([1], dtype=np.int64), dense_shape], 0))
- values.append(reshaped_sparse)
- return sparse_ops.sparse_concat(0, values)
- else:
- return tuple([
- self._structuredRaggedSparseElement(substructure, shapes, dtype,
- padded_shape)
- for substructure in structure
- ])
-
- @parameterized.named_parameters(
- ("1", None, np.int64([[1], [2], [3]]), dtypes.bool, [-1]),
- ("2", None, np.int64([[1], [2], [3]]), dtypes.int32, [-1]),
- ("3", None, np.int64([[1], [2], [3]]), dtypes.float32, [-1]),
- ("4", None, np.int64([[1], [2], [3]]), dtypes.string, [-1]),
- ("5", None, np.int64([[1, 3], [2, 2], [3, 1]]), dtypes.int32, [-1, -1]),
- ("6", None, np.int64([[1, 3, 1], [3, 1, 3]]), dtypes.int32, [-1, -1, -1]),
- ("7", (None, None, None), np.int64([[1], [2], [3]]), dtypes.int32, [-1]),
- ("8", (None,
- (None, None)), np.int64([[1], [2], [3]]), dtypes.int32, [-1]),
- ("9", None, np.int64([[1], [2], [3]]), dtypes.int32, [-1]),
- ("10", None, np.int64([[1], [2], [3]]), dtypes.int32, np.int64([10])),
- )
- def testWindowDatasetPaddedBatchSparse(self, structure, shapes, dtype,
- padded_shape):
- """Tests padded batching of sparse tensor windows.
-
- Args:
- structure: the input structure
- shapes: the input shapes
- dtype: the input data type
- padded_shape: the shape to pad the output to
- """
-
- def fn(*args):
- if len(args) == 1 and not isinstance(args[0], tuple):
- return batching.padded_batch_window(args[0], padded_shape)
-
- return tuple([
- fn(*arg) if isinstance(arg, tuple) else batching.padded_batch_window(
- arg, padded_shape) for arg in args
- ])
-
- dataset = self._structuredRaggedSparseDataset(
- structure, shapes, dtype).apply(grouping.window_dataset(
- len(shapes))).apply(grouping._map_x_dataset(fn))
- get_next = dataset.make_one_shot_iterator().get_next()
- with self.cached_session() as sess:
- expected = sess.run(
- self._structuredRaggedSparseElement(structure, shapes, dtype,
- padded_shape))
- actual = sess.run(get_next)
- self._assertEqual(expected, actual)
-
- @parameterized.named_parameters(
- ("1", np.int64([[1], [2], [3]]), [-1]),
- ("2", np.int64([[1, 3], [2, 2], [3, 1]]), [-1, -1]),
- ("3", np.int64([[3, 1, 3], [1, 3, 1]]), [-1, -1, -1]),
- )
- def testWindowDatasetPaddedBatchSparseDynamicShape(self, shapes,
- padded_shape):
- """Tests padded batching of dynamically shaped sparse tensor windows.
-
- Args:
- shapes: the input shapes
- padded_shape: the shape to pad the output to
- """
-
- shapes_t = array_ops.placeholder(dtypes.int32)
- dataset = dataset_ops.Dataset.from_tensor_slices(shapes_t).map(
- lambda shape: array_ops.zeros(shape, dtype=dtypes.int32)).map(
- self._make_dense_to_sparse_fn(False)
- ).apply(grouping.window_dataset(len(shapes))).apply(
- grouping._map_x_dataset(
- lambda x: batching.padded_batch_window(x, padded_shape)))
- iterator = dataset.make_initializable_iterator()
- init_op = iterator.initializer
- get_next = iterator.get_next()
- with self.cached_session() as sess:
- sess.run(init_op, {shapes_t: shapes})
- expected = sess.run(
- self._structuredRaggedSparseElement(None, shapes, dtypes.int32,
- padded_shape))
- actual = sess.run(get_next)
- self._assertEqual(expected, actual)
-
- @parameterized.named_parameters(
- ("1", np.int64([[1]]), [0]),
- ("2", np.int64([[10], [20]]), [15]),
- )
- def testWindowDatasetPaddedBatchSparseInvalid(self, shapes, padded_shape):
- """Tests invalid padded batching of sparse tensor windows.
-
- Args:
- shapes: the input shapes
- padded_shape: the shape to pad the output to
- """
-
- dataset = dataset_ops.Dataset.from_tensor_slices(shapes).map(
- lambda shape: array_ops.zeros(shape, dtype=dtypes.int32)).map(
- self._make_dense_to_sparse_fn(False)
- ).apply(grouping.window_dataset(len(shapes))).apply(
- grouping._map_x_dataset(
- lambda x: batching.padded_batch_window(x, padded_shape)))
- get_next = dataset.make_one_shot_iterator().get_next()
- with self.cached_session() as sess:
- with self.assertRaises(errors.InvalidArgumentError):
- sess.run(get_next)
-
-
-if __name__ == "__main__":
- test.main()
diff --git a/tensorflow/contrib/data/python/ops/BUILD b/tensorflow/contrib/data/python/ops/BUILD
index a14781c..34dc237 100644
--- a/tensorflow/contrib/data/python/ops/BUILD
+++ b/tensorflow/contrib/data/python/ops/BUILD
@@ -16,10 +16,7 @@
srcs = ["counter.py"],
srcs_version = "PY2AND3",
deps = [
- ":scan_ops",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:framework_ops",
- "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/experimental/ops:counter",
],
)
@@ -28,12 +25,7 @@
srcs = ["get_single_element.py"],
srcs_version = "PY2AND3",
deps = [
- ":grouping",
- "//tensorflow/python:dataset_ops_gen",
- "//tensorflow/python/data/ops:dataset_ops",
- "//tensorflow/python/data/util:nest",
- "//tensorflow/python/data/util:sparse",
- "//third_party/py/numpy",
+ "//tensorflow/python/data/experimental/ops:get_single_element",
],
)
@@ -44,10 +36,7 @@
],
srcs_version = "PY2AND3",
deps = [
- "//tensorflow/python:dataset_ops_gen",
- "//tensorflow/python:framework_ops",
- "//tensorflow/python:training",
- "//tensorflow/python/data/ops:iterator_ops",
+ "//tensorflow/python/data/experimental/ops:iterator_ops",
],
)
@@ -58,15 +47,7 @@
],
srcs_version = "PY2AND3",
deps = [
- "//tensorflow/python:constant_op",
- "//tensorflow/python:dataset_ops_gen",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:framework_ops",
- "//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/experimental/ops:random_ops",
],
)
@@ -78,19 +59,19 @@
srcs_version = "PY2AND3",
deps = [
":batching",
- ":gen_dataset_ops",
":interleave_ops",
- ":optimization",
":parsing_ops",
":shuffle_ops",
"//tensorflow/python:constant_op",
"//tensorflow/python:dataset_ops_gen",
"//tensorflow/python:dtypes",
+ "//tensorflow/python:experimental_dataset_ops_gen",
"//tensorflow/python:framework_ops",
"//tensorflow/python:lib",
"//tensorflow/python:platform",
"//tensorflow/python:tensor_shape",
"//tensorflow/python:util",
+ "//tensorflow/python/data/experimental/ops:readers",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/ops:readers",
"//tensorflow/python/data/util:convert",
@@ -106,7 +87,7 @@
],
srcs_version = "PY2AND3",
deps = [
- "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/experimental/ops:shuffle_ops",
],
)
@@ -125,6 +106,7 @@
"//tensorflow/python:math_ops",
"//tensorflow/python:tensor_shape",
"//tensorflow/python:tensor_util",
+ "//tensorflow/python/data/experimental/ops:batching",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/util:convert",
"//tensorflow/python/data/util:nest",
@@ -138,8 +120,7 @@
srcs = ["enumerate_ops.py"],
srcs_version = "PY2AND3",
deps = [
- "//tensorflow/python:dtypes",
- "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/experimental/ops:enumerate_ops",
],
)
@@ -148,11 +129,7 @@
srcs = ["error_ops.py"],
srcs_version = "PY2AND3",
deps = [
- ":contrib_op_loader",
- ":gen_dataset_ops",
- "//tensorflow/python/data/ops:dataset_ops",
- "//tensorflow/python/data/util:nest",
- "//tensorflow/python/data/util:sparse",
+ "//tensorflow/python/data/experimental/ops:error_ops",
],
)
@@ -161,16 +138,7 @@
srcs = ["grouping.py"],
srcs_version = "PY2AND3",
deps = [
- "//tensorflow/python:array_ops",
- "//tensorflow/python:check_ops",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:framework_ops",
- "//tensorflow/python:function",
- "//tensorflow/python:math_ops",
- "//tensorflow/python:tensor_shape",
- "//tensorflow/python/data/ops:dataset_ops",
- "//tensorflow/python/data/util:nest",
- "//tensorflow/python/data/util:sparse",
+ "//tensorflow/python/data/experimental/ops:grouping",
],
)
@@ -179,32 +147,7 @@
srcs = ["interleave_ops.py"],
srcs_version = "PY2AND3",
deps = [
- ":contrib_op_loader",
- ":gen_dataset_ops",
- ":random_ops",
- "//tensorflow/contrib/stateless",
- "//tensorflow/python:array_ops",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:framework_ops",
- "//tensorflow/python:math_ops",
- "//tensorflow/python:util",
- "//tensorflow/python/data/ops:readers",
- "//tensorflow/python/data/util:nest",
- "//tensorflow/python/data/util:sparse",
- ],
-)
-
-py_library(
- name = "optimization",
- srcs = ["optimization.py"],
- srcs_version = "PY2AND3",
- deps = [
- ":contrib_op_loader",
- ":gen_dataset_ops",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:framework_ops",
- "//tensorflow/python/data/util:nest",
- "//tensorflow/python/data/util:sparse",
+ "//tensorflow/python/data/experimental/ops:interleave_ops",
],
)
@@ -213,25 +156,7 @@
srcs = ["parsing_ops.py"],
srcs_version = "PY2AND3",
deps = [
- "//tensorflow/python:dataset_ops_gen",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:framework_ops",
- "//tensorflow/python:parsing_ops",
- "//tensorflow/python:sparse_tensor",
- "//tensorflow/python:tensor_shape",
- "//tensorflow/python/data/ops:dataset_ops",
- "//tensorflow/python/data/util:nest",
- ],
-)
-
-py_library(
- name = "map_defun",
- srcs = ["map_defun.py"],
- srcs_version = "PY2AND3",
- deps = [
- "//tensorflow/python:dataset_ops_gen",
- "//tensorflow/python:framework_ops",
- "//tensorflow/python:tensor_shape",
+ "//tensorflow/python/data/experimental/ops:parsing_ops",
],
)
@@ -240,18 +165,7 @@
srcs = ["resampling.py"],
srcs_version = "PY2AND3",
deps = [
- ":batching",
- ":interleave_ops",
- ":scan_ops",
- "//tensorflow/python:array_ops",
- "//tensorflow/python:control_flow_ops",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:framework_ops",
- "//tensorflow/python:logging_ops",
- "//tensorflow/python:math_ops",
- "//tensorflow/python:random_ops",
- "//tensorflow/python/data/ops:dataset_ops",
- "//third_party/py/numpy",
+ "//tensorflow/python/data/experimental/ops:resampling",
],
)
@@ -260,12 +174,7 @@
srcs = ["scan_ops.py"],
srcs_version = "PY2AND3",
deps = [
- "//tensorflow/python:dataset_ops_gen",
- "//tensorflow/python:framework_ops",
- "//tensorflow/python:function",
- "//tensorflow/python/data/ops:dataset_ops",
- "//tensorflow/python/data/util:nest",
- "//tensorflow/python/data/util:sparse",
+ "//tensorflow/python/data/experimental/ops:scan_ops",
],
)
@@ -285,32 +194,11 @@
)
py_library(
- name = "stats_ops",
- srcs = ["stats_ops.py"],
- srcs_version = "PY2AND3",
- deps = [
- "//tensorflow/python:dataset_ops_gen",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:framework_ops",
- "//tensorflow/python/data/ops:dataset_ops",
- "//tensorflow/python/data/ops:iterator_ops",
- "//tensorflow/python/data/util:nest",
- "//tensorflow/python/data/util:sparse",
- ],
-)
-
-py_library(
name = "threadpool",
srcs = ["threadpool.py"],
srcs_version = "PY2AND3",
deps = [
- ":contrib_op_loader",
- ":gen_dataset_ops",
- "//tensorflow/python:resource_variable_ops",
- "//tensorflow/python/data/ops:dataset_ops",
- "//tensorflow/python/data/util:nest",
- "//tensorflow/python/data/util:sparse",
- "//tensorflow/python/eager:context",
+ "//tensorflow/python/data/experimental/ops:threadpool",
],
)
@@ -321,12 +209,7 @@
],
srcs_version = "PY2AND3",
deps = [
- ":contrib_op_loader",
- ":gen_dataset_ops",
- "//tensorflow/python:dtypes",
- "//tensorflow/python/data/ops:dataset_ops",
- "//tensorflow/python/data/util:nest",
- "//tensorflow/python/data/util:sparse",
+ "//tensorflow/python/data/experimental/ops:unique",
],
)
@@ -337,56 +220,7 @@
],
srcs_version = "PY2AND3",
deps = [
- "//tensorflow/python:dtypes",
- "//tensorflow/python/data/ops:dataset_ops",
- ],
-)
-
-tf_gen_op_wrapper_py(
- name = "gen_dataset_ops",
- out = "gen_dataset_ops.py",
- deps = [
- "//tensorflow/contrib/data:dataset_ops_op_lib",
- "//tensorflow/contrib/data:indexed_dataset_ops_op_lib",
- ],
-)
-
-tf_kernel_library(
- name = "dataset_ops_kernels",
- deps = [
- "//tensorflow/contrib/data/kernels:dataset_kernels",
- "//tensorflow/core:framework",
- ],
- alwayslink = 1,
-)
-
-tf_custom_op_py_library(
- name = "contrib_op_loader",
- srcs = ["contrib_op_loader.py"],
- dso = ["//tensorflow/contrib/data:_dataset_ops.so"],
- kernels = [
- ":dataset_ops_kernels",
- "//tensorflow/contrib/data:indexed_dataset_ops_op_lib",
- "//tensorflow/contrib/data:dataset_ops_op_lib",
- ],
- srcs_version = "PY2AND3",
- deps = [
- ":gen_dataset_ops",
- "//tensorflow/contrib/util:util_py",
- "//tensorflow/python:platform",
- ],
-)
-
-py_library(
- name = "indexed_dataset_ops",
- srcs = ["indexed_dataset_ops.py"],
- deps = [
- ":contrib_op_loader",
- ":gen_dataset_ops",
- "//tensorflow/python:framework_ops",
- "//tensorflow/python/data/ops:dataset_ops",
- "//tensorflow/python/data/util:nest",
- "//tensorflow/python/data/util:sparse",
+ "//tensorflow/python/data/experimental/ops:writers",
],
)
@@ -394,11 +228,7 @@
name = "prefetching_ops",
srcs = ["prefetching_ops.py"],
deps = [
- ":contrib_op_loader",
- "//tensorflow/python:framework_ops",
- "//tensorflow/python/data/ops:dataset_ops",
- "//tensorflow/python/data/util:nest",
- "//tensorflow/python/data/util:sparse",
+ "//tensorflow/python/data/experimental/ops:prefetching_ops",
],
)
@@ -411,17 +241,14 @@
":error_ops",
":get_single_element",
":grouping",
- ":indexed_dataset_ops",
":interleave_ops",
- ":map_defun",
- ":optimization",
":prefetching_ops",
+ ":random_ops",
":readers",
":resampling",
":scan_ops",
":shuffle_ops",
":sliding",
- ":stats_ops",
":threadpool",
":unique",
":writers",
diff --git a/tensorflow/contrib/data/python/ops/batching.py b/tensorflow/contrib/data/python/ops/batching.py
index 7a0f221..8c60459 100644
--- a/tensorflow/contrib/data/python/ops/batching.py
+++ b/tensorflow/contrib/data/python/ops/batching.py
@@ -17,134 +17,14 @@
from __future__ import division
from __future__ import print_function
-import numpy as np
-
-from tensorflow.contrib.data.python.ops import get_single_element
-from tensorflow.contrib.data.python.ops import grouping
from tensorflow.contrib.framework import with_shape
-from tensorflow.python.data.ops import dataset_ops
-from tensorflow.python.data.util import convert
+from tensorflow.python.data.experimental.ops import batching
from tensorflow.python.data.util import nest
-from tensorflow.python.data.util import sparse
-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 check_ops
-from tensorflow.python.ops import control_flow_ops
-from tensorflow.python.ops import gen_array_ops
-from tensorflow.python.ops import gen_dataset_ops
-from tensorflow.python.ops import math_ops
-from tensorflow.python.ops import sparse_ops
from tensorflow.python.util import deprecation
-def batch_window(dataset):
- """Batches a window of tensors.
-
- Args:
- dataset: the input dataset.
-
- Returns:
- A `Tensor` representing the batch of the entire input dataset.
- """
- if isinstance(dataset.output_classes, tuple):
- raise TypeError("Input dataset expected to have a single component")
- if dataset.output_classes is ops.Tensor:
- return _batch_dense_window(dataset)
- elif dataset.output_classes is sparse_tensor.SparseTensor:
- return _batch_sparse_window(dataset)
- else:
- raise TypeError("Unsupported dataset type: %s" % dataset.output_classes)
-
-
-def _batch_dense_window(dataset):
- """Batches a window of dense tensors."""
-
- def key_fn(_):
- return np.int64(0)
-
- def shape_init_fn(_):
- return array_ops.shape(first_element)
-
- def shape_reduce_fn(state, value):
- check_ops.assert_equal(state, array_ops.shape(value))
- return state
-
- def finalize_fn(state):
- return state
-
- if dataset.output_shapes.is_fully_defined():
- shape = dataset.output_shapes
- else:
- first_element = get_single_element.get_single_element(dataset.take(1))
- shape_reducer = grouping.Reducer(shape_init_fn, shape_reduce_fn,
- finalize_fn)
- shape = get_single_element.get_single_element(
- dataset.apply(grouping.group_by_reducer(key_fn, shape_reducer)))
-
- def batch_init_fn(_):
- batch_shape = array_ops.concat([[0], shape], 0)
- return gen_array_ops.empty(batch_shape, dtype=dataset.output_types)
-
- def batch_reduce_fn(state, value):
- return array_ops.concat([state, [value]], 0)
-
- batch_reducer = grouping.Reducer(batch_init_fn, batch_reduce_fn, finalize_fn)
- return get_single_element.get_single_element(
- dataset.apply(grouping.group_by_reducer(key_fn, batch_reducer)))
-
-
-def _batch_sparse_window(dataset):
- """Batches a window of sparse tensors."""
-
- def key_fn(_):
- return np.int64(0)
-
- def shape_init_fn(_):
- return first_element.dense_shape
-
- def shape_reduce_fn(state, value):
- check_ops.assert_equal(state, value.dense_shape)
- return state
-
- def finalize_fn(state):
- return state
-
- if dataset.output_shapes.is_fully_defined():
- shape = dataset.output_shapes
- else:
- first_element = get_single_element.get_single_element(dataset.take(1))
- shape_reducer = grouping.Reducer(shape_init_fn, shape_reduce_fn,
- finalize_fn)
- shape = get_single_element.get_single_element(
- dataset.apply(grouping.group_by_reducer(key_fn, shape_reducer)))
-
- def batch_init_fn(_):
- indices_shape = array_ops.concat([[0], [array_ops.size(shape) + 1]], 0)
- return sparse_tensor.SparseTensor(
- indices=gen_array_ops.empty(indices_shape, dtype=dtypes.int64),
- values=constant_op.constant([], shape=[0], dtype=dataset.output_types),
- dense_shape=array_ops.concat(
- [np.array([0], dtype=np.int64),
- math_ops.cast(shape, dtypes.int64)], 0))
-
- def batch_reduce_fn(state, value):
- return sparse_ops.sparse_concat(0, [state, value])
-
- def reshape_fn(value):
- return sparse_ops.sparse_reshape(
- value,
- array_ops.concat([np.array([1], dtype=np.int64), value.dense_shape], 0))
-
- batch_reducer = grouping.Reducer(batch_init_fn, batch_reduce_fn, finalize_fn)
- return get_single_element.get_single_element(
- dataset.map(reshape_fn).apply(
- grouping.group_by_reducer(key_fn, batch_reducer)))
-
-
+@deprecation.deprecated(
+ None, "Use `tf.data.experimental.dense_to_sparse_batch(...)`.")
def dense_to_sparse_batch(batch_size, row_shape):
"""A transformation that batches ragged elements into `tf.SparseTensor`s.
@@ -187,201 +67,10 @@
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
-
- def _apply_fn(dataset):
- return _DenseToSparseBatchDataset(dataset, batch_size, row_shape)
-
- return _apply_fn
+ return batching.dense_to_sparse_batch(batch_size, row_shape)
-def padded_batch_window(dataset, padded_shape, padding_value=None):
- """Batches a window of tensors with padding.
-
- Args:
- dataset: the input dataset.
- padded_shape: (Optional.) `tf.TensorShape` or `tf.int64` vector tensor-like
- object representing the shape to which the input elements should be padded
- prior to batching. Any unknown dimensions (e.g. `tf.Dimension(None)` in a
- `tf.TensorShape` or `-1` in a tensor-like object) will be padded to the
- maximum size of that dimension in each batch.
- padding_value: (Optional.) A scalar-shaped `tf.Tensor`, representing the
- padding value to use. Defaults are `0` for numeric types and the empty
- string for string types. If `dataset` contains `tf.SparseTensor`, this
- value is ignored.
-
- Returns:
- A `Tensor` representing the batch of the entire input dataset.
-
- Raises:
- ValueError: if invalid arguments are provided.
- """
- if not issubclass(dataset.output_classes,
- (ops.Tensor, sparse_tensor.SparseTensor)):
- raise TypeError("Input dataset expected to have a single tensor component")
- if issubclass(dataset.output_classes, (ops.Tensor)):
- return _padded_batch_dense_window(dataset, padded_shape, padding_value)
- elif issubclass(dataset.output_classes, (sparse_tensor.SparseTensor)):
- if padding_value is not None:
- raise ValueError("Padding value not allowed for sparse tensors")
- return _padded_batch_sparse_window(dataset, padded_shape)
- else:
- raise TypeError("Unsupported dataset type: %s" % dataset.output_classes)
-
-
-def _padded_batch_dense_window(dataset, padded_shape, padding_value=None):
- """Batches a window of dense tensors with padding."""
-
- padded_shape = math_ops.cast(
- convert.partial_shape_to_tensor(padded_shape), dtypes.int32)
-
- def key_fn(_):
- return np.int64(0)
-
- def max_init_fn(_):
- return padded_shape
-
- def max_reduce_fn(state, value):
- """Computes the maximum shape to pad to."""
- condition = math_ops.reduce_all(
- math_ops.logical_or(
- math_ops.less_equal(array_ops.shape(value), padded_shape),
- math_ops.equal(padded_shape, -1)))
- assert_op = control_flow_ops.Assert(condition, [
- "Actual shape greater than padded shape: ",
- array_ops.shape(value), padded_shape
- ])
- with ops.control_dependencies([assert_op]):
- return math_ops.maximum(state, array_ops.shape(value))
-
- def finalize_fn(state):
- return state
-
- # Compute the padded shape.
- max_reducer = grouping.Reducer(max_init_fn, max_reduce_fn, finalize_fn)
- padded_shape = get_single_element.get_single_element(
- dataset.apply(grouping.group_by_reducer(key_fn, max_reducer)))
-
- if padding_value is None:
- if dataset.output_types == dtypes.string:
- padding_value = ""
- elif dataset.output_types == dtypes.bool:
- padding_value = False
- elif dataset.output_types == dtypes.variant:
- raise TypeError("Unable to create padding for field of type 'variant'")
- else:
- padding_value = 0
-
- def batch_init_fn(_):
- batch_shape = array_ops.concat(
- [np.array([0], dtype=np.int32), padded_shape], 0)
- return gen_array_ops.empty(batch_shape, dtype=dataset.output_types)
-
- def batch_reduce_fn(state, value):
- return array_ops.concat([state, [value]], 0)
-
- def pad_fn(value):
- shape = array_ops.shape(value)
- left = array_ops.zeros_like(shape)
- right = padded_shape - shape
- return array_ops.pad(
- value, array_ops.stack([left, right], 1), constant_values=padding_value)
-
- batch_reducer = grouping.Reducer(batch_init_fn, batch_reduce_fn, finalize_fn)
- return get_single_element.get_single_element(
- dataset.map(pad_fn).apply(
- grouping.group_by_reducer(key_fn, batch_reducer)))
-
-
-def _padded_batch_sparse_window(dataset, padded_shape):
- """Batches a window of sparse tensors with padding."""
-
- def key_fn(_):
- return np.int64(0)
-
- def max_init_fn(_):
- return convert.partial_shape_to_tensor(padded_shape)
-
- def max_reduce_fn(state, value):
- """Computes the maximum shape to pad to."""
- condition = math_ops.reduce_all(
- math_ops.logical_or(
- math_ops.less_equal(value.dense_shape, padded_shape),
- math_ops.equal(padded_shape, -1)))
- assert_op = control_flow_ops.Assert(condition, [
- "Actual shape greater than padded shape: ", value.dense_shape,
- padded_shape
- ])
- with ops.control_dependencies([assert_op]):
- return math_ops.maximum(state, value.dense_shape)
-
- def finalize_fn(state):
- return state
-
- # Compute the padded shape.
- max_reducer = grouping.Reducer(max_init_fn, max_reduce_fn, finalize_fn)
- padded_shape = get_single_element.get_single_element(
- dataset.apply(grouping.group_by_reducer(key_fn, max_reducer)))
-
- def batch_init_fn(_):
- indices_shape = array_ops.concat([[0], [array_ops.size(padded_shape) + 1]],
- 0)
- return sparse_tensor.SparseTensor(
- indices=gen_array_ops.empty(indices_shape, dtype=dtypes.int64),
- values=constant_op.constant([], shape=[0], dtype=dataset.output_types),
- dense_shape=array_ops.concat(
- [np.array([0], dtype=np.int64), padded_shape], 0))
-
- def batch_reduce_fn(state, value):
- padded_value = sparse_tensor.SparseTensor(
- indices=value.indices, values=value.values, dense_shape=padded_shape)
- reshaped_value = sparse_ops.sparse_reshape(
- padded_value,
- array_ops.concat(
- [np.array([1], dtype=np.int64), padded_value.dense_shape], 0))
- return sparse_ops.sparse_concat(0, [state, reshaped_value])
-
- reducer = grouping.Reducer(batch_init_fn, batch_reduce_fn, finalize_fn)
- return get_single_element.get_single_element(
- dataset.apply(grouping.group_by_reducer(key_fn, reducer)))
-
-
-class _UnbatchDataset(dataset_ops.UnaryDataset):
- """A dataset that splits the elements of its input into multiple elements."""
-
- def __init__(self, input_dataset):
- """See `unbatch()` for more details."""
- super(_UnbatchDataset, self).__init__(input_dataset)
- flat_shapes = nest.flatten(input_dataset.output_shapes)
- if any(s.ndims == 0 for s in flat_shapes):
- raise ValueError("Cannot unbatch an input with scalar components.")
- known_batch_dim = tensor_shape.Dimension(None)
- for s in flat_shapes:
- try:
- known_batch_dim = known_batch_dim.merge_with(s[0])
- except ValueError:
- raise ValueError("Cannot unbatch an input whose components have "
- "different batch sizes.")
- self._input_dataset = input_dataset
-
- def _as_variant_tensor(self):
- return gen_dataset_ops.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
-
-
+@deprecation.deprecated(None, "Use `tf.data.experimental.unbatch()`.")
def unbatch():
"""Splits elements of a dataset into multiple elements on the batch dimension.
@@ -403,39 +92,7 @@
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
-
- 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.
- def normalize(arg, *rest):
- if rest:
- return sparse.serialize_many_sparse_tensors((arg,) + rest)
- else:
- return sparse.serialize_many_sparse_tensors(arg)
-
- normalized_dataset = dataset.map(normalize)
-
- # NOTE(mrry): Our `map()` has lost information about the sparseness
- # of any SparseTensor components, so re-apply the structure of the
- # original dataset.
- restructured_dataset = _RestructuredDataset(
- normalized_dataset,
- dataset.output_types,
- dataset.output_shapes,
- dataset.output_classes,
- allow_unsafe_cast=True)
- return _UnbatchDataset(restructured_dataset)
-
- return _apply_fn
+ return batching.unbatch()
@deprecation.deprecated(
@@ -514,135 +171,8 @@
return _apply_fn
-class _DenseToSparseBatchDataset(dataset_ops.UnaryDataset):
- """A `Dataset` that batches ragged dense elements into `tf.SparseTensor`s."""
-
- def __init__(self, input_dataset, batch_size, row_shape):
- """See `Dataset.dense_to_sparse_batch()` for more details."""
- super(_DenseToSparseBatchDataset, self).__init__(input_dataset)
- if not isinstance(input_dataset.output_types, dtypes.DType):
- raise TypeError("DenseToSparseDataset requires an input whose elements "
- "have a single component, whereas the input has %r." %
- input_dataset.output_types)
- self._input_dataset = input_dataset
- self._batch_size = batch_size
- self._row_shape = row_shape
-
- def _as_variant_tensor(self):
- return gen_dataset_ops.dense_to_sparse_batch_dataset(
- self._input_dataset._as_variant_tensor(), # pylint: disable=protected-access
- self._batch_size,
- row_shape=convert.partial_shape_to_tensor(self._row_shape),
- **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
-
-
-class _RestructuredDataset(dataset_ops.UnaryDataset):
- """An internal helper for changing the structure and shape of a dataset."""
-
- def __init__(self,
- dataset,
- output_types,
- output_shapes=None,
- output_classes=None,
- allow_unsafe_cast=False):
- """Creates a new dataset with the given output types and shapes.
-
- The given `dataset` must have a structure that is convertible:
- * `dataset.output_types` must be the same as `output_types` module nesting.
- * Each shape in `dataset.output_shapes` must be compatible with each shape
- in `output_shapes` (if given).
-
- Note: This helper permits "unsafe casts" for shapes, equivalent to using
- `tf.Tensor.set_shape()` where domain-specific knowledge is available.
-
- Args:
- dataset: A `Dataset` object.
- output_types: A nested structure of `tf.DType` objects.
- output_shapes: (Optional.) A nested structure of `tf.TensorShape` objects.
- If omitted, the shapes will be inherited from `dataset`.
- output_classes: (Optional.) A nested structure of class types.
- If omitted, the class types will be inherited from `dataset`.
- allow_unsafe_cast: (Optional.) If `True`, the caller may switch the
- reported output types and shapes of the restructured dataset, e.g. to
- switch a sparse tensor represented as `tf.variant` to its user-visible
- type and shape.
-
- Raises:
- ValueError: If either `output_types` or `output_shapes` is not compatible
- with the structure of `dataset`.
- """
- super(_RestructuredDataset, self).__init__(dataset)
- self._input_dataset = dataset
-
- if not allow_unsafe_cast:
- # Validate that the types are compatible.
- output_types = nest.map_structure(dtypes.as_dtype, output_types)
- flat_original_types = nest.flatten(dataset.output_types)
- flat_new_types = nest.flatten(output_types)
- if flat_original_types != flat_new_types:
- raise ValueError(
- "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))
- else:
- if not allow_unsafe_cast:
- # Validate that the shapes are compatible.
- nest.assert_same_structure(output_types, output_shapes)
- flat_original_shapes = nest.flatten(dataset.output_shapes)
- flat_new_shapes = nest.flatten_up_to(output_types, output_shapes)
-
- for original_shape, new_shape in zip(flat_original_shapes,
- flat_new_shapes):
- if not original_shape.is_compatible_with(new_shape):
- raise ValueError(
- "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_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
-
- 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
-
-
+# TODO(b/116817045): Move this to `tf.data.experimental` when the `with_shape()`
+# function is available in the core.
def assert_element_shape(expected_shapes):
"""Assert the shape of this `Dataset`.
@@ -687,7 +217,8 @@
def _apply_fn(dataset):
output_shapes = _merge_output_shapes(dataset.output_shapes,
expected_shapes)
- return _RestructuredDataset(
+ # pylint: disable=protected-access
+ return batching._RestructuredDataset(
dataset.map(_check_shape),
dataset.output_types,
output_shapes=output_shapes,
@@ -696,49 +227,7 @@
return _apply_fn
-class _MapAndBatchDataset(dataset_ops.MapDataset):
- """A `Dataset` that maps a function over a batch of elements."""
-
- def __init__(self, input_dataset, map_func, batch_size, num_parallel_calls,
- drop_remainder):
- """See `Dataset.map()` for details."""
- super(_MapAndBatchDataset, self).__init__(input_dataset, map_func)
- self._batch_size_t = ops.convert_to_tensor(
- batch_size, dtype=dtypes.int64, name="batch_size")
- self._num_parallel_calls_t = ops.convert_to_tensor(
- num_parallel_calls, dtype=dtypes.int64, name="num_parallel_calls")
- self._drop_remainder_t = ops.convert_to_tensor(
- drop_remainder, dtype=dtypes.bool, name="drop_remainder")
-
- self._batch_size = batch_size
- self._drop_remainder = drop_remainder
-
- def _as_variant_tensor(self):
- # pylint: disable=protected-access
- input_resource = self._input_dataset._as_variant_tensor()
- return gen_dataset_ops.map_and_batch_dataset_v2(
- input_resource,
- self._map_func.captured_inputs,
- f=self._map_func,
- batch_size=self._batch_size_t,
- num_parallel_calls=self._num_parallel_calls_t,
- drop_remainder=self._drop_remainder_t,
- **dataset_ops.flat_structure(self))
- # pylint: enable=protected-access
-
- @property
- def output_shapes(self):
- dim = self._batch_size if self._drop_remainder else None
- return nest.pack_sequence_as(self._output_shapes, [
- tensor_shape.vector(dim).concatenate(s)
- for s in nest.flatten(self._output_shapes)
- ])
-
- @property
- def output_types(self):
- return self._output_types
-
-
+@deprecation.deprecated(None, "Use `tf.data.experimental.map_and_batch(...)`.")
def map_and_batch(map_func,
batch_size,
num_parallel_batches=None,
@@ -779,17 +268,5 @@
ValueError: If both `num_parallel_batches` and `num_parallel_calls` are
specified.
"""
-
- if num_parallel_batches is None and num_parallel_calls is None:
- num_parallel_calls = batch_size
- elif num_parallel_batches is not None and num_parallel_calls is None:
- num_parallel_calls = batch_size * num_parallel_batches
- elif num_parallel_batches is not None and num_parallel_calls is not None:
- raise ValueError("The `num_parallel_batches` and `num_parallel_calls` "
- "arguments are mutually exclusive.")
-
- def _apply_fn(dataset):
- return _MapAndBatchDataset(dataset, map_func, batch_size,
- num_parallel_calls, drop_remainder)
-
- return _apply_fn
+ return batching.map_and_batch(map_func, batch_size, num_parallel_batches,
+ drop_remainder, num_parallel_calls)
diff --git a/tensorflow/contrib/data/python/ops/counter.py b/tensorflow/contrib/data/python/ops/counter.py
index 6ef65f9..4ff5bf3 100644
--- a/tensorflow/contrib/data/python/ops/counter.py
+++ b/tensorflow/contrib/data/python/ops/counter.py
@@ -17,13 +17,12 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.contrib.data.python.ops import scan_ops
-
-from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.experimental.ops import counter
from tensorflow.python.framework import dtypes
-from tensorflow.python.framework import ops
+from tensorflow.python.util import deprecation
+@deprecation.deprecated(None, "Use `tf.data.experimental.Counter(...)`.")
def Counter(start=0, step=1, dtype=dtypes.int64):
"""Creates a `Dataset` that counts from `start` in steps of size `step`.
@@ -46,8 +45,4 @@
Returns:
A `Dataset` of scalar `dtype` elements.
"""
- with ops.name_scope("counter"):
- start = ops.convert_to_tensor(start, dtype=dtype, name="start")
- step = ops.convert_to_tensor(step, dtype=dtype, name="step")
- return dataset_ops.Dataset.from_tensors(0).repeat(None).apply(
- scan_ops.scan(start, lambda state, _: (state + step, state)))
+ return counter.Counter(start, step, dtype)
diff --git a/tensorflow/contrib/data/python/ops/enumerate_ops.py b/tensorflow/contrib/data/python/ops/enumerate_ops.py
index 490281e..a21da4d 100644
--- a/tensorflow/contrib/data/python/ops/enumerate_ops.py
+++ b/tensorflow/contrib/data/python/ops/enumerate_ops.py
@@ -17,12 +17,13 @@
from __future__ import division
from __future__ import print_function
-import numpy as np
-from tensorflow.python.data.ops import dataset_ops
-from tensorflow.python.framework import dtypes
+from tensorflow.python.data.experimental.ops import enumerate_ops
+from tensorflow.python.util import deprecation
+@deprecation.deprecated(None,
+ "Use `tf.data.experimental.enumerate_dataset(...)`.")
def enumerate_dataset(start=0):
"""A transformation that enumerate the elements of a dataset.
@@ -49,10 +50,4 @@
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
-
- def _apply_fn(dataset):
- max_value = np.iinfo(dtypes.int64.as_numpy_dtype).max
- return dataset_ops.Dataset.zip((dataset_ops.Dataset.range(start, max_value),
- dataset))
-
- return _apply_fn
+ return enumerate_ops.enumerate_dataset(start)
diff --git a/tensorflow/contrib/data/python/ops/error_ops.py b/tensorflow/contrib/data/python/ops/error_ops.py
index 615dbca..0559a2e 100644
--- a/tensorflow/contrib/data/python/ops/error_ops.py
+++ b/tensorflow/contrib/data/python/ops/error_ops.py
@@ -17,11 +17,11 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.contrib.data.python.ops import contrib_op_loader # pylint: disable=unused-import
-from tensorflow.contrib.data.python.ops import gen_dataset_ops
-from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.experimental.ops import error_ops
+from tensorflow.python.util import deprecation
+@deprecation.deprecated(None, "Use `tf.data.experimental.ignore_errors()`.")
def ignore_errors():
"""Creates a `Dataset` from another `Dataset` and silently ignores any errors.
@@ -44,34 +44,4 @@
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
-
- def _apply_fn(dataset):
- return _IgnoreErrorsDataset(dataset)
-
- return _apply_fn
-
-
-class _IgnoreErrorsDataset(dataset_ops.UnaryDataset):
- """A `Dataset` that silently ignores errors when computing its input."""
-
- def __init__(self, input_dataset):
- """See `Dataset.ignore_errors()` for details."""
- super(_IgnoreErrorsDataset, self).__init__(input_dataset)
- self._input_dataset = input_dataset
-
- def _as_variant_tensor(self):
- return gen_dataset_ops.ignore_errors_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 self._input_dataset.output_shapes
-
- @property
- def output_types(self):
- return self._input_dataset.output_types
+ return error_ops.ignore_errors()
diff --git a/tensorflow/contrib/data/python/ops/get_single_element.py b/tensorflow/contrib/data/python/ops/get_single_element.py
index a6713b0..58ad9ee 100644
--- a/tensorflow/contrib/data/python/ops/get_single_element.py
+++ b/tensorflow/contrib/data/python/ops/get_single_element.py
@@ -19,13 +19,13 @@
import numpy as np
-from tensorflow.contrib.data.python.ops import grouping
+from tensorflow.python.data.experimental.ops import get_single_element as experimental_get_single_element
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 import deprecation
+@deprecation.deprecated(None,
+ "Use `tf.data.experimental.get_single_element(...)`.")
def get_single_element(dataset):
"""Returns the single element in `dataset` as a nested structure of tensors.
@@ -61,18 +61,10 @@
InvalidArgumentError (at runtime): if `dataset` does not contain exactly
one element.
"""
- if not isinstance(dataset, dataset_ops.Dataset):
- 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
- **dataset_ops.flat_structure(dataset)))
- return sparse.deserialize_sparse_tensors(
- nested_ret, dataset.output_types, dataset.output_shapes,
- dataset.output_classes)
+ return experimental_get_single_element.get_single_element(dataset)
+@deprecation.deprecated(None, "Use `tf.data.Dataset.reduce(...)`.")
def reduce_dataset(dataset, reducer):
"""Returns the result of reducing the `dataset` using `reducer`.
@@ -90,11 +82,4 @@
if not isinstance(dataset, dataset_ops.Dataset):
raise TypeError("`dataset` must be a `tf.data.Dataset` object.")
- # The sentinel dataset is used in case the reduced dataset is empty.
- sentinel_dataset = dataset_ops.Dataset.from_tensors(
- reducer.finalize_func(reducer.init_func(np.int64(0))))
- reduced_dataset = dataset.apply(
- grouping.group_by_reducer(lambda x: np.int64(0), reducer))
-
- return get_single_element(
- reduced_dataset.concatenate(sentinel_dataset).take(1))
+ return dataset.reduce(reducer.init_func(np.int64(0)), reducer.reduce_func)
diff --git a/tensorflow/contrib/data/python/ops/grouping.py b/tensorflow/contrib/data/python/ops/grouping.py
index 7cae33b..a99dc2f 100644
--- a/tensorflow/contrib/data/python/ops/grouping.py
+++ b/tensorflow/contrib/data/python/ops/grouping.py
@@ -17,20 +17,13 @@
from __future__ import division
from __future__ import print_function
-import numpy as np
-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 ops
-from tensorflow.python.framework import tensor_shape
-from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import check_ops
-from tensorflow.python.ops import gen_dataset_ops
-from tensorflow.python.ops import math_ops
+from tensorflow.python.data.experimental.ops import grouping
+from tensorflow.python.util import deprecation
+@deprecation.deprecated(None,
+ "Use `tf.data.experimental.group_by_reducer(...)`.")
def group_by_reducer(key_func, reducer):
"""A transformation that groups elements and performs a reduction.
@@ -52,14 +45,11 @@
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
-
- def _apply_fn(dataset):
- """Function from `Dataset` to `Dataset` that applies the transformation."""
- return _GroupByReducerDataset(dataset, key_func, reducer)
-
- return _apply_fn
+ return grouping.group_by_reducer(key_func, reducer)
+@deprecation.deprecated(None,
+ "Use `tf.data.experimental.group_by_window(...)`.")
def group_by_window(key_func,
reduce_func,
window_size=None,
@@ -98,27 +88,12 @@
ValueError: if neither or both of {`window_size`, `window_size_func`} are
passed.
"""
- if (window_size is not None and window_size_func or
- not (window_size is not None or window_size_func)):
- raise ValueError("Must pass either window_size or window_size_func.")
-
- if window_size is not None:
-
- def constant_window_func(unused_key):
- return ops.convert_to_tensor(window_size, dtype=dtypes.int64)
-
- window_size_func = constant_window_func
-
- assert window_size_func is not None
-
- def _apply_fn(dataset):
- """Function from `Dataset` to `Dataset` that applies the transformation."""
- return _GroupByWindowDataset(dataset, key_func, reduce_func,
- window_size_func)
-
- return _apply_fn
+ return grouping.group_by_window(key_func, reduce_func, window_size,
+ window_size_func)
+@deprecation.deprecated(
+ None, "Use `tf.data.experimental.bucket_by_sequence_length(...)`.")
def bucket_by_sequence_length(element_length_func,
bucket_boundaries,
bucket_batch_sizes,
@@ -163,342 +138,12 @@
Raises:
ValueError: if `len(bucket_batch_sizes) != len(bucket_boundaries) + 1`.
"""
- with ops.name_scope("bucket_by_seq_length"):
- if len(bucket_batch_sizes) != (len(bucket_boundaries) + 1):
- raise ValueError(
- "len(bucket_batch_sizes) must equal len(bucket_boundaries) + 1")
-
- batch_sizes = constant_op.constant(bucket_batch_sizes, dtype=dtypes.int64)
-
- def element_to_bucket_id(*args):
- """Return int64 id of the length bucket for this element."""
- seq_length = element_length_func(*args)
-
- boundaries = list(bucket_boundaries)
- buckets_min = [np.iinfo(np.int32).min] + boundaries
- buckets_max = boundaries + [np.iinfo(np.int32).max]
- conditions_c = math_ops.logical_and(
- math_ops.less_equal(buckets_min, seq_length),
- math_ops.less(seq_length, buckets_max))
- bucket_id = math_ops.reduce_min(array_ops.where(conditions_c))
-
- return bucket_id
-
- def window_size_fn(bucket_id):
- # The window size is set to the batch size for this bucket
- window_size = batch_sizes[bucket_id]
- return window_size
-
- def make_padded_shapes(shapes, none_filler=None):
- padded = []
- for shape in nest.flatten(shapes):
- shape = tensor_shape.TensorShape(shape)
- shape = [
- none_filler if d.value is None else d
- for d in shape
- ]
- padded.append(shape)
- return nest.pack_sequence_as(shapes, padded)
-
- def batching_fn(bucket_id, grouped_dataset):
- """Batch elements in dataset."""
- batch_size = window_size_fn(bucket_id)
- if no_padding:
- return grouped_dataset.batch(batch_size)
- none_filler = None
- if pad_to_bucket_boundary:
- err_msg = ("When pad_to_bucket_boundary=True, elements must have "
- "length < max(bucket_boundaries).")
- check = check_ops.assert_less(
- bucket_id,
- constant_op.constant(len(bucket_batch_sizes) - 1,
- dtype=dtypes.int64),
- message=err_msg)
- with ops.control_dependencies([check]):
- boundaries = constant_op.constant(bucket_boundaries,
- dtype=dtypes.int64)
- bucket_boundary = boundaries[bucket_id]
- none_filler = bucket_boundary - 1
- shapes = make_padded_shapes(
- padded_shapes or grouped_dataset.output_shapes,
- none_filler=none_filler)
- return grouped_dataset.padded_batch(batch_size, shapes, padding_values)
-
- def _apply_fn(dataset):
- return dataset.apply(
- group_by_window(element_to_bucket_id, batching_fn,
- window_size_func=window_size_fn))
-
- return _apply_fn
+ return grouping.bucket_by_sequence_length(
+ element_length_func, bucket_boundaries, bucket_batch_sizes, padded_shapes,
+ padding_values, pad_to_bucket_boundary, no_padding)
-def _map_x_dataset(map_func):
- """A transformation that maps `map_func` across its input.
-
- This transformation is similar to `tf.data.Dataset.map`, but in addition to
- supporting dense and sparse tensor inputs, it also supports dataset inputs.
-
- Args:
- map_func: A function mapping a nested structure of tensors and/or datasets
- (having shapes and types defined by `self.output_shapes` and
- `self.output_types`) to another nested structure of tensors and/or
- datasets.
-
- Returns:
- Dataset: A `Dataset`.
- """
-
- def _apply_fn(dataset):
- """Function from `Dataset` to `Dataset` that applies the transformation."""
- return _MapXDataset(dataset, map_func)
-
- return _apply_fn
-
-
-# TODO(b/115382007) Remove this once canned reducers move to core.
-def window_dataset(window_size):
- """A transformation that creates window datasets from the input dataset.
-
- The resulting datasets will contain `window_size` elements (or
- `N % window_size` for the last dataset if `window_size` does not divide the
- number of input elements `N` evenly).
-
- Args:
- window_size: A `tf.int64` scalar `tf.Tensor`, representing the number of
- consecutive elements of the input dataset to combine into a window.
-
- Returns:
- Dataset: A `Dataset`.
- """
-
- def _apply_fn(dataset):
- return dataset_ops.WindowDataset(
- dataset,
- size=window_size,
- shift=window_size,
- stride=1,
- drop_remainder=False)
-
- return _apply_fn
-
-
-class _GroupByReducerDataset(dataset_ops.UnaryDataset):
- """A `Dataset` that groups its input and performs a reduction."""
-
- def __init__(self, input_dataset, key_func, reducer):
- """See `group_by_reducer()` for details."""
- super(_GroupByReducerDataset, self).__init__(input_dataset)
-
- self._input_dataset = input_dataset
-
- self._make_key_func(key_func, input_dataset)
- self._make_init_func(reducer.init_func)
- self._make_reduce_func(reducer.reduce_func, input_dataset)
- self._make_finalize_func(reducer.finalize_func)
-
- def _make_key_func(self, key_func, input_dataset):
- """Make wrapping Defun for key_func."""
- wrapped_func = dataset_ops.StructuredFunctionWrapper(
- key_func, "tf.contrib.data.group_by_reducer()", input_dataset)
- if not (
- wrapped_func.output_types == dtypes.int64 and
- wrapped_func.output_shapes.is_compatible_with(tensor_shape.scalar())):
- 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.function
-
- def _make_init_func(self, init_func):
- """Make wrapping Defun for init_func."""
- wrapped_func = dataset_ops.StructuredFunctionWrapper(
- init_func, "tf.contrib.data.group_by_reducer()",
- input_classes=ops.Tensor, input_shapes=tensor_shape.scalar(),
- input_types=dtypes.int64)
- self._init_func = wrapped_func.function
- self._state_classes = wrapped_func.output_classes
- self._state_shapes = wrapped_func.output_shapes
- self._state_types = wrapped_func.output_types
-
- 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`.
- need_to_rerun = True
- while need_to_rerun:
-
- wrapped_func = dataset_ops.StructuredFunctionWrapper(
- reduce_func, "tf.contrib.data.group_by_reducer()",
- 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),
- 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)):
- if not issubclass(new_state_class, state_class):
- raise TypeError(
- "The element classes for the new state must match the initial "
- "state. Expected %s; got %s." %
- (self._state_classes, wrapped_func.output_classes))
-
- # 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)):
- 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))
-
- # Extract shape information from the returned values.
- flat_state_shapes = nest.flatten(self._state_shapes)
- flat_new_state_shapes = nest.flatten(wrapped_func.output_shapes)
- weakened_state_shapes = [
- original.most_specific_compatible_shape(new)
- for original, new in zip(flat_state_shapes, flat_new_state_shapes)
- ]
-
- need_to_rerun = False
- for original_shape, weakened_shape in zip(flat_state_shapes,
- weakened_state_shapes):
- if original_shape.ndims is not None and (
- weakened_shape.ndims is None or
- original_shape.as_list() != weakened_shape.as_list()):
- need_to_rerun = True
- break
-
- if need_to_rerun:
- self._state_shapes = nest.pack_sequence_as(self._state_shapes,
- weakened_state_shapes)
-
- self._reduce_func = wrapped_func.function
- self._reduce_func.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, "tf.contrib.data.group_by_reducer()",
- input_classes=self._state_classes, input_shapes=self._state_shapes,
- input_types=self._state_types)
- self._finalize_func = wrapped_func.function
- self._output_classes = wrapped_func.output_classes
- self._output_shapes = wrapped_func.output_shapes
- self._output_types = wrapped_func.output_types
-
- @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 _as_variant_tensor(self):
- return gen_dataset_ops.group_by_reducer_dataset(
- self._input_dataset._as_variant_tensor(), # pylint: disable=protected-access
- self._key_func.captured_inputs,
- self._init_func.captured_inputs,
- self._reduce_func.captured_inputs,
- self._finalize_func.captured_inputs,
- key_func=self._key_func,
- init_func=self._init_func,
- reduce_func=self._reduce_func,
- finalize_func=self._finalize_func,
- **dataset_ops.flat_structure(self))
-
-
-class _GroupByWindowDataset(dataset_ops.UnaryDataset):
- """A `Dataset` that groups its input and performs a windowed reduction."""
-
- def __init__(self, input_dataset, key_func, reduce_func, window_size_func):
- """See `group_by_window()` for details."""
- super(_GroupByWindowDataset, self).__init__(input_dataset)
-
- self._input_dataset = input_dataset
-
- self._make_key_func(key_func, input_dataset)
- self._make_reduce_func(reduce_func, input_dataset)
- self._make_window_size_func(window_size_func)
-
- def _make_window_size_func(self, window_size_func):
- """Make wrapping Defun for window_size_func."""
- def window_size_func_wrapper(key):
- return ops.convert_to_tensor(window_size_func(key), dtype=dtypes.int64)
- wrapped_func = dataset_ops.StructuredFunctionWrapper(
- window_size_func_wrapper, "tf.contrib.data.group_by_window()",
- 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())):
- raise ValueError(
- "`window_size_func` must return a single tf.int64 scalar tensor.")
- self._window_size_func = wrapped_func.function
-
- 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(
- key_func_wrapper, "tf.contrib.data.group_by_window()", input_dataset)
- if not (
- wrapped_func.output_types == dtypes.int64 and
- wrapped_func.output_shapes.is_compatible_with(tensor_shape.scalar())):
- raise ValueError(
- "`key_func` must return a single tf.int64 scalar tensor.")
- self._key_func = wrapped_func.function
-
- def _make_reduce_func(self, reduce_func, input_dataset):
- """Make wrapping Defun for reduce_func."""
- nested_dataset = dataset_ops._NestedDatasetComponent(input_dataset) # pylint: disable=protected-access
- wrapped_func = dataset_ops.StructuredFunctionWrapper(
- reduce_func, "tf.contrib.data.reduce_by_window()",
- input_classes=(ops.Tensor, nested_dataset),
- input_shapes=(tensor_shape.scalar(), nested_dataset),
- input_types=(dtypes.int64, nested_dataset),
- experimental_nested_dataset_support=True)
- if not isinstance(
- wrapped_func.output_classes, dataset_ops._NestedDatasetComponent): # pylint: disable=protected-access
- raise TypeError("`reduce_func` must return a `Dataset` object.")
- self._output_classes = wrapped_func.output_classes.output_classes
- self._output_types = wrapped_func.output_types.output_types
- self._output_shapes = wrapped_func.output_shapes.output_shapes
- self._reduce_func = wrapped_func.function
-
- @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 _as_variant_tensor(self):
- return gen_dataset_ops.group_by_window_dataset(
- self._input_dataset._as_variant_tensor(), # pylint: disable=protected-access
- self._key_func.captured_inputs,
- self._reduce_func.captured_inputs,
- self._window_size_func.captured_inputs,
- key_func=self._key_func,
- reduce_func=self._reduce_func,
- window_size_func=self._window_size_func,
- **dataset_ops.flat_structure(self))
-
-
-class Reducer(object):
+class Reducer(grouping.Reducer):
"""A reducer is used for reducing a set of elements.
A reducer is represented as a tuple of the three functions:
@@ -507,58 +152,6 @@
3) finalization function: state => result
"""
+ @deprecation.deprecated(None, "Use `tf.data.experimental.Reducer(...)`.")
def __init__(self, init_func, reduce_func, finalize_func):
- self._init_func = init_func
- self._reduce_func = reduce_func
- self._finalize_func = finalize_func
-
- @property
- def init_func(self):
- return self._init_func
-
- @property
- def reduce_func(self):
- return self._reduce_func
-
- @property
- def finalize_func(self):
- return self._finalize_func
-
-
-class _MapXDataset(dataset_ops.UnaryDataset):
- """A `Dataset` that maps a function over elements in its input."""
-
- def __init__(self, input_dataset, map_func):
- """See `map_x_dataset()` for details."""
- super(_MapXDataset, self).__init__(input_dataset)
- self._input_dataset = input_dataset
-
- wrapped_func = dataset_ops.StructuredFunctionWrapper(
- map_func,
- "tf.contrib.data.map_x_dataset()",
- input_dataset,
- experimental_nested_dataset_support=True)
- self._output_classes = wrapped_func.output_classes
- self._output_shapes = wrapped_func.output_shapes
- self._output_types = wrapped_func.output_types
- self._map_func = wrapped_func.function
-
- def _as_variant_tensor(self):
- input_t = self._input_dataset._as_variant_tensor() # pylint: disable=protected-access
- return gen_dataset_ops.map_dataset(
- input_t,
- self._map_func.captured_inputs,
- f=self._map_func,
- **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
+ super(Reducer, self).__init__(init_func, reduce_func, finalize_func)
diff --git a/tensorflow/contrib/data/python/ops/interleave_ops.py b/tensorflow/contrib/data/python/ops/interleave_ops.py
index bfa3fdf..f50da4d 100644
--- a/tensorflow/contrib/data/python/ops/interleave_ops.py
+++ b/tensorflow/contrib/data/python/ops/interleave_ops.py
@@ -17,21 +17,12 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.contrib import stateless
-from tensorflow.contrib.data.python.ops import contrib_op_loader # pylint: disable=unused-import
-from tensorflow.contrib.data.python.ops import gen_dataset_ops
-from tensorflow.contrib.data.python.ops import random_ops
-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.framework import dtypes
-from tensorflow.python.framework import ops
-from tensorflow.python.framework import tensor_shape
-from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import math_ops
+from tensorflow.python.data.experimental.ops import interleave_ops
from tensorflow.python.util import deprecation
+@deprecation.deprecated(None,
+ "Use `tf.data.experimental.parallel_interleave(...)`.")
def parallel_interleave(map_func,
cycle_length,
block_length=1,
@@ -81,12 +72,9 @@
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
- def _apply_fn(dataset):
- return readers.ParallelInterleaveDataset(
- dataset, map_func, cycle_length, block_length, sloppy,
- buffer_output_elements, prefetch_input_elements)
-
- return _apply_fn
+ return interleave_ops.parallel_interleave(
+ map_func, cycle_length, block_length, sloppy, buffer_output_elements,
+ prefetch_input_elements)
@deprecation.deprecated(
@@ -140,61 +128,12 @@
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
- def _apply_fn(dataset):
- return readers.ParallelInterleaveDataset(
- dataset,
- map_func,
- cycle_length,
- block_length,
- sloppy=True,
- buffer_output_elements=None,
- prefetch_input_elements=None)
-
- return _apply_fn
+ return interleave_ops.parallel_interleave(
+ map_func, cycle_length, block_length, sloppy=True)
-class _DirectedInterleaveDataset(dataset_ops.Dataset):
- """A substitute for `Dataset.interleave()` on a fixed list of datasets."""
-
- def __init__(self, selector_input, data_inputs):
- self._selector_input = selector_input
- self._data_inputs = list(data_inputs)
-
- for data_input in data_inputs[1:]:
- if (data_input.output_types != data_inputs[0].output_types or
- data_input.output_classes != data_inputs[0].output_classes):
- raise TypeError("All datasets must have the same type and class.")
-
- def _as_variant_tensor(self):
- # pylint: disable=protected-access
- return gen_dataset_ops.directed_interleave_dataset(
- self._selector_input._as_variant_tensor(),
- [data_input._as_variant_tensor() for data_input in self._data_inputs],
- **dataset_ops.flat_structure(self))
- # pylint: enable=protected-access
-
- def _inputs(self):
- 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
-
-
+@deprecation.deprecated(None,
+ "Use `tf.data.experimental.sample_from_datasets(...)`.")
def sample_from_datasets(datasets, weights=None, seed=None):
"""Samples elements at random from the datasets in `datasets`.
@@ -218,64 +157,11 @@
ValueError: If the `weights` argument is specified and does not match the
length of the `datasets` element.
"""
- num_datasets = len(datasets)
- if not isinstance(weights, dataset_ops.Dataset):
- if weights is None:
- # Select inputs with uniform probability.
- logits = [[1.0] * num_datasets]
-
- else:
- # Use the given `weights` as the probability of choosing the respective
- # input.
- weights = ops.convert_to_tensor(weights, name="weights")
- if weights.dtype not in (dtypes.float32, dtypes.float64):
- raise TypeError("`weights` must be convertible to a tensor of "
- "`tf.float32` or `tf.float64` elements.")
- if not weights.shape.is_compatible_with([num_datasets]):
- raise ValueError(
- "`weights` must be a vector of length `len(datasets)`.")
-
- # The `stateless_multinomial()` op expects log-probabilities, as opposed
- # to weights.
- logits = array_ops.expand_dims(math_ops.log(weights, name="logits"), 0)
-
- # NOTE(mrry): We only specialize when `weights` is not a `Dataset`. When it
- # is a `Dataset`, it is possible that evaluating it has a side effect the
- # user depends on.
- if len(datasets) == 1:
- return datasets[0]
-
- def select_dataset_constant_logits(seed):
- return array_ops.squeeze(
- stateless.stateless_multinomial(logits, 1, seed=seed), axis=[0, 1])
-
- selector_input = dataset_ops.MapDataset(
- random_ops.RandomDataset(seed).batch(2),
- select_dataset_constant_logits,
- use_inter_op_parallelism=False)
-
- else:
- # Use each element of the given `weights` dataset as the probability of
- # choosing the respective input.
-
- # The `stateless_multinomial()` op expects log-probabilities, as opposed to
- # weights.
- logits_ds = weights.map(lambda *p: math_ops.log(p, name="logits"))
-
- def select_dataset_varying_logits(logits, seed):
- return array_ops.squeeze(
- stateless.stateless_multinomial(logits, 1, seed=seed), axis=[0, 1])
-
- logits_and_seeds = dataset_ops.Dataset.zip(
- (logits_ds, random_ops.RandomDataset(seed).batch(2)))
- selector_input = dataset_ops.MapDataset(
- logits_and_seeds,
- select_dataset_varying_logits,
- use_inter_op_parallelism=False)
-
- return _DirectedInterleaveDataset(selector_input, datasets)
+ return interleave_ops.sample_from_datasets(datasets, weights, seed)
+@deprecation.deprecated(None,
+ "Use `tf.data.experimental.choose_from_datasets(...)`.")
def choose_from_datasets(datasets, choice_dataset):
"""Creates a dataset that deterministically chooses elements from `datasets`.
@@ -311,10 +197,4 @@
TypeError: If the `datasets` or `choice_dataset` arguments have the wrong
type.
"""
- if not (choice_dataset.output_types == dtypes.int64
- and choice_dataset.output_shapes.is_compatible_with(
- tensor_shape.scalar())
- and choice_dataset.output_classes == ops.Tensor):
- raise TypeError("`choice_dataset` must be a dataset of scalar "
- "`tf.int64` tensors.")
- return _DirectedInterleaveDataset(choice_dataset, datasets)
+ return interleave_ops.choose_from_datasets(datasets, choice_dataset)
diff --git a/tensorflow/contrib/data/python/ops/iterator_ops.py b/tensorflow/contrib/data/python/ops/iterator_ops.py
index 18515e2..48c325c 100644
--- a/tensorflow/contrib/data/python/ops/iterator_ops.py
+++ b/tensorflow/contrib/data/python/ops/iterator_ops.py
@@ -16,15 +16,13 @@
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
-from tensorflow.python.data.ops import iterator_ops
-from tensorflow.python.framework import ops
-from tensorflow.python.ops import gen_dataset_ops
-from tensorflow.python.training import basic_session_run_hooks
-from tensorflow.python.training import checkpoint_management
-from tensorflow.python.training import saver as saver_lib
-from tensorflow.python.training import session_run_hook
+
+from tensorflow.python.data.experimental.ops import iterator_ops
+from tensorflow.python.util import deprecation
+@deprecation.deprecated(
+ None, "Use `tf.data.experimental.make_saveable_from_iterator(...)`.")
def make_saveable_from_iterator(iterator):
"""Returns a SaveableObject for saving/restore iterator state using Saver.
@@ -60,27 +58,10 @@
Note: Not all iterators support checkpointing yet. Attempting to save the
state of an unsupported iterator will throw an error.
"""
- return _Saveable(iterator._iterator_resource) # pylint: disable=protected-access
+ return iterator_ops.make_saveable_from_iterator(iterator)
-class _Saveable(saver_lib.BaseSaverBuilder.SaveableObject):
- """SaveableObject for saving/restoring iterator state."""
-
- def __init__(self, iterator_resource):
- serialized_iterator = gen_dataset_ops.serialize_iterator(iterator_resource)
- specs = [
- saver_lib.BaseSaverBuilder.SaveSpec(serialized_iterator, "",
- iterator_resource.name + "-state")
- ]
- super(_Saveable, self).__init__(iterator_resource, specs,
- iterator_resource.name)
-
- def restore(self, restored_tensors, unused_restored_shapes):
- with ops.colocate_with(self.op):
- return gen_dataset_ops.deserialize_iterator(self.op, restored_tensors[0])
-
-
-class CheckpointInputPipelineHook(session_run_hook.SessionRunHook):
+class CheckpointInputPipelineHook(iterator_ops.CheckpointInputPipelineHook):
"""Checkpoints input pipeline state every N steps or seconds.
This hook saves the state of the iterators in the `Graph` so that when
@@ -125,135 +106,7 @@
collector when building the eval graph.
"""
+ @deprecation.deprecated(
+ None, "Use `tf.data.experimental.CheckpointInputPipelineHook(...)`.")
def __init__(self, estimator):
- """Initializes a `CheckpointInputPipelineHook`.
-
- Args:
- estimator: Estimator.
-
- Raises:
- ValueError: One of `save_steps` or `save_secs` should be set.
- ValueError: At most one of saver or scaffold should be set.
- """
- # `checkpoint_basename` is "input.ckpt" for non-distributed pipelines or
- # of the form "input_<task_type>_<task_id>.ckpt" for distributed pipelines.
- # Note: The default `checkpoint_basename` used by `CheckpointSaverHook` is
- # "model.ckpt". We intentionally choose the input pipeline checkpoint prefix
- # to be different to avoid conflicts with the model checkpoint.
-
- # pylint: disable=protected-access
- checkpoint_prefix = "input"
- if estimator._config.num_worker_replicas > 1:
- # Distributed setting.
- suffix = "_{}_{}".format(estimator._config.task_type,
- estimator._config.task_id)
- checkpoint_prefix += suffix
- # pylint: enable=protected-access
-
- # We use a composition paradigm instead of inheriting from
- # `CheckpointSaverHook` because `Estimator` does an `isinstance` check
- # to check whether a `CheckpointSaverHook` is already present in the list
- # of hooks and if not, adds one. Inheriting from `CheckpointSaverHook`
- # would thwart this behavior. This hook checkpoints *only the iterators*
- # and not the graph variables.
- self._checkpoint_saver_hook = basic_session_run_hooks.CheckpointSaverHook(
- estimator.model_dir,
- save_secs=estimator._config.save_checkpoints_secs, # pylint: disable=protected-access
- save_steps=estimator._config.save_checkpoints_steps, # pylint: disable=protected-access
- checkpoint_basename=checkpoint_prefix + ".ckpt")
-
- # Name for the protocol buffer file that will contain the list of most
- # recent checkpoints stored as a `CheckpointState` protocol buffer.
- # This file, kept in the same directory as the checkpoint files, is
- # automatically managed by the `Saver` to keep track of recent checkpoints.
- # The default name used by the `Saver` for this file is "checkpoint". Here
- # we use the name "checkpoint_<checkpoint_prefix>" so that in case the
- # `checkpoint_dir` is the same as the model checkpoint directory, there are
- # no conflicts during restore.
- self._latest_filename = "checkpoint_" + checkpoint_prefix
- self._first_run = True
-
- def begin(self):
- # Build a Saver that saves all iterators in the `GLOBAL_ITERATORS`
- # collection if no `Saver` or `Scaffold` is provided.
- # pylint: disable=protected-access
- if (self._checkpoint_saver_hook._saver is None and
- self._checkpoint_saver_hook._scaffold is None):
- iterators = ops.get_collection(iterator_ops.GLOBAL_ITERATORS)
- saveables = [_Saveable(i) for i in iterators]
- self._checkpoint_saver_hook._saver = _CustomSaver(saveables,
- self._latest_filename)
- # pylint: enable=protected-access
- self._checkpoint_saver_hook.begin()
-
- def _restore_or_save_initial_ckpt(self, session):
- # Ideally this should be run in after_create_session but is not for the
- # following reason:
- # Currently there is no way of enforcing an order of running the
- # `SessionRunHooks`. Hence it is possible that the `_DatasetInitializerHook`
- # is run *after* this hook. That is troublesome because
- # 1. If a checkpoint exists and this hook restores it, the initializer hook
- # will override it.
- # 2. If no checkpoint exists, this hook will try to save an initialized
- # iterator which will result in an exception.
- #
- # As a temporary fix we enter the following implicit contract between this
- # hook and the _DatasetInitializerHook.
- # 1. The _DatasetInitializerHook initializes the iterator in the call to
- # after_create_session.
- # 2. This hook saves the iterator on the first call to `before_run()`, which
- # is guaranteed to happen after `after_create_session()` of all hooks
- # have been run.
-
- # Check if there is an existing checkpoint. If so, restore from it.
- # pylint: disable=protected-access
- latest_checkpoint_path = checkpoint_management.latest_checkpoint(
- self._checkpoint_saver_hook._checkpoint_dir,
- latest_filename=self._latest_filename)
- if latest_checkpoint_path:
- self._checkpoint_saver_hook._get_saver().restore(session,
- latest_checkpoint_path)
- else:
- # The checkpoint saved here is the state at step "global_step".
- # Note: We do not save the GraphDef or MetaGraphDef here.
- global_step = session.run(self._checkpoint_saver_hook._global_step_tensor)
- self._checkpoint_saver_hook._save(session, global_step)
- self._checkpoint_saver_hook._timer.update_last_triggered_step(global_step)
- # pylint: enable=protected-access
-
- def before_run(self, run_context):
- if self._first_run:
- self._restore_or_save_initial_ckpt(run_context.session)
- self._first_run = False
- return self._checkpoint_saver_hook.before_run(run_context)
-
- def after_run(self, run_context, run_values):
- self._checkpoint_saver_hook.after_run(run_context, run_values)
-
- def end(self, session):
- self._checkpoint_saver_hook.end(session)
-
-
-class _CustomSaver(saver_lib.Saver):
- """`Saver` with a different default `latest_filename`.
-
- This is used in the `CheckpointInputPipelineHook` to avoid conflicts with
- the model ckpt saved by the `CheckpointSaverHook`.
- """
-
- def __init__(self, var_list, latest_filename):
- super(_CustomSaver, self).__init__(var_list)
- self._latest_filename = latest_filename
-
- def save(self,
- sess,
- save_path,
- global_step=None,
- latest_filename=None,
- meta_graph_suffix="meta",
- write_meta_graph=True,
- write_state=True,
- strip_default_attrs=False):
- return super(_CustomSaver, self).save(
- sess, save_path, global_step, latest_filename or self._latest_filename,
- meta_graph_suffix, write_meta_graph, write_state, strip_default_attrs)
+ super(CheckpointInputPipelineHook, self).__init__(estimator)
diff --git a/tensorflow/contrib/data/python/ops/parsing_ops.py b/tensorflow/contrib/data/python/ops/parsing_ops.py
index cfbba70..3aeee9d 100644
--- a/tensorflow/contrib/data/python/ops/parsing_ops.py
+++ b/tensorflow/contrib/data/python/ops/parsing_ops.py
@@ -17,92 +17,12 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.python.data.ops import dataset_ops
-from tensorflow.python.data.util import nest
-from tensorflow.python.framework import dtypes
-from tensorflow.python.framework import ops
-from tensorflow.python.framework import sparse_tensor
-from tensorflow.python.ops import gen_dataset_ops
-from tensorflow.python.ops import parsing_ops
+from tensorflow.python.data.experimental.ops import parsing_ops
+from tensorflow.python.util import deprecation
-class _ParseExampleDataset(dataset_ops.UnaryDataset):
- """A `Dataset` that parses `example` dataset into a `dict` dataset."""
-
- 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)):
- raise TypeError("Input dataset should be a dataset of vectors of strings")
- self._num_parallel_calls = num_parallel_calls
- # pylint: disable=protected-access
- self._features = parsing_ops._prepend_none_dimension(features)
- # sparse_keys and dense_keys come back sorted here.
- (sparse_keys, sparse_types, dense_keys, dense_types, dense_defaults,
- dense_shapes) = parsing_ops._features_to_raw_params(
- self._features, [
- parsing_ops.VarLenFeature, parsing_ops.SparseFeature,
- parsing_ops.FixedLenFeature, parsing_ops.FixedLenSequenceFeature
- ])
- # TODO(b/112859642): Pass sparse_index and sparse_values for SparseFeature.
- (_, dense_defaults_vec, sparse_keys, sparse_types, dense_keys, dense_shapes,
- dense_shape_as_shape) = parsing_ops._process_raw_parameters(
- None, dense_defaults, sparse_keys, sparse_types, dense_keys,
- dense_types, dense_shapes)
- # pylint: enable=protected-access
- self._sparse_keys = sparse_keys
- self._sparse_types = sparse_types
- self._dense_keys = dense_keys
- self._dense_defaults = dense_defaults_vec
- self._dense_shapes = dense_shapes
- self._dense_types = dense_types
- dense_output_shapes = [
- self._input_dataset.output_shapes.concatenate(shape)
- for shape in dense_shape_as_shape
- ]
- sparse_output_shapes = [
- self._input_dataset.output_shapes.concatenate([None])
- for _ in range(len(sparse_keys))
- ]
-
- self._output_shapes = dict(
- zip(self._dense_keys + self._sparse_keys,
- dense_output_shapes + sparse_output_shapes))
- self._output_types = dict(
- zip(self._dense_keys + self._sparse_keys,
- self._dense_types + self._sparse_types))
- self._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))
- ]))
-
- def _as_variant_tensor(self):
- return gen_dataset_ops.parse_example_dataset(
- self._input_dataset._as_variant_tensor(), # pylint: disable=protected-access
- self._num_parallel_calls,
- self._dense_defaults,
- self._sparse_keys,
- self._dense_keys,
- self._sparse_types,
- self._dense_shapes,
- **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
-
-
-# TODO(b/111553342): add arguments names and example names as well.
+@deprecation.deprecated(
+ None, "Use `tf.data.experimental.parse_example_dataset(...)`.")
def parse_example_dataset(features, num_parallel_calls=1):
"""A transformation that parses `Example` protos into a `dict` of tensors.
@@ -130,21 +50,4 @@
Raises:
ValueError: if features argument is None.
"""
- if features is None:
- raise ValueError("Missing: features was %s." % features)
-
- def _apply_fn(dataset):
- """Function from `Dataset` to `Dataset` that applies the transformation."""
- out_dataset = _ParseExampleDataset(dataset, features, num_parallel_calls)
- if any([
- isinstance(feature, parsing_ops.SparseFeature)
- for _, feature in features.items()
- ]):
- # pylint: disable=protected-access
- # pylint: disable=g-long-lambda
- out_dataset = out_dataset.map(
- lambda x: parsing_ops._construct_sparse_tensors_for_sparse_features(
- features, x), num_parallel_calls=num_parallel_calls)
- return out_dataset
-
- return _apply_fn
+ return parsing_ops.parse_example_dataset(features, num_parallel_calls)
diff --git a/tensorflow/contrib/data/python/ops/prefetching_ops.py b/tensorflow/contrib/data/python/ops/prefetching_ops.py
index f994425..adfb390 100644
--- a/tensorflow/contrib/data/python/ops/prefetching_ops.py
+++ b/tensorflow/contrib/data/python/ops/prefetching_ops.py
@@ -17,321 +17,12 @@
from __future__ import division
from __future__ import print_function
-import warnings
-
-from tensorflow.contrib.data.python.ops import contrib_op_loader # pylint: disable=unused-import
-from tensorflow.contrib.data.python.ops import gen_dataset_ops
-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.framework import device as framework_device
-from tensorflow.python.framework import dtypes
-from tensorflow.python.framework import 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 functional_ops
-from tensorflow.python.ops import gen_dataset_ops as core_gen_dataset_ops
-from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.data.experimental.ops import prefetching_ops
+from tensorflow.python.util import deprecation
-def function_buffering_resource(string_arg,
- target_device,
- f,
- buffer_size,
- output_types,
- container="",
- shared_name=None,
- name=None):
- """Creates a FunctionBufferingResource.
-
- A FunctionBufferingResource fills up a buffer by calling a function `f` on
- `target_device`. `f` should take in only a single string argument as input.
-
- Args:
- string_arg: The single string argument to the function.
- target_device: The device to run `f` on.
- f: The function to be executed.
- buffer_size: Size of the buffer to be populated.
- output_types: The output types generated by the function.
- container: (Optional) string. Defaults to "".
- shared_name: (Optional) string.
- name: (Optional) string to name the op.
-
- Returns:
- Handle to a FunctionBufferingResource.
- """
- if shared_name is None:
- shared_name = ""
- return gen_dataset_ops.function_buffering_resource(
- string_arg=string_arg,
- target_device=target_device,
- shared_name=shared_name,
- f=f,
- buffer_size=buffer_size,
- container=container,
- name=name,
- output_types=output_types)
-
-
-def function_buffering_resource_get_next(function_buffer_resource,
- output_types,
- name=None):
- return gen_dataset_ops.function_buffering_resource_get_next(
- function_buffer_resource=function_buffer_resource,
- output_types=output_types,
- name=name)
-
-
-def function_buffering_resource_reset(function_buffer_resource, name=None):
- return gen_dataset_ops.function_buffering_resource_reset(
- function_buffer_resource=function_buffer_resource, name=name)
-
-
-# pylint: disable=protected-access
-class _PrefetchToDeviceIterator(object):
- """A replacement for `tf.data.Iterator` that prefetches to another device.
-
- Args:
- input_dataset: The input dataset
- one_shot: If true, we make a one shot iterator that's already initialized.
- device: A fully specified device string where we want to prefetch to
- buffer_size: Size of the prefetching buffer.
- shared_name: (Optional.) If non-empty, the returned iterator will be
- shared under the given name across multiple sessions that share the
- same devices (e.g. when using a remote server).
-
- Returns:
- An Iterator type object.
- """
-
- def __init__(self,
- input_dataset,
- one_shot,
- device,
- buffer_size,
- shared_name=None):
- self._input_dataset = input_dataset
- self._get_next_call_count = 0
- self._one_shot = one_shot
- if shared_name is None:
- shared_name = ""
-
- if self._one_shot:
- self._input_iterator = input_dataset.make_one_shot_iterator()
- else:
- self._input_iterator = iterator_ops.Iterator.from_structure(
- self._input_dataset.output_types, self._input_dataset.output_shapes,
- shared_name, self._input_dataset.output_classes)
- input_iterator_handle = self._input_iterator.string_handle()
-
- @function.Defun(dtypes.string)
- def _prefetch_fn(handle):
- """Prefetches one element from `input_iterator`."""
- remote_iterator = iterator_ops.Iterator.from_string_handle(
- handle, self._input_iterator.output_types,
- self._input_iterator.output_shapes,
- self._input_iterator.output_classes)
- ret = remote_iterator.get_next()
- return nest.flatten(sparse.serialize_sparse_tensors(ret))
-
- iterator_device = gen_dataset_ops.iterator_get_device(
- self._input_iterator._iterator_resource)
-
- with ops.device(device):
- self._buffering_resource = function_buffering_resource(
- f=_prefetch_fn,
- target_device=iterator_device,
- string_arg=input_iterator_handle,
- buffer_size=buffer_size,
- shared_name=shared_name,
- output_types=nest.flatten(
- sparse.as_dense_types(self._input_dataset.output_types,
- self._input_dataset.output_classes)))
-
- if not self._one_shot:
- reset_op = function_buffering_resource_reset(self._buffering_resource)
- with ops.control_dependencies([reset_op]):
- self._initializer = self._input_iterator.make_initializer(
- self._input_dataset)
-
- def get_next(self, name=None):
- """See `tf.data.Iterator.get_next`."""
- self._get_next_call_count += 1
- if self._get_next_call_count > iterator_ops.GET_NEXT_CALL_WARNING_THRESHOLD:
- warnings.warn(iterator_ops.GET_NEXT_CALL_WARNING_MESSAGE)
-
- flat_ret = gen_dataset_ops.function_buffering_resource_get_next(
- self._buffering_resource,
- output_types=nest.flatten(sparse.as_dense_types(
- self.output_types, self.output_classes)), name=name)
-
- ret = sparse.deserialize_sparse_tensors(
- nest.pack_sequence_as(self.output_types, flat_ret),
- self.output_types, self.output_shapes, self.output_classes)
-
- for tensor, shape in zip(
- nest.flatten(ret), nest.flatten(self.output_shapes)):
- if isinstance(tensor, ops.Tensor):
- tensor.set_shape(shape)
-
- return ret
-
- @property
- def initializer(self):
- if self._one_shot:
- raise NotImplementedError("Can't initialize a one_shot_iterator")
- return self._initializer
-
- @property
- def output_classes(self):
- return self._input_dataset.output_classes
-
- @property
- def output_shapes(self):
- return self._input_dataset.output_shapes
-
- @property
- def output_types(self):
- return self._input_dataset.output_types
-
-
-class _PrefetchToDeviceEagerIterator(iterator_ops.EagerIterator):
- """A replacement for `tf.data.Iterator` that prefetches to another device.
-
- Args:
- input_dataset: The input dataset
- one_shot: If true, we make a one shot iterator that's already initialized.
- device: A fully specified device string where we want to prefetch to
- buffer_size: Size of the prefetching buffer.
- shared_name: (Optional.) If non-empty, the returned iterator will be
- shared under the given name across multiple sessions that share the
- same devices (e.g. when using a remote server).
-
- Returns:
- An Iterator type object.
- """
-
- def __init__(self,
- input_dataset,
- device,
- buffer_size):
- with ops.device("/device:CPU:0"):
- super(_PrefetchToDeviceEagerIterator, self).__init__(input_dataset)
- input_iterator_handle = core_gen_dataset_ops.iterator_to_string_handle(
- self._resource)
-
- self._device = device
-
- @function.Defun(dtypes.string)
- def _prefetch_fn(handle):
- """Prefetches one element from `input_iterator`."""
- remote_iterator = iterator_ops.Iterator.from_string_handle(
- handle, self.output_types, self.output_shapes, self.output_classes)
- ret = remote_iterator.get_next()
- return nest.flatten(sparse.serialize_sparse_tensors(ret))
-
- _prefetch_fn.add_to_graph(None)
-
- with ops.device(device):
- self._buffering_resource = function_buffering_resource(
- f=_prefetch_fn,
- output_types=self._flat_output_types,
- target_device=gen_dataset_ops.iterator_get_device(self._resource),
- string_arg=input_iterator_handle,
- buffer_size=buffer_size,
- shared_name=iterator_ops._generate_shared_name(
- "function_buffer_resource"))
-
- def _next_internal(self):
- """Returns a nested structure of `tf.Tensor`s containing the next element.
- """
- # This runs in sync mode as iterators use an error status to communicate
- # that there is no more data to iterate over.
- # TODO(b/77291417): Fix
- with context.execution_mode(context.SYNC):
- with ops.device(self._device):
- ret = gen_dataset_ops.function_buffering_resource_get_next(
- function_buffer_resource=self._buffering_resource,
- output_types=self._flat_output_types)
- return sparse.deserialize_sparse_tensors(
- nest.pack_sequence_as(self._output_types, ret), self._output_types,
- self._output_shapes, self._output_classes)
-# pylint: enable=protected-access
-
-
-class _PrefetchToDeviceDataset(dataset_ops.UnaryDataset):
- """A `Dataset` whose iterator prefetches elements to another device."""
-
- def __init__(self, input_dataset, device, buffer_size):
- super(_PrefetchToDeviceDataset, self).__init__(input_dataset)
- self._input_dataset = input_dataset
- self._device = device
- self._buffer_size = buffer_size if buffer_size is not None else 1
-
- # The static analysis cannot tell that the eager iterator's superclass has
- # a `next()` method.
- # pylint: disable=non-iterator-returned
- def __iter__(self):
- """Creates an `Iterator` for enumerating the elements of this dataset.
-
- The returned iterator implements the Python iterator protocol and therefore
- can only be used in eager mode.
-
- Returns:
- An `Iterator` over the elements of this dataset.
-
- Raises:
- RuntimeError: If eager execution is enabled.
- """
- if context.executing_eagerly():
- return _PrefetchToDeviceEagerIterator(self._input_dataset, self._device,
- self._buffer_size)
- else:
- raise RuntimeError("dataset.__iter__() is only supported when eager "
- "execution is enabled.")
- # pylint: enable=non-iterator-returned
-
- def make_one_shot_iterator(self):
- if context.executing_eagerly():
- return _PrefetchToDeviceEagerIterator(self._input_dataset, self._device,
- self._buffer_size)
- else:
- return _PrefetchToDeviceIterator(self._input_dataset, one_shot=True,
- device=self._device,
- buffer_size=self._buffer_size)
-
- def make_initializable_iterator(self, shared_name=None):
- return _PrefetchToDeviceIterator(
- self._input_dataset,
- one_shot=False,
- device=self._device,
- buffer_size=self._buffer_size,
- shared_name=shared_name)
-
- def _as_variant_tensor(self):
- # TODO(mrry): Raise this error earlier (e.g. when one of the Dataset
- # transformation methods is called.
- # TODO(mrry): Investigate support for chaining further transformations after
- # the prefetch, including GPU support.
- raise NotImplementedError("`prefetch_to_device()` must be the last "
- "transformation in a dataset pipeline.")
-
- @property
- def output_types(self):
- return self._input_dataset.output_types
-
- @property
- def output_shapes(self):
- return self._input_dataset.output_shapes
-
- @property
- def output_classes(self):
- return self._input_dataset.output_classes
-
-
+@deprecation.deprecated(None,
+ "Use `tf.data.experimental.prefetch_to_device(...)`.")
def prefetch_to_device(device, buffer_size=None):
"""A transformation that prefetches dataset values to the given `device`.
@@ -347,12 +38,10 @@
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
- def _apply_fn(dataset):
- return _PrefetchToDeviceDataset(dataset, device, buffer_size)
-
- return _apply_fn
+ return prefetching_ops.prefetch_to_device(device, buffer_size)
+@deprecation.deprecated(None, "Use `tf.data.experimental.copy_to_device(...)`.")
def copy_to_device(target_device, source_device="/cpu:0"):
"""A transformation that copies dataset elements to the given `target_device`.
@@ -364,353 +53,4 @@
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
-
- def _apply_fn(dataset):
- return _CopyToDeviceDataset(
- dataset, target_device=target_device, source_device=source_device)
-
- return _apply_fn
-
-
-# TODO(rohanj): Use the _input_hostmem attr on the RemoteCall ops to indicate
-# all inputs to the Op are in host memory, thereby avoiding some unnecessary
-# Sends and Recvs.
-class _CopyToDeviceDataset(dataset_ops.UnaryDataset):
- """A `Dataset` that copies elements to another device."""
-
- def __init__(self, input_dataset, target_device, source_device="/cpu:0"):
- """Constructs a _CopyToDeviceDataset.
-
- Args:
- input_dataset: `Dataset` to be copied
- target_device: The name of the device to which elements would be copied.
- source_device: Device where input_dataset would be placed.
- """
- super(_CopyToDeviceDataset, self).__init__(input_dataset)
- self._input_dataset = input_dataset
- self._target_device = target_device
- spec = framework_device.DeviceSpec().from_string(self._target_device)
- self._is_gpu_target = (spec.device_type == "GPU")
- 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.
-
- Returns:
- A `string` tensor that encapsulates the iterator created.
- """
- # pylint: disable=protected-access
- ds_variant = self._input_dataset._as_variant_tensor()
- resource = core_gen_dataset_ops.anonymous_iterator(
- output_types=self._flat_output_types,
- output_shapes=self._flat_output_shapes)
- with ops.control_dependencies(
- [core_gen_dataset_ops.make_iterator(ds_variant, resource)]):
- return core_gen_dataset_ops.iterator_to_string_handle(resource)
-
- @function.Defun()
- def _remote_init_func():
- return functional_ops.remote_call(
- target=self._source_device,
- args=_init_func.captured_inputs,
- Tout=[dtypes.string],
- f=_init_func)
-
- self._init_func = _remote_init_func
- self._init_captured_args = _remote_init_func.captured_inputs
-
- @function.Defun(dtypes.string)
- def _next_func(string_handle):
- """Calls get_next for created iterator.
-
- Args:
- string_handle: An iterator string handle created by _init_func
- Returns:
- The elements generated from `input_dataset`
- """
- with ops.device(self._source_device_string):
- 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))
-
- @function.Defun(dtypes.string)
- def _remote_next_func(string_handle):
- return functional_ops.remote_call(
- target=self._source_device,
- args=[string_handle] + _next_func.captured_inputs,
- Tout=self._flat_output_types,
- f=_next_func)
-
- self._next_func = _remote_next_func
- self._next_captured_args = _remote_next_func.captured_inputs
-
- @function.Defun(dtypes.string)
- def _finalize_func(string_handle):
- """Destroys the iterator resource created.
-
- Args:
- string_handle: An iterator string handle created by _init_func
- Returns:
- Tensor constant 0
- """
- iterator_resource = core_gen_dataset_ops.iterator_from_string_handle_v2(
- string_handle,
- output_types=self._flat_output_types,
- output_shapes=self._flat_output_shapes)
- with ops.control_dependencies([
- resource_variable_ops.destroy_resource_op(
- iterator_resource, ignore_lookup_error=True)]):
- return array_ops.constant(0, dtypes.int64)
-
- @function.Defun(dtypes.string)
- def _remote_finalize_func(string_handle):
- return functional_ops.remote_call(
- target=self._source_device,
- args=[string_handle] + _finalize_func.captured_inputs,
- Tout=[dtypes.int64],
- f=_finalize_func)
-
- self._finalize_func = _remote_finalize_func
- self._finalize_captured_args = _remote_finalize_func.captured_inputs
-
- g = ops.get_default_graph()
- _remote_init_func.add_to_graph(g)
- _remote_next_func.add_to_graph(g)
- _remote_finalize_func.add_to_graph(g)
- # pylint: enable=protected-scope
-
- # The one_shot_iterator implementation needs a 0 arg _make_dataset function
- # that thereby captures all the inputs required to create the dataset. Since
- # there are strings that are inputs to the GeneratorDataset which can't be
- # placed on a GPU, this fails for the GPU case. Therefore, disabling it for
- # GPU
- def make_one_shot_iterator(self):
- if self._is_gpu_target:
- raise ValueError("Cannot create a one shot iterator when using "
- "`tf.contrib.data.copy_to_device()` on GPU. Please use "
- "`Dataset.make_initializable_iterator()` instead.")
- else:
- return super(_CopyToDeviceDataset, self).make_one_shot_iterator()
-
- def _as_variant_tensor(self):
- with ops.device(self._target_device):
- return core_gen_dataset_ops.generator_dataset(
- self._init_captured_args,
- self._next_captured_args,
- self._finalize_captured_args,
- 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)
-
- @property
- def output_types(self):
- return self._input_dataset.output_types
-
- @property
- def output_shapes(self):
- return self._input_dataset.output_shapes
-
- @property
- def output_classes(self):
- return self._input_dataset.output_classes
-
-
-class _PerDeviceGenerator(dataset_ops.Dataset):
- """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):
- 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))
-
- multi_device_iterator_string_handle = (
- gen_dataset_ops.multi_device_iterator_to_string_handle(
- multi_device_iterator_resource))
-
- @function.Defun()
- def _init_func():
- return multi_device_iterator_string_handle
-
- @function.Defun()
- def _remote_init_func():
- return functional_ops.remote_call(
- target=source_device,
- args=_init_func.captured_inputs,
- Tout=[dtypes.string],
- f=_init_func)
-
- self._init_func = _remote_init_func
- self._init_captured_args = _remote_init_func.captured_inputs
-
- @function.Defun(dtypes.string)
- def _next_func(string_handle):
- 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))
- 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)
-
- @function.Defun(dtypes.string)
- def _remote_next_func(string_handle):
- return functional_ops.remote_call(
- target=source_device,
- args=[string_handle] + _next_func.captured_inputs,
- Tout=self._flat_output_types,
- f=_next_func)
-
- self._next_func = _remote_next_func
- self._next_captured_args = _remote_next_func.captured_inputs
-
- @function.Defun(dtypes.string)
- def _finalize_func(unused_string_handle):
- return array_ops.constant(0, dtypes.int64)
-
- @function.Defun(dtypes.string)
- def _remote_finalize_func(string_handle):
- return functional_ops.remote_call(
- target=source_device,
- args=[string_handle] + _finalize_func.captured_inputs,
- Tout=[dtypes.int64],
- f=_finalize_func)
-
- self._finalize_func = _remote_finalize_func
- self._finalize_captured_args = _remote_finalize_func.captured_inputs
-
- def _as_variant_tensor(self):
- with ops.device(self._target_device):
- return core_gen_dataset_ops.generator_dataset(
- self._init_captured_args,
- self._next_captured_args,
- self._finalize_captured_args,
- 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)
-
- 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
-
-
-class MultiDeviceIterator(object):
- """An iterator over multiple devices."""
-
- def __init__(self,
- dataset,
- devices,
- max_buffer_size=1,
- prefetch_buffer_size=1,
- source_device="/cpu:0"):
- """Constructs a MultiDeviceIterator.
-
- Args:
- dataset: The input dataset to be iterated over.
- devices: The list of devices to fetch data to.
- max_buffer_size: Maximum size of the host side per device buffer to keep.
- prefetch_buffer_size: if > 1, then we setup a buffer on each device
- to prefetch into.
- source_device: The host device to place the `dataset` on.
- """
- self._dataset = dataset
- self._devices = devices
- 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 = (
- gen_dataset_ops.multi_device_iterator(
- devices=self._devices,
- shared_name="",
- container="",
- output_types=self._flat_output_types,
- output_shapes=self._flat_output_shapes))
-
- # The incarnation ID is used to ensure consistency between the per-device
- # iterators and the multi-device iterator.
- self._incarnation_id = gen_dataset_ops.multi_device_iterator_init(
- self._dataset._as_variant_tensor(), # pylint: disable=protected-access
- self._multi_device_iterator_resource,
- max_buffer_size=max_buffer_size)
-
- # TODO(rohanj): Explore the possibility of the MultiDeviceIterator to
- # initialize the device side of the pipeline. This would allow the
- # MultiDeviceIterator to choose, for example, to move some transformations
- # into the device side from its input. It might be useful in rewriting.
- # Create the per device iterators.
- self._device_iterators = []
- i = 0
- for device in 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)
- if prefetch_buffer_size > 0:
- ds = ds.prefetch(prefetch_buffer_size)
- with ops.device(device):
- self._device_iterators.append(ds.make_initializable_iterator())
- i += 1
-
- device_iterator_initializers = [
- iterator.initializer for iterator in self._device_iterators
- ]
- self._initializer = control_flow_ops.group(*device_iterator_initializers)
-
- def get_next(self):
- result = []
- i = 0
- for device in self._devices:
- with ops.device(device):
- result.append(self._device_iterators[i].get_next())
- i += 1
- return result
-
- @property
- def initializer(self):
- return self._initializer
+ return prefetching_ops.copy_to_device(target_device, source_device)
diff --git a/tensorflow/contrib/data/python/ops/random_ops.py b/tensorflow/contrib/data/python/ops/random_ops.py
index 344a076..2c95125 100644
--- a/tensorflow/contrib/data/python/ops/random_ops.py
+++ b/tensorflow/contrib/data/python/ops/random_ops.py
@@ -17,36 +17,14 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.python.data.ops import dataset_ops
-from tensorflow.python.data.util import random_seed
-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_dataset_ops
+from tensorflow.python.data.experimental.ops import random_ops
+from tensorflow.python.util import deprecation
-class RandomDataset(dataset_ops.DatasetSource):
+class RandomDataset(random_ops.RandomDataset):
"""A `Dataset` of pseudorandom values."""
+ @deprecation.deprecated(
+ None, "Use `tf.data.experimental.RandomDataset(...)`.")
def __init__(self, seed=None):
- """A `Dataset` of pseudorandom values."""
- super(RandomDataset, self).__init__()
- self._seed, self._seed2 = random_seed.get_seed(seed)
-
- def _as_variant_tensor(self):
- return gen_dataset_ops.random_dataset(
- seed=self._seed,
- seed2=self._seed2,
- **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
+ super(RandomDataset, self).__init__(seed)
diff --git a/tensorflow/contrib/data/python/ops/readers.py b/tensorflow/contrib/data/python/ops/readers.py
index d9d06e2..4601376 100644
--- a/tensorflow/contrib/data/python/ops/readers.py
+++ b/tensorflow/contrib/data/python/ops/readers.py
@@ -17,295 +17,20 @@
from __future__ import division
from __future__ import print_function
-import collections
-import csv
-
-import numpy as np
-
-from tensorflow.contrib.data.python.ops import batching
-from tensorflow.contrib.data.python.ops import gen_dataset_ops as contrib_gen_dataset_ops
-from tensorflow.contrib.data.python.ops import interleave_ops
-from tensorflow.contrib.data.python.ops import optimization
-from tensorflow.contrib.data.python.ops import parsing_ops
-from tensorflow.contrib.data.python.ops import shuffle_ops
+from tensorflow.python.data.experimental.ops import optimization
+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 convert
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 ops
from tensorflow.python.framework import tensor_shape
-from tensorflow.python.lib.io import file_io
-from tensorflow.python.ops import gen_dataset_ops
-from tensorflow.python.platform import gfile
+from tensorflow.python.ops import gen_experimental_dataset_ops
from tensorflow.python.util import deprecation
-_ACCEPTABLE_CSV_TYPES = (dtypes.float32, dtypes.float64, dtypes.int32,
- dtypes.int64, dtypes.string)
-
-def _is_valid_int32(str_val):
- try:
- # Checks equality to prevent int32 overflow
- return dtypes.int32.as_numpy_dtype(str_val) == dtypes.int64.as_numpy_dtype(
- str_val)
- except (ValueError, OverflowError):
- return False
-
-
-def _is_valid_int64(str_val):
- try:
- dtypes.int64.as_numpy_dtype(str_val)
- return True
- except (ValueError, OverflowError):
- return False
-
-
-def _is_valid_float(str_val, float_dtype):
- try:
- return float_dtype.as_numpy_dtype(str_val) < np.inf
- except ValueError:
- return False
-
-
-def _infer_type(str_val, na_value, prev_type):
- """Given a string, infers its tensor type.
-
- Infers the type of a value by picking the least 'permissive' type possible,
- while still allowing the previous type inference for this column to be valid.
-
- Args:
- str_val: String value to infer the type of.
- na_value: Additional string to recognize as a NA/NaN CSV value.
- prev_type: Type previously inferred based on values of this column that
- we've seen up till now.
- Returns:
- Inferred dtype.
- """
- if str_val in ("", na_value):
- # If the field is null, it gives no extra information about its type
- return prev_type
-
- type_list = [
- dtypes.int32, dtypes.int64, dtypes.float32, dtypes.float64, dtypes.string
- ] # list of types to try, ordered from least permissive to most
-
- type_functions = [
- _is_valid_int32,
- _is_valid_int64,
- lambda str_val: _is_valid_float(str_val, dtypes.float32),
- lambda str_val: _is_valid_float(str_val, dtypes.float64),
- lambda str_val: True,
- ] # Corresponding list of validation functions
-
- for i in range(len(type_list)):
- validation_fn = type_functions[i]
- if validation_fn(str_val) and (prev_type is None or
- prev_type in type_list[:i + 1]):
- return type_list[i]
-
-
-def _next_csv_row(filenames, num_cols, field_delim, use_quote_delim, header):
- """Generator that yields rows of CSV file(s) in order."""
- for fn in filenames:
- with file_io.FileIO(fn, "r") as f:
- rdr = csv.reader(
- f,
- delimiter=field_delim,
- quoting=csv.QUOTE_MINIMAL if use_quote_delim else csv.QUOTE_NONE)
- if header:
- next(rdr) # Skip header lines
-
- for csv_row in rdr:
- if len(csv_row) != num_cols:
- raise ValueError(
- "Problem inferring types: CSV row has different number of fields "
- "than expected.")
- yield csv_row
-
-
-def _infer_column_defaults(filenames, num_cols, field_delim, use_quote_delim,
- na_value, header, num_rows_for_inference,
- select_columns):
- """Infers column types from the first N valid CSV records of files."""
- if select_columns is None:
- select_columns = range(num_cols)
- inferred_types = [None] * len(select_columns)
-
- for i, csv_row in enumerate(
- _next_csv_row(filenames, num_cols, field_delim, use_quote_delim, header)):
- if num_rows_for_inference is not None and i >= num_rows_for_inference:
- break
-
- for j, col_index in enumerate(select_columns):
- inferred_types[j] = _infer_type(csv_row[col_index], na_value,
- inferred_types[j])
-
- # Replace None's with a default type
- inferred_types = [t or dtypes.string for t in inferred_types]
- # Default to 0 or '' for null values
- return [
- constant_op.constant([0 if t is not dtypes.string else ""], dtype=t)
- for t in inferred_types
- ]
-
-
-def _infer_column_names(filenames, field_delim, use_quote_delim):
- """Infers column names from first rows of files."""
- csv_kwargs = {
- "delimiter": field_delim,
- "quoting": csv.QUOTE_MINIMAL if use_quote_delim else csv.QUOTE_NONE
- }
- with file_io.FileIO(filenames[0], "r") as f:
- try:
- column_names = next(csv.reader(f, **csv_kwargs))
- except StopIteration:
- raise ValueError(("Received StopIteration when reading the header line "
- "of %s. Empty file?") % filenames[0])
-
- for name in filenames[1:]:
- with file_io.FileIO(name, "r") as f:
- try:
- if next(csv.reader(f, **csv_kwargs)) != column_names:
- raise ValueError(
- "Files have different column names in the header row.")
- except StopIteration:
- raise ValueError(("Received StopIteration when reading the header line "
- "of %s. Empty file?") % filenames[0])
- return column_names
-
-
-def _get_sorted_col_indices(select_columns, column_names):
- """Transforms select_columns argument into sorted column indices."""
- names_to_indices = {n: i for i, n in enumerate(column_names)}
- num_cols = len(column_names)
- for i, v in enumerate(select_columns):
- if isinstance(v, int):
- if v < 0 or v >= num_cols:
- raise ValueError(
- "Column index %d specified in select_columns out of valid range." %
- v)
- continue
- if v not in names_to_indices:
- raise ValueError(
- "Value '%s' specified in select_columns not a valid column index or "
- "name." % v)
- select_columns[i] = names_to_indices[v]
-
- # Sort and ensure there are no duplicates
- result = sorted(set(select_columns))
- if len(result) != len(select_columns):
- raise ValueError("select_columns contains duplicate columns")
- return result
-
-
-def _maybe_shuffle_and_repeat(
- dataset, num_epochs, shuffle, shuffle_buffer_size, shuffle_seed):
- """Optionally shuffle and repeat dataset, as requested."""
- if num_epochs != 1 and shuffle:
- # Use shuffle_and_repeat for perf
- return dataset.apply(
- shuffle_ops.shuffle_and_repeat(shuffle_buffer_size, num_epochs,
- shuffle_seed))
- elif shuffle:
- return dataset.shuffle(shuffle_buffer_size, shuffle_seed)
- elif num_epochs != 1:
- return dataset.repeat(num_epochs)
- return dataset
-
-
-def make_tf_record_dataset(file_pattern,
- batch_size,
- parser_fn=None,
- num_epochs=None,
- shuffle=True,
- shuffle_buffer_size=None,
- shuffle_seed=None,
- prefetch_buffer_size=optimization.AUTOTUNE,
- num_parallel_reads=None,
- num_parallel_parser_calls=None,
- drop_final_batch=False):
- """Reads and optionally parses TFRecord files into a dataset.
-
- Provides common functionality such as batching, optional parsing, shuffling,
- and performant defaults.
-
- Args:
- file_pattern: List of files or patterns of TFRecord file paths.
- See `tf.gfile.Glob` for pattern rules.
- batch_size: An int representing the number of records to combine
- in a single batch.
- parser_fn: (Optional.) A function accepting string input to parse
- and process the record contents. This function must map records
- to components of a fixed shape, so they may be batched. By
- default, uses the record contents unmodified.
- num_epochs: (Optional.) An int specifying the number of times this
- dataset is repeated. If None (the default), cycles through the
- dataset forever.
- shuffle: (Optional.) A bool that indicates whether the input
- should be shuffled. Defaults to `True`.
- shuffle_buffer_size: (Optional.) Buffer size to use for
- shuffling. A large buffer size ensures better shuffling, but
- increases memory usage and startup time.
- shuffle_seed: (Optional.) Randomization seed to use for shuffling.
- prefetch_buffer_size: (Optional.) An int specifying the number of
- feature batches to prefetch for performance improvement.
- Defaults to auto-tune. Set to 0 to disable prefetching.
- num_parallel_reads: (Optional.) Number of threads used to read
- records from files. By default or if set to a value >1, the
- results will be interleaved.
- num_parallel_parser_calls: (Optional.) Number of parallel
- records to parse in parallel. Defaults to an automatic selection.
- drop_final_batch: (Optional.) Whether the last batch should be
- dropped in case its size is smaller than `batch_size`; the
- default behavior is not to drop the smaller batch.
-
- Returns:
- A dataset, where each element matches the output of `parser_fn`
- except it will have an additional leading `batch-size` dimension,
- or a `batch_size`-length 1-D tensor of strings if `parser_fn` is
- unspecified.
- """
- files = dataset_ops.Dataset.list_files(
- file_pattern, shuffle=shuffle, seed=shuffle_seed)
-
- if num_parallel_reads is None:
- # Note: We considered auto-tuning this value, but there is a concern
- # that this affects the mixing of records from different files, which
- # could affect training convergence/accuracy, so we are defaulting to
- # a constant for now.
- num_parallel_reads = 24
- dataset = core_readers.TFRecordDataset(
- files, num_parallel_reads=num_parallel_reads)
-
- if shuffle_buffer_size is None:
- # TODO(josh11b): Auto-tune this value when not specified
- shuffle_buffer_size = 10000
- dataset = _maybe_shuffle_and_repeat(
- dataset, num_epochs, shuffle, shuffle_buffer_size, shuffle_seed)
-
- # NOTE(mrry): We set `drop_final_batch=True` when `num_epochs is None` to
- # improve the shape inference, because it makes the batch dimension static.
- # It is safe to do this because in that case we are repeating the input
- # indefinitely, and all batches will be full-sized.
- drop_final_batch = drop_final_batch or num_epochs is None
-
- if parser_fn is None:
- dataset = dataset.batch(batch_size, drop_remainder=drop_final_batch)
- else:
- # TODO(josh11b): if num_parallel_parser_calls is None, use some function
- # of num cores instead of map_and_batch's default behavior of one batch.
- dataset = dataset.apply(batching.map_and_batch(
- parser_fn, batch_size, num_parallel_calls=num_parallel_parser_calls,
- drop_remainder=drop_final_batch))
-
- if prefetch_buffer_size == 0:
- return dataset
- else:
- return dataset.prefetch(buffer_size=prefetch_buffer_size)
-
-
+@deprecation.deprecated(None,
+ "Use `tf.data.experimental.make_csv_dataset(...)`.")
def make_csv_dataset(
file_pattern,
batch_size,
@@ -387,7 +112,6 @@
prefetch_buffer_size: An int specifying the number of feature
batches to prefetch for performance improvement. Recommended value is the
number of batches consumed per training step. Defaults to auto-tune.
-
num_parallel_reads: Number of threads used to read CSV records from files.
If >1, the results will be interleaved.
sloppy: If `True`, reading performance will be improved at
@@ -411,106 +135,18 @@
Raises:
ValueError: If any of the arguments is malformed.
"""
- # Create dataset of all matching filenames
- filenames = _get_file_names(file_pattern, False)
- dataset = dataset_ops.Dataset.from_tensor_slices(filenames)
- if shuffle:
- dataset = dataset.shuffle(len(filenames), shuffle_seed)
-
- # Clean arguments; figure out column names and defaults
-
- if column_names is None:
- if not header:
- raise ValueError("Cannot infer column names without a header line.")
- # If column names are not provided, infer from the header lines
- column_names = _infer_column_names(filenames, field_delim, use_quote_delim)
- if len(column_names) != len(set(column_names)):
- raise ValueError("Cannot have duplicate column names.")
-
- if select_columns is not None:
- select_columns = _get_sorted_col_indices(select_columns, column_names)
-
- if column_defaults is not None:
- column_defaults = [
- constant_op.constant([], dtype=x) if x in _ACCEPTABLE_CSV_TYPES else x
- for x in column_defaults
- ]
- else:
- # If column defaults are not provided, infer from records at graph
- # construction time
- column_defaults = _infer_column_defaults(
- filenames, len(column_names), field_delim, use_quote_delim, na_value,
- header, num_rows_for_inference, select_columns)
-
- if select_columns is not None and len(column_defaults) != len(select_columns):
- raise ValueError(
- "If specified, column_defaults and select_columns must have same "
- "length."
- )
- if select_columns is not None and len(column_names) > len(select_columns):
- # Pick the relevant subset of column names
- column_names = [column_names[i] for i in select_columns]
-
- if label_name is not None and label_name not in column_names:
- raise ValueError("`label_name` provided must be one of the columns.")
-
- def filename_to_dataset(filename):
- return CsvDataset(
- filename,
- record_defaults=column_defaults,
- field_delim=field_delim,
- use_quote_delim=use_quote_delim,
- na_value=na_value,
- select_cols=select_columns,
- header=header,
- compression_type=compression_type,
- )
-
- def map_fn(*columns):
- """Organizes columns into a features dictionary.
-
- Args:
- *columns: list of `Tensor`s corresponding to one csv record.
- Returns:
- An OrderedDict of feature names to values for that particular record. If
- label_name is provided, extracts the label feature to be returned as the
- second element of the tuple.
- """
- features = collections.OrderedDict(zip(column_names, columns))
- if label_name is not None:
- label = features.pop(label_name)
- return features, label
- return features
-
- # Read files sequentially (if num_parallel_reads=1) or in parallel
- dataset = dataset.apply(
- interleave_ops.parallel_interleave(
- filename_to_dataset, cycle_length=num_parallel_reads, sloppy=sloppy))
-
- dataset = _maybe_shuffle_and_repeat(
- dataset, num_epochs, shuffle, shuffle_buffer_size, shuffle_seed)
-
- # Apply batch before map for perf, because map has high overhead relative
- # to the size of the computation in each map.
- # NOTE(mrry): We set `drop_remainder=True` when `num_epochs is None` to
- # improve the shape inference, because it makes the batch dimension static.
- # It is safe to do this because in that case we are repeating the input
- # indefinitely, and all batches will be full-sized.
- dataset = dataset.batch(batch_size=batch_size,
- drop_remainder=num_epochs is None)
- dataset = dataset_ops.MapDataset(
- dataset, map_fn, use_inter_op_parallelism=False)
- dataset = dataset.prefetch(prefetch_buffer_size)
-
- return dataset
+ return readers.make_csv_dataset(
+ file_pattern, batch_size, column_names, column_defaults, label_name,
+ select_columns, field_delim, use_quote_delim, na_value, header,
+ num_epochs, shuffle, shuffle_buffer_size, shuffle_seed,
+ prefetch_buffer_size, num_parallel_reads, sloppy, num_rows_for_inference,
+ compression_type)
-_DEFAULT_READER_BUFFER_SIZE_BYTES = 4 * 1024 * 1024 # 4 MB
-
-
-class CsvDataset(dataset_ops.DatasetSource):
+class CsvDataset(readers.CsvDataset):
"""A Dataset comprising lines from one or more CSV files."""
+ @deprecation.deprecated(None, "Use `tf.data.experimental.CsvDataset(...)`.")
def __init__(self,
filenames,
record_defaults,
@@ -521,140 +157,13 @@
use_quote_delim=True,
na_value="",
select_cols=None):
- """Creates a `CsvDataset` by reading and decoding CSV files.
-
- The elements of this dataset correspond to records from the file(s).
- RFC 4180 format is expected for CSV files
- (https://tools.ietf.org/html/rfc4180)
- Note that we allow leading and trailing spaces with int or float field.
+ super(CsvDataset, self).__init__(
+ filenames, record_defaults, compression_type, buffer_size, header,
+ field_delim, use_quote_delim, na_value, select_cols)
- For example, suppose we have a file 'my_file0.csv' with four CSV columns of
- different data types:
- ```
- abcdefg,4.28E10,5.55E6,12
- hijklmn,-5.3E14,,2
- ```
-
- We can construct a CsvDataset from it as follows:
- ```python
- dataset = tf.contrib.data.CsvDataset(
- "my_file*.csv",
- [tf.float32, # Required field, use dtype or empty tensor
- tf.constant([0.0], dtype=tf.float32), # Optional field, default to 0.0
- tf.int32, # Required field, use dtype or empty tensor
- ],
- select_cols=[1,2,3] # Only parse last three columns
- )
- ```
-
- The expected output of its iterations is:
- ```python
- next_element = dataset.make_one_shot_iterator().get_next()
- with tf.Session() as sess:
- while True:
- try:
- print(sess.run(next_element))
- except tf.errors.OutOfRangeError:
- break
-
- >> (4.28e10, 5.55e6, 12)
- >> (-5.3e14, 0.0, 2)
- ```
-
- Args:
- filenames: A `tf.string` tensor containing one or more filenames.
- record_defaults: A list of default values for the CSV fields. Each item in
- the list is either a valid CSV `DType` (float32, float64, int32, int64,
- string), or a `Tensor` object with one of the above types. One per
- column of CSV data, with either a scalar `Tensor` default value for the
- column if it is optional, or `DType` or empty `Tensor` if required. If
- both this and `select_columns` are specified, these must have the same
- lengths, and `column_defaults` is assumed to be sorted in order of
- increasing column index.
- compression_type: (Optional.) A `tf.string` scalar evaluating to one of
- `""` (no compression), `"ZLIB"`, or `"GZIP"`. Defaults to no
- compression.
- buffer_size: (Optional.) A `tf.int64` scalar denoting the number of bytes
- to buffer while reading files. Defaults to 4MB.
- header: (Optional.) A `tf.bool` scalar indicating whether the CSV file(s)
- have header line(s) that should be skipped when parsing. Defaults to
- `False`.
- field_delim: (Optional.) A `tf.string` scalar containing the delimiter
- character that separates fields in a record. Defaults to `","`.
- use_quote_delim: (Optional.) A `tf.bool` scalar. If `False`, treats
- double quotation marks as regular characters inside of string fields
- (ignoring RFC 4180, Section 2, Bullet 5). Defaults to `True`.
- na_value: (Optional.) A `tf.string` scalar indicating a value that will
- be treated as NA/NaN.
- select_cols: (Optional.) A sorted list of column indices to select from
- the input data. If specified, only this subset of columns will be
- parsed. Defaults to parsing all columns.
- """
- super(CsvDataset, self).__init__()
- self._filenames = ops.convert_to_tensor(
- filenames, dtype=dtypes.string, name="filenames")
- self._compression_type = convert.optional_param_to_tensor(
- "compression_type",
- compression_type,
- argument_default="",
- argument_dtype=dtypes.string)
- record_defaults = [
- constant_op.constant([], dtype=x) if x in _ACCEPTABLE_CSV_TYPES else x
- for x in record_defaults
- ]
- self._record_defaults = ops.convert_n_to_tensor(
- record_defaults, name="record_defaults")
- self._buffer_size = convert.optional_param_to_tensor(
- "buffer_size", buffer_size, _DEFAULT_READER_BUFFER_SIZE_BYTES)
- self._header = ops.convert_to_tensor(
- header, dtype=dtypes.bool, name="header")
- self._field_delim = ops.convert_to_tensor(
- field_delim, dtype=dtypes.string, name="field_delim")
- self._use_quote_delim = ops.convert_to_tensor(
- use_quote_delim, dtype=dtypes.bool, name="use_quote_delim")
- self._na_value = ops.convert_to_tensor(
- na_value, dtype=dtypes.string, name="na_value")
- self._select_cols = convert.optional_param_to_tensor(
- "select_cols",
- select_cols,
- 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)))
-
- def _as_variant_tensor(self):
- # Constructs graph node for the dataset op.
- return contrib_gen_dataset_ops.csv_dataset(
- filenames=self._filenames,
- record_defaults=self._record_defaults,
- buffer_size=self._buffer_size,
- header=self._header,
- output_shapes=self._output_shapes,
- field_delim=self._field_delim,
- use_quote_delim=self._use_quote_delim,
- na_value=self._na_value,
- select_cols=self._select_cols,
- compression_type=self._compression_type,
- )
-
- @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
-
-
+@deprecation.deprecated(
+ None, "Use `tf.data.experimental.make_batched_features_dataset(...)`.")
def make_batched_features_dataset(file_pattern,
batch_size,
features,
@@ -759,57 +268,15 @@
Raises:
ValueError: If `label_key` is not one of the `features` keys.
"""
- # Create dataset of all matching filenames
- filenames = _get_file_names(file_pattern, False)
- dataset = dataset_ops.Dataset.from_tensor_slices(filenames)
- if shuffle:
- dataset = dataset.shuffle(len(filenames), shuffle_seed)
-
- # Read `Example` records from files as tensor objects.
- if reader_args is None:
- reader_args = []
-
- # Read files sequentially (if reader_num_threads=1) or in parallel
- dataset = dataset.apply(
- interleave_ops.parallel_interleave(
- lambda filename: reader(filename, *reader_args),
- cycle_length=reader_num_threads,
- sloppy=sloppy_ordering))
-
- # Extract values if the `Example` tensors are stored as key-value tuples.
- if dataset.output_types == (dtypes.string, dtypes.string):
- dataset = dataset_ops.MapDataset(
- dataset, lambda _, v: v, use_inter_op_parallelism=False)
-
- # Apply dataset repeat and shuffle transformations.
- dataset = _maybe_shuffle_and_repeat(
- dataset, num_epochs, shuffle, shuffle_buffer_size, shuffle_seed)
-
- # NOTE(mrry): We set `drop_remainder=True` when `num_epochs is None` to
- # improve the shape inference, because it makes the batch dimension static.
- # It is safe to do this because in that case we are repeating the input
- # indefinitely, and all batches will be full-sized.
- dataset = dataset.batch(
- batch_size, drop_remainder=drop_final_batch or num_epochs is None)
-
- # Parse `Example` tensors to a dictionary of `Feature` tensors.
- dataset = dataset.apply(
- parsing_ops.parse_example_dataset(
- features, num_parallel_calls=parser_num_threads))
-
- if label_key:
- if label_key not in features:
- raise ValueError(
- "The `label_key` provided (%r) must be one of the `features` keys." %
- label_key)
- dataset = dataset.map(lambda x: (x, x.pop(label_key)))
-
- dataset = dataset.prefetch(prefetch_buffer_size)
- return dataset
+ return readers.make_batched_features_dataset(
+ file_pattern, batch_size, features, reader, label_key, reader_args,
+ num_epochs, shuffle, shuffle_buffer_size, shuffle_seed,
+ prefetch_buffer_size, reader_num_threads, parser_num_threads,
+ sloppy_ordering, drop_final_batch)
-@deprecation.deprecated(None,
- "Use `tf.contrib.data.make_batched_features_dataset`")
+@deprecation.deprecated(
+ None, "Use `tf.data.experimental.make_batched_features_dataset(...)`")
def read_batch_features(file_pattern,
batch_size,
features,
@@ -879,7 +346,7 @@
Returns:
A dict from keys in features to `Tensor` or `SparseTensor` objects.
"""
- dataset = make_batched_features_dataset(
+ dataset = readers.make_batched_features_dataset(
file_pattern,
batch_size,
features,
@@ -893,96 +360,13 @@
return outputs
-def _get_file_names(file_pattern, shuffle):
- """Parse list of file names from pattern, optionally shuffled.
-
- Args:
- file_pattern: File glob pattern, or list of glob patterns.
- shuffle: Whether to shuffle the order of file names.
-
- Returns:
- List of file names matching `file_pattern`.
-
- Raises:
- ValueError: If `file_pattern` is empty, or pattern matches no files.
- """
- if isinstance(file_pattern, list):
- if not file_pattern:
- raise ValueError("File pattern is empty.")
- file_names = []
- for entry in file_pattern:
- file_names.extend(gfile.Glob(entry))
- else:
- file_names = list(gfile.Glob(file_pattern))
-
- if not file_names:
- raise ValueError("No files match %s." % file_pattern)
-
- # Sort files so it will be deterministic for unit tests.
- if not shuffle:
- file_names = sorted(file_names)
- return file_names
-
-
-class SqlDataset(dataset_ops.DatasetSource):
+class SqlDataset(readers.SqlDataset):
"""A `Dataset` consisting of the results from a SQL query."""
+ @deprecation.deprecated(None, "Use `tf.data.experimental.SqlDataset(...)`.")
def __init__(self, driver_name, data_source_name, query, output_types):
- """Creates a `SqlDataset`.
-
- `SqlDataset` allows a user to read data from the result set of a SQL query.
- For example:
-
- ```python
- dataset = tf.contrib.data.SqlDataset("sqlite", "/foo/bar.sqlite3",
- "SELECT name, age FROM people",
- (tf.string, tf.int32))
- iterator = dataset.make_one_shot_iterator()
- next_element = iterator.get_next()
- # Prints the rows of the result set of the above query.
- while True:
- try:
- print(sess.run(next_element))
- except tf.errors.OutOfRangeError:
- break
- ```
-
- Args:
- driver_name: A 0-D `tf.string` tensor containing the database type.
- Currently, the only supported value is 'sqlite'.
- data_source_name: A 0-D `tf.string` tensor containing a connection string
- to connect to the database.
- query: A 0-D `tf.string` tensor containing the SQL query to execute.
- output_types: A tuple of `tf.DType` objects representing the types of the
- columns returned by `query`.
- """
- super(SqlDataset, self).__init__()
- self._driver_name = ops.convert_to_tensor(
- driver_name, dtype=dtypes.string, name="driver_name")
- self._data_source_name = ops.convert_to_tensor(
- 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
-
- def _as_variant_tensor(self):
- return gen_dataset_ops.sql_dataset(self._driver_name,
- self._data_source_name, self._query,
- 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
+ super(SqlDataset, self).__init__(
+ driver_name, data_source_name, query, output_types)
class LMDBDataset(dataset_ops.DatasetSource):
@@ -1013,7 +397,7 @@
filenames, dtype=dtypes.string, name="filenames")
def _as_variant_tensor(self):
- return contrib_gen_dataset_ops.lmdb_dataset(
+ return gen_experimental_dataset_ops.experimental_lmdb_dataset(
self._filenames,
output_types=nest.flatten(self.output_types),
output_shapes=nest.flatten(self.output_shapes))
diff --git a/tensorflow/contrib/data/python/ops/resampling.py b/tensorflow/contrib/data/python/ops/resampling.py
index 75642f1..29d7752 100644
--- a/tensorflow/contrib/data/python/ops/resampling.py
+++ b/tensorflow/contrib/data/python/ops/resampling.py
@@ -17,22 +17,12 @@
from __future__ import division
from __future__ import print_function
-import numpy as np
-
-from tensorflow.contrib.data.python.ops import batching
-from tensorflow.contrib.data.python.ops import interleave_ops
-from tensorflow.contrib.data.python.ops import scan_ops
-from tensorflow.python.data.ops import dataset_ops
-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_ops
-from tensorflow.python.ops import logging_ops
-from tensorflow.python.ops import math_ops
-from tensorflow.python.ops import random_ops
+from tensorflow.python.data.experimental.ops import resampling
+from tensorflow.python.util import deprecation
+@deprecation.deprecated(None,
+ "Use `tf.data.experimental.rejection_resample(...)`.")
def rejection_resample(class_func, target_dist, initial_dist=None, seed=None):
"""A transformation that resamples a dataset to achieve a target distribution.
@@ -52,243 +42,5 @@
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
- def _apply_fn(dataset):
- """Function from `Dataset` to `Dataset` that applies the transformation."""
- target_dist_t = ops.convert_to_tensor(target_dist, name="target_dist")
- class_values_ds = dataset.map(class_func)
-
- # Get initial distribution.
- if initial_dist is not None:
- initial_dist_t = ops.convert_to_tensor(initial_dist, name="initial_dist")
- acceptance_dist, prob_of_original = (
- _calculate_acceptance_probs_with_mixing(initial_dist_t,
- target_dist_t))
- initial_dist_ds = dataset_ops.Dataset.from_tensors(
- initial_dist_t).repeat()
- acceptance_dist_ds = dataset_ops.Dataset.from_tensors(
- acceptance_dist).repeat()
- prob_of_original_ds = dataset_ops.Dataset.from_tensors(
- prob_of_original).repeat()
- else:
- initial_dist_ds = _estimate_initial_dist_ds(
- target_dist_t, class_values_ds)
- acceptance_and_original_prob_ds = initial_dist_ds.map(
- lambda initial: _calculate_acceptance_probs_with_mixing(
- initial, target_dist_t))
- acceptance_dist_ds = acceptance_and_original_prob_ds.map(
- lambda accept_prob, _: accept_prob)
- prob_of_original_ds = acceptance_and_original_prob_ds.map(
- lambda _, prob_original: prob_original)
- filtered_ds = _filter_ds(dataset, acceptance_dist_ds, initial_dist_ds,
- class_values_ds, seed)
- # Prefetch filtered dataset for speed.
- filtered_ds = filtered_ds.prefetch(3)
-
- prob_original_static = _get_prob_original_static(
- initial_dist_t, target_dist_t) if initial_dist is not None else None
- if prob_original_static == 1:
- return dataset_ops.Dataset.zip((class_values_ds, dataset))
- elif prob_original_static == 0:
- return filtered_ds
- else:
- return interleave_ops.sample_from_datasets(
- [dataset_ops.Dataset.zip((class_values_ds, dataset)), filtered_ds],
- weights=prob_of_original_ds.map(lambda prob: [(prob, 1.0 - prob)]),
- seed=seed)
-
- return _apply_fn
-
-
-def _get_prob_original_static(initial_dist_t, target_dist_t):
- """Returns the static probability of sampling from the original.
-
- `tensor_util.constant_value(prob_of_original)` returns `None` if it encounters
- an Op that it isn't defined for. We have some custom logic to avoid this.
-
- Args:
- initial_dist_t: A tensor of the initial distribution.
- target_dist_t: A tensor of the target distribution.
-
- Returns:
- The probability of sampling from the original distribution as a constant,
- if it is a constant, or `None`.
- """
- init_static = tensor_util.constant_value(initial_dist_t)
- target_static = tensor_util.constant_value(target_dist_t)
-
- if init_static is None or target_static is None:
- return None
- else:
- return np.min(target_static / init_static)
-
-
-def _filter_ds(dataset, acceptance_dist_ds, initial_dist_ds, class_values_ds,
- seed):
- """Filters a dataset based on per-class acceptance probabilities.
-
- Args:
- dataset: The dataset to be filtered.
- acceptance_dist_ds: A dataset of acceptance probabilities.
- initial_dist_ds: A dataset of the initial probability distribution, given or
- estimated.
- class_values_ds: A dataset of the corresponding classes.
- seed: (Optional.) Python integer seed for the resampler.
-
- Returns:
- A dataset of (class value, data) after filtering.
- """
- def maybe_warn_on_large_rejection(accept_dist, initial_dist):
- proportion_rejected = math_ops.reduce_sum((1 - accept_dist) * initial_dist)
- return control_flow_ops.cond(
- math_ops.less(proportion_rejected, .5),
- lambda: accept_dist,
- lambda: logging_ops.Print( # pylint: disable=g-long-lambda
- accept_dist, [proportion_rejected, initial_dist, accept_dist],
- message="Proportion of examples rejected by sampler is high: ",
- summarize=100,
- first_n=10))
-
- acceptance_dist_ds = (dataset_ops.Dataset.zip((acceptance_dist_ds,
- initial_dist_ds))
- .map(maybe_warn_on_large_rejection))
-
- def _gather_and_copy(class_val, acceptance_prob, data):
- return class_val, array_ops.gather(acceptance_prob, class_val), data
-
- current_probabilities_and_class_and_data_ds = dataset_ops.Dataset.zip(
- (class_values_ds, acceptance_dist_ds, dataset)).map(_gather_and_copy)
- filtered_ds = (
- current_probabilities_and_class_and_data_ds
- .filter(lambda _1, p, _2: random_ops.random_uniform([], seed=seed) < p))
- return filtered_ds.map(lambda class_value, _, data: (class_value, data))
-
-
-def _estimate_initial_dist_ds(
- target_dist_t, class_values_ds, dist_estimation_batch_size=32,
- smoothing_constant=10):
- num_classes = (target_dist_t.shape[0].value or
- array_ops.shape(target_dist_t)[0])
- initial_examples_per_class_seen = array_ops.fill(
- [num_classes], np.int64(smoothing_constant))
-
- def update_estimate_and_tile(num_examples_per_class_seen, c):
- updated_examples_per_class_seen, dist = _estimate_data_distribution(
- c, num_examples_per_class_seen)
- tiled_dist = array_ops.tile(
- array_ops.expand_dims(dist, 0), [dist_estimation_batch_size, 1])
- return updated_examples_per_class_seen, tiled_dist
-
- initial_dist_ds = (class_values_ds.batch(dist_estimation_batch_size)
- .apply(scan_ops.scan(initial_examples_per_class_seen,
- update_estimate_and_tile))
- .apply(batching.unbatch()))
-
- return initial_dist_ds
-
-
-def _get_target_to_initial_ratio(initial_probs, target_probs):
- # Add tiny to initial_probs to avoid divide by zero.
- denom = (initial_probs + np.finfo(initial_probs.dtype.as_numpy_dtype).tiny)
- return target_probs / denom
-
-
-def _estimate_data_distribution(c, num_examples_per_class_seen):
- """Estimate data distribution as labels are seen.
-
- Args:
- c: The class labels. Type `int32`, shape `[batch_size]`.
- num_examples_per_class_seen: Type `int64`, shape `[num_classes]`,
- containing counts.
-
- Returns:
- num_examples_per_lass_seen: Updated counts. Type `int64`, shape
- `[num_classes]`.
- dist: The updated distribution. Type `float32`, shape `[num_classes]`.
- """
- num_classes = num_examples_per_class_seen.get_shape()[0].value
- # Update the class-count based on what labels are seen in batch.
- num_examples_per_class_seen = math_ops.add(
- num_examples_per_class_seen, math_ops.reduce_sum(
- array_ops.one_hot(c, num_classes, dtype=dtypes.int64), 0))
- init_prob_estimate = math_ops.truediv(
- num_examples_per_class_seen,
- math_ops.reduce_sum(num_examples_per_class_seen))
- dist = math_ops.cast(init_prob_estimate, dtypes.float32)
- return num_examples_per_class_seen, dist
-
-
-def _calculate_acceptance_probs_with_mixing(initial_probs, target_probs):
- """Calculates the acceptance probabilities and mixing ratio.
-
- In this case, we assume that we can *either* sample from the original data
- distribution with probability `m`, or sample from a reshaped distribution
- that comes from rejection sampling on the original distribution. This
- rejection sampling is done on a per-class basis, with `a_i` representing the
- probability of accepting data from class `i`.
-
- This method is based on solving the following analysis for the reshaped
- distribution:
-
- Let F be the probability of a rejection (on any example).
- Let p_i be the proportion of examples in the data in class i (init_probs)
- Let a_i is the rate the rejection sampler should *accept* class i
- Let t_i is the target proportion in the minibatches for class i (target_probs)
-
- ```
- F = sum_i(p_i * (1-a_i))
- = 1 - sum_i(p_i * a_i) using sum_i(p_i) = 1
- ```
-
- An example with class `i` will be accepted if `k` rejections occur, then an
- example with class `i` is seen by the rejector, and it is accepted. This can
- be written as follows:
-
- ```
- t_i = sum_k=0^inf(F^k * p_i * a_i)
- = p_i * a_j / (1 - F) using geometric series identity, since 0 <= F < 1
- = p_i * a_i / sum_j(p_j * a_j) using F from above
- ```
-
- Note that the following constraints hold:
- ```
- 0 <= p_i <= 1, sum_i(p_i) = 1
- 0 <= a_i <= 1
- 0 <= t_i <= 1, sum_i(t_i) = 1
- ```
-
- A solution for a_i in terms of the other variables is the following:
- ```a_i = (t_i / p_i) / max_i[t_i / p_i]```
-
- If we try to minimize the amount of data rejected, we get the following:
-
- M_max = max_i [ t_i / p_i ]
- M_min = min_i [ t_i / p_i ]
-
- The desired probability of accepting data if it comes from class `i`:
-
- a_i = (t_i/p_i - m) / (M_max - m)
-
- The desired probability of pulling a data element from the original dataset,
- rather than the filtered one:
-
- m = M_min
-
- Args:
- initial_probs: A Tensor of the initial probability distribution, given or
- estimated.
- target_probs: A Tensor of the corresponding classes.
-
- Returns:
- (A 1D Tensor with the per-class acceptance probabilities, the desired
- probability of pull from the original distribution.)
- """
- ratio_l = _get_target_to_initial_ratio(initial_probs, target_probs)
- max_ratio = math_ops.reduce_max(ratio_l)
- min_ratio = math_ops.reduce_min(ratio_l)
-
- # Target prob to sample from original distribution.
- m = min_ratio
-
- # TODO(joelshor): Simplify fraction, if possible.
- a_i = (ratio_l - m) / (max_ratio - m)
- return a_i, m
+ return resampling.rejection_resample(class_func, target_dist, initial_dist,
+ seed)
diff --git a/tensorflow/contrib/data/python/ops/scan_ops.py b/tensorflow/contrib/data/python/ops/scan_ops.py
index c52582c..0ca9fdd 100644
--- a/tensorflow/contrib/data/python/ops/scan_ops.py
+++ b/tensorflow/contrib/data/python/ops/scan_ops.py
@@ -17,137 +17,11 @@
from __future__ import division
from __future__ import print_function
-import collections
-
-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.framework import ops
-from tensorflow.python.framework import sparse_tensor
-from tensorflow.python.ops import gen_dataset_ops
+from tensorflow.python.data.experimental.ops import scan_ops
+from tensorflow.python.util import deprecation
-class _ScanDataset(dataset_ops.UnaryDataset):
- """A dataset that scans a function across its input."""
-
- def __init__(self, input_dataset, initial_state, scan_func):
- """See `scan()` for details."""
- super(_ScanDataset, self).__init__(input_dataset)
- self._input_dataset = input_dataset
-
- with ops.name_scope("initial_state"):
- # Convert any `SparseTensorValue`s to `SparseTensor`s and all other
- # values to tensors.
- self._initial_state = nest.pack_sequence_as(initial_state, [
- sparse_tensor.SparseTensor.from_value(t)
- if sparse_tensor.is_sparse(t) else ops.convert_to_tensor(
- t, name="component_%d" % i)
- for i, t in enumerate(nest.flatten(initial_state))
- ])
-
- # 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
-
- # Iteratively rerun the scan function until reaching a fixed point on
- # `self._state_shapes`.
- need_to_rerun = True
- while need_to_rerun:
-
- wrapped_func = dataset_ops.StructuredFunctionWrapper(
- scan_func, "tf.contrib.data.scan()",
- 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),
- add_to_graph=False)
- if not (
- isinstance(wrapped_func.output_types, collections.Sequence) and
- len(wrapped_func.output_types) == 2):
- raise TypeError("The scan function must return a pair comprising the "
- "new state and the output value.")
-
- 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(
- nest.flatten(new_state_classes),
- nest.flatten(self._state_classes)):
- if not issubclass(new_state_class, 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))
-
- # 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:
- raise TypeError(
- "The element types for the new state must match the initial "
- "state. Expected %s; got %s." %
- (self._state_types, new_state_types))
-
- # Extract shape information from the returned values.
- new_state_shapes, self._output_shapes = wrapped_func.output_shapes
-
- flat_state_shapes = nest.flatten(self._state_shapes)
- flat_new_state_shapes = nest.flatten(new_state_shapes)
- weakened_state_shapes = [
- original.most_specific_compatible_shape(new)
- for original, new in zip(flat_state_shapes, flat_new_state_shapes)
- ]
-
- need_to_rerun = False
- for original_shape, weakened_shape in zip(flat_state_shapes,
- weakened_state_shapes):
- if original_shape.ndims is not None and (
- weakened_shape.ndims is None or
- original_shape.as_list() != weakened_shape.as_list()):
- need_to_rerun = True
- break
-
- if need_to_rerun:
- self._state_shapes = nest.pack_sequence_as(self._state_shapes,
- weakened_state_shapes)
-
- self._scan_func = wrapped_func.function
- self._scan_func.add_to_graph(ops.get_default_graph())
-
- def _as_variant_tensor(self):
- input_t = self._input_dataset._as_variant_tensor() # pylint: disable=protected-access
- return gen_dataset_ops.scan_dataset(
- input_t,
- nest.flatten(sparse.serialize_sparse_tensors(self._initial_state)),
- self._scan_func.captured_inputs,
- f=self._scan_func,
- **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
-
-
+@deprecation.deprecated(None, "Use `tf.data.experimental.scan(...)`.")
def scan(initial_state, scan_func):
"""A transformation that scans a function across an input dataset.
@@ -168,7 +42,4 @@
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
- def _apply_fn(dataset):
- return _ScanDataset(dataset, initial_state, scan_func)
-
- return _apply_fn
+ return scan_ops.scan(initial_state, scan_func)
diff --git a/tensorflow/contrib/data/python/ops/shuffle_ops.py b/tensorflow/contrib/data/python/ops/shuffle_ops.py
index 985d1d8..329b34f 100644
--- a/tensorflow/contrib/data/python/ops/shuffle_ops.py
+++ b/tensorflow/contrib/data/python/ops/shuffle_ops.py
@@ -17,54 +17,12 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.python.data.ops import dataset_ops
-from tensorflow.python.data.util import random_seed
-from tensorflow.python.framework import constant_op
-from tensorflow.python.framework import dtypes
-from tensorflow.python.framework import ops
-from tensorflow.python.ops import gen_dataset_ops
+from tensorflow.python.data.experimental.ops import shuffle_ops
+from tensorflow.python.util import deprecation
-class _ShuffleAndRepeatDataset(dataset_ops.UnaryDataset):
- """A `Dataset` that fuses `shuffle` and `repeat`."""
-
- def __init__(self, input_dataset, buffer_size, count=None, seed=None):
- super(_ShuffleAndRepeatDataset, self).__init__(input_dataset)
- self._input_dataset = input_dataset
- self._buffer_size = ops.convert_to_tensor(
- buffer_size, dtype=dtypes.int64, name="buffer_size")
- if count is None:
- self._count = constant_op.constant(-1, dtype=dtypes.int64, name="count")
- else:
- self._count = ops.convert_to_tensor(
- count, dtype=dtypes.int64, name="count")
- self._seed, self._seed2 = random_seed.get_seed(seed)
-
- def _as_variant_tensor(self):
- # pylint: disable=protected-access
- input_resource = self._input_dataset._as_variant_tensor()
- return gen_dataset_ops.shuffle_and_repeat_dataset(
- input_resource,
- buffer_size=self._buffer_size,
- count=self._count,
- seed=self._seed,
- seed2=self._seed2,
- **dataset_ops.flat_structure(self))
- # pylint: enable=protected-access
-
- @property
- def output_classes(self):
- return self._input_dataset.output_classes
-
- @property
- def output_shapes(self):
- return self._input_dataset.output_shapes
-
- @property
- def output_types(self):
- return self._input_dataset.output_types
-
-
+@deprecation.deprecated(None,
+ "Use `tf.data.experimental.shuffle_and_repeat(...)`.")
def shuffle_and_repeat(buffer_size, count=None, seed=None):
"""Shuffles and repeats a Dataset returning a new permutation for each epoch.
@@ -93,8 +51,4 @@
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
-
- def _apply_fn(dataset): # pylint: disable=missing-docstring
- return _ShuffleAndRepeatDataset(dataset, buffer_size, count, seed)
-
- return _apply_fn
+ return shuffle_ops.shuffle_and_repeat(buffer_size, count, seed)
diff --git a/tensorflow/contrib/data/python/ops/threadpool.py b/tensorflow/contrib/data/python/ops/threadpool.py
index 9d165ad..20cceb4 100644
--- a/tensorflow/contrib/data/python/ops/threadpool.py
+++ b/tensorflow/contrib/data/python/ops/threadpool.py
@@ -17,89 +17,6 @@
from __future__ import division
from __future__ import print_function
-import threading
-
-from tensorflow.contrib.data.python.ops import contrib_op_loader # pylint: disable=unused-import
-from tensorflow.contrib.data.python.ops import gen_dataset_ops
-from tensorflow.python.data.ops import dataset_ops
-from tensorflow.python.eager import context
-from tensorflow.python.ops import resource_variable_ops
-
-_uid_counter = 0
-_uid_lock = threading.Lock()
-
-
-def _generate_shared_name(prefix):
- with _uid_lock:
- global _uid_counter
- uid = _uid_counter
- _uid_counter += 1
- return "{}{}".format(prefix, uid)
-
-
-# TODO(b/73383364): Properly export in the `tf.contrib.data` API when stable
-# or make private / remove.
-class PrivateThreadPool(object):
- """A stateful resource that represents a private thread pool."""
-
- def __init__(self, num_threads, display_name=None,
- max_intra_op_parallelism=1):
- """Creates a `PrivateThreadPool` with the given number of threads."""
- if context.executing_eagerly():
- shared_name = _generate_shared_name("privatethreadpool")
- self._resource = gen_dataset_ops.thread_pool_handle(
- num_threads=num_threads,
- max_intra_op_parallelism=max_intra_op_parallelism,
- display_name=display_name,
- shared_name=shared_name)
- self._resource_deleter = resource_variable_ops.EagerResourceDeleter(
- handle=self._resource, handle_device=context.context().device_name)
- else:
- self._resource = gen_dataset_ops.thread_pool_handle(
- num_threads=num_threads,
- max_intra_op_parallelism=max_intra_op_parallelism,
- display_name=display_name)
-
-
-class _ThreadPoolDataset(dataset_ops.UnaryDataset):
- """A `Dataset` that acts as an identity, and sets a custom threadpool."""
-
- def __init__(self, input_dataset, thread_pool):
- super(_ThreadPoolDataset, self).__init__(input_dataset)
- self._input_dataset = input_dataset
- self._thread_pool = thread_pool
-
- def _as_variant_tensor(self):
- return gen_dataset_ops.thread_pool_dataset(
- self._input_dataset._as_variant_tensor(), # pylint: disable=protected-access
- self._thread_pool._resource, # pylint: disable=protected-access
- **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
-
-
-# TODO(b/73383364): Properly export in the `tf.contrib.data` API when stable
-# or make private / remove.
-def override_threadpool(dataset, thread_pool):
- """Returns a new dataset that uses the given thread pool for its operations.
-
- Args:
- dataset: A `tf.data.Dataset` object.
- thread_pool: A `PrivateThreadPool` object.
-
- Returns:
- A dataset containing the same values as `dataset`, but which uses
- `thread_pool` to compute any of its parallel operations (such as
- `tf.data.Dataset.map`).
- """
- return _ThreadPoolDataset(dataset, thread_pool)
+# pylint: disable=unused-import
+from tensorflow.python.data.experimental.ops.threadpool import override_threadpool
+from tensorflow.python.data.experimental.ops.threadpool import PrivateThreadPool
diff --git a/tensorflow/contrib/data/python/ops/unique.py b/tensorflow/contrib/data/python/ops/unique.py
index bad67a5..909d06c 100644
--- a/tensorflow/contrib/data/python/ops/unique.py
+++ b/tensorflow/contrib/data/python/ops/unique.py
@@ -17,12 +17,11 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.contrib.data.python.ops import contrib_op_loader # pylint: disable=unused-import
-from tensorflow.contrib.data.python.ops import gen_dataset_ops
-from tensorflow.python.data.ops import dataset_ops
-from tensorflow.python.framework import dtypes
+from tensorflow.python.data.experimental.ops import unique as experimental_unique
+from tensorflow.python.util import deprecation
+@deprecation.deprecated(None, "Use `tf.data.experimental.unique()`.")
def unique():
"""Creates a `Dataset` from another `Dataset`, discarding duplicates.
@@ -40,39 +39,4 @@
A `Dataset` transformation function, which can be passed to
`tf.data.Dataset.apply`.
"""
-
- def _apply_fn(dataset):
- return _UniqueDataset(dataset)
-
- return _apply_fn
-
-
-class _UniqueDataset(dataset_ops.UnaryDataset):
- """A `Dataset` contains the unique elements from its input."""
-
- def __init__(self, input_dataset):
- """See `unique()` for details."""
- super(_UniqueDataset, self).__init__(input_dataset)
- self._input_dataset = input_dataset
- if input_dataset.output_types not in (dtypes.int32, dtypes.int64,
- dtypes.string):
- raise TypeError(
- "`tf.contrib.data.unique()` only supports inputs with a single "
- "`tf.int32`, `tf.int64`, or `tf.string` component.")
-
- def _as_variant_tensor(self):
- return gen_dataset_ops.unique_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 self._input_dataset.output_shapes
-
- @property
- def output_types(self):
- return self._input_dataset.output_types
+ return experimental_unique.unique()
diff --git a/tensorflow/contrib/data/python/ops/writers.py b/tensorflow/contrib/data/python/ops/writers.py
index c455fdc..42fb69b 100644
--- a/tensorflow/contrib/data/python/ops/writers.py
+++ b/tensorflow/contrib/data/python/ops/writers.py
@@ -17,42 +17,14 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.python.data.ops import dataset_ops
-from tensorflow.python.data.util import convert
-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_dataset_ops
+from tensorflow.python.data.experimental.ops import writers
+from tensorflow.python.util import deprecation
-class TFRecordWriter(object):
+class TFRecordWriter(writers.TFRecordWriter):
"""Writes data to a TFRecord file."""
+ @deprecation.deprecated(
+ None, "Use `tf.data.experimental.TFRecordWriter(...)`.")
def __init__(self, filename, compression_type=None):
- self._filename = ops.convert_to_tensor(
- filename, dtypes.string, name="filename")
- self._compression_type = convert.optional_param_to_tensor(
- "compression_type",
- compression_type,
- argument_default="",
- argument_dtype=dtypes.string)
-
- def write(self, dataset):
- """Returns a `tf.Operation` to write a dataset to a file.
-
- Args:
- dataset: a `tf.data.Dataset` whose elements are to be written to a file
-
- Returns:
- A `tf.Operation` that, when run, writes contents of `dataset` to a file.
- """
- if not isinstance(dataset, dataset_ops.Dataset):
- raise TypeError("`dataset` must be a `tf.data.Dataset` object.")
- if (dataset.output_types != dtypes.string or
- dataset.output_shapes != tensor_shape.scalar()):
- raise TypeError(
- "`dataset` must produce scalar `DT_STRING` tensors whereas it "
- "produces shape {0} and types {1}".format(dataset.output_shapes,
- dataset.output_types))
- return gen_dataset_ops.dataset_to_tf_record(
- dataset._as_variant_tensor(), self._filename, self._compression_type) # pylint: disable=protected-access
+ super(TFRecordWriter, self).__init__(filename, compression_type)
diff --git a/tensorflow/contrib/decision_trees/proto/BUILD b/tensorflow/contrib/decision_trees/proto/BUILD
index 3b50a48..06940a9 100644
--- a/tensorflow/contrib/decision_trees/proto/BUILD
+++ b/tensorflow/contrib/decision_trees/proto/BUILD
@@ -17,7 +17,6 @@
name = "generic_tree_model",
srcs = ["generic_tree_model.proto"],
cc_api_version = 2,
- java_api_version = 2,
visibility = ["//visibility:public"],
)
diff --git a/tensorflow/contrib/distribute/README.md b/tensorflow/contrib/distribute/README.md
index 91a27f9..2e02576 100644
--- a/tensorflow/contrib/distribute/README.md
+++ b/tensorflow/contrib/distribute/README.md
@@ -231,7 +231,8 @@
important to shuffle your dataset in your `input_fn`.
`MirroredStrategy` will insert a `tf.dataset.Dataset.shard` call in you
-`input_fn`. As a result, each worker gets a fraction of your input data.
+`input_fn` if `auto_shard_dataset` is set to `True`. As a result, each worker
+gets a fraction of your input data.
### Performance Tips
diff --git a/tensorflow/contrib/distribute/python/BUILD b/tensorflow/contrib/distribute/python/BUILD
index 7eead6e..cfb9d42 100644
--- a/tensorflow/contrib/distribute/python/BUILD
+++ b/tensorflow/contrib/distribute/python/BUILD
@@ -22,14 +22,15 @@
visibility = ["//tensorflow:internal"],
deps = [
":input_ops",
+ ":prefetching_ops_v2",
"//tensorflow/python:array_ops",
"//tensorflow/python:control_flow_ops",
"//tensorflow/python:device_util",
"//tensorflow/python:distribute",
"//tensorflow/python:framework_ops",
+ "//tensorflow/python:resource_variable_ops",
"//tensorflow/python:training",
"//tensorflow/python:util",
- "//tensorflow/python/data/ops:multi_device_iterator_ops",
"//tensorflow/python/eager:context",
"//tensorflow/python/training/checkpointable:base",
"@six_archive//:six",
@@ -453,7 +454,7 @@
cuda_py_test(
name = "estimator_training_test",
- size = "large",
+ size = "enormous",
srcs = ["estimator_training_test.py"],
additional_deps = [
":combinations",
@@ -648,6 +649,32 @@
)
py_library(
+ name = "prefetching_ops_v2",
+ srcs = ["prefetching_ops_v2.py"],
+ deps = [
+ "//tensorflow/contrib/data/python/ops:prefetching_ops",
+ "//tensorflow/python:experimental_dataset_ops_gen",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/util:nest",
+ "//tensorflow/python/data/util:sparse",
+ ],
+)
+
+cuda_py_test(
+ name = "prefetching_ops_v2_test",
+ srcs = ["prefetching_ops_v2_test.py"],
+ additional_deps = [
+ ":prefetching_ops_v2",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python:framework_test_lib",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/ops:iterator_ops",
+ ],
+)
+
+py_library(
name = "input_ops",
srcs = ["input_ops.py"],
visibility = ["//tensorflow:internal"],
diff --git a/tensorflow/contrib/distribute/python/collective_all_reduce_strategy.py b/tensorflow/contrib/distribute/python/collective_all_reduce_strategy.py
index c900b41..9809204 100644
--- a/tensorflow/contrib/distribute/python/collective_all_reduce_strategy.py
+++ b/tensorflow/contrib/distribute/python/collective_all_reduce_strategy.py
@@ -216,7 +216,7 @@
"""Configures the object.
Args:
- session_config: a @{tf.ConfigProto}
+ session_config: a `tf.ConfigProto`
cluster_spec: a dict, ClusterDef or ClusterSpec object specifying the
cluster configurations.
task_type: the current task type, such as "worker".
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 33ffbf6..6796a23 100644
--- a/tensorflow/contrib/distribute/python/collective_all_reduce_strategy_test.py
+++ b/tensorflow/contrib/distribute/python/collective_all_reduce_strategy_test.py
@@ -128,7 +128,8 @@
# TODO(yuefengz): support non-Mirrored variable as destinations.
g = d.reduce(
variable_scope.VariableAggregation.SUM, g, destinations=v)
- with ops.control_dependencies(d.unwrap(d.update(v, update, g))):
+ with ops.control_dependencies(
+ d.update(v, update, g, grouped=False)):
after_list.append(d.read_var(v))
return before_list, after_list
diff --git a/tensorflow/contrib/distribute/python/combinations.py b/tensorflow/contrib/distribute/python/combinations.py
index 244d1fc..82ca041 100644
--- a/tensorflow/contrib/distribute/python/combinations.py
+++ b/tensorflow/contrib/distribute/python/combinations.py
@@ -59,6 +59,7 @@
from tensorflow.python.training import adam
from tensorflow.python.training import distribution_strategy_context
from tensorflow.python.training import gradient_descent
+from tensorflow.python.training import rmsprop
from tensorflow.python.util import tf_inspect
@@ -354,6 +355,8 @@
"GradientDescentV1", lambda: gradient_descent.GradientDescentOptimizer(0.2))
adagrad_optimizer_v1_fn = NamedObject(
"AdagradV1", lambda: adagrad.AdagradOptimizer(0.001))
+rmsprop_optimizer_v1_fn = NamedObject(
+ "RmsPropV1", lambda: rmsprop.RMSPropOptimizer(0.001))
optimizers_v1 = [adam_optimizer_v1_fn, gradient_descent_optimizer_v1_fn,
adagrad_optimizer_v1_fn]
diff --git a/tensorflow/contrib/distribute/python/examples/simple_estimator_example.py b/tensorflow/contrib/distribute/python/examples/simple_estimator_example.py
index 44a69ed..79a9803 100644
--- a/tensorflow/contrib/distribute/python/examples/simple_estimator_example.py
+++ b/tensorflow/contrib/distribute/python/examples/simple_estimator_example.py
@@ -22,6 +22,8 @@
import tensorflow as tf
+from tensorflow.python.keras import metrics as metrics_module
+
def build_model_fn_optimizer():
"""Simple model_fn with optimizer."""
@@ -45,7 +47,10 @@
return y * y
if mode == tf.estimator.ModeKeys.EVAL:
- return tf.estimator.EstimatorSpec(mode, loss=loss_fn())
+ acc_obj = metrics_module.BinaryAccuracy()
+ acc_obj.update_state(labels, labels)
+ return tf.estimator.EstimatorSpec(
+ mode, loss=loss_fn(), eval_metric_ops={"Accuracy": acc_obj})
assert mode == tf.estimator.ModeKeys.TRAIN
@@ -61,18 +66,26 @@
["/device:GPU:0", "/device:GPU:1"])
config = tf.estimator.RunConfig(train_distribute=distribution,
eval_distribute=distribution)
+ # Since there are 2 devices and 10 samples, we set steps=5.
+ steps = 5
- def input_fn():
+ def train_input_fn():
features = tf.data.Dataset.from_tensors([[1.]]).repeat(10)
labels = tf.data.Dataset.from_tensors([1.]).repeat(10)
return tf.data.Dataset.zip((features, labels))
estimator = tf.estimator.Estimator(
model_fn=build_model_fn_optimizer(), config=config)
- estimator.train(input_fn=input_fn, steps=10)
+ estimator.train(input_fn=train_input_fn, steps=steps)
- eval_result = estimator.evaluate(input_fn=input_fn, steps=10)
+ def eval_input_fn():
+ features = tf.data.Dataset.from_tensors([[1.]]).repeat(10)
+ labels = tf.data.Dataset.from_tensors([1.]).repeat(10)
+ return tf.data.Dataset.zip((features, labels))
+
+ eval_result = estimator.evaluate(input_fn=eval_input_fn, steps=steps)
print("Eval result: {}".format(eval_result))
+ assert eval_result["Accuracy"] == 1.0
def predict_input_fn():
predict_features = tf.data.Dataset.from_tensors([[1.]]).repeat(10)
diff --git a/tensorflow/contrib/distribute/python/keras_test.py b/tensorflow/contrib/distribute/python/keras_test.py
index 2e6cd43..3aab2c5 100644
--- a/tensorflow/contrib/distribute/python/keras_test.py
+++ b/tensorflow/contrib/distribute/python/keras_test.py
@@ -173,13 +173,42 @@
return dataset.batch(batch_size)
-def all_combinations():
+def get_model():
+ x = keras.layers.Input(shape=(3,), name='input')
+ y = keras.layers.Dense(4, name='dense')(x)
+ model = keras.Model(x, y)
+ return model
+
+
+def get_dataset(distribution):
+ inputs = np.zeros((10, 3), dtype=np.float32)
+ targets = np.zeros((10, 4), dtype=np.float32)
+ dataset = dataset_ops.Dataset.from_tensor_slices((inputs, targets))
+ dataset = dataset.repeat(100)
+ dataset = batch_wrapper(dataset, 10, distribution)
+ return dataset
+
+
+strategies = [combinations.default_strategy,
+ combinations.one_device_strategy,
+ combinations.mirrored_strategy_with_gpu_and_cpu,
+ combinations.mirrored_strategy_with_two_gpus,
+ combinations.tpu_strategy_one_step]
+
+
+def strategy_combinations():
return combinations.combine(
- distribution=[combinations.default_strategy,
- combinations.one_device_strategy,
- combinations.mirrored_strategy_with_gpu_and_cpu,
- combinations.mirrored_strategy_with_two_gpus,
- combinations.tpu_strategy_one_step],
+ distribution=strategies,
+ mode=['graph'])
+
+
+def strategy_and_optimizer_combinations():
+ return combinations.combine(
+ distribution=strategies,
+ optimizer=[combinations.adagrad_optimizer_v1_fn,
+ combinations.adam_optimizer_v1_fn,
+ combinations.gradient_descent_optimizer_v1_fn,
+ combinations.rmsprop_optimizer_v1_fn],
mode=['graph'])
@@ -205,6 +234,7 @@
keras_model = simple_functional_model()
keras_model.compile(
loss='categorical_crossentropy',
+ metrics=[keras.metrics.CategoricalAccuracy()],
optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.01))
config = run_config_lib.RunConfig(tf_random_seed=_RANDOM_SEED,
model_dir=self._base_dir,
@@ -229,6 +259,7 @@
keras_model = simple_sequential_model()
keras_model.compile(
loss='categorical_crossentropy',
+ metrics=[keras.metrics.CategoricalAccuracy()],
optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.01))
config = run_config_lib.RunConfig(tf_random_seed=_RANDOM_SEED,
model_dir=self._base_dir,
@@ -358,13 +389,11 @@
def test_calling_model_with_numpy_arrays(self):
with self.cached_session():
- x = keras.layers.Input(shape=(3,), name='input')
- y = keras.layers.Dense(4, name='dense')(x)
- model = keras.Model(x, y)
+ model = get_model()
optimizer = gradient_descent.GradientDescentOptimizer(0.001)
loss = 'mse'
- metrics = ['mae']
+ metrics = ['mae', keras.metrics.CategoricalAccuracy()]
strategy = mirrored_strategy.MirroredStrategy(['/device:GPU:1',
'/device:GPU:0'])
model.compile(optimizer, loss, metrics=metrics, distribute=strategy)
@@ -390,23 +419,17 @@
# with batch_size
model.predict(inputs, batch_size=8)
- @combinations.generate(all_combinations())
+ @combinations.generate(strategy_combinations())
def test_calling_model_on_same_dataset(self, distribution):
with self.cached_session():
- x = keras.layers.Input(shape=(3,), name='input')
- y = keras.layers.Dense(4, name='dense')(x)
- model = keras.Model(x, y)
+ model = get_model()
optimizer = gradient_descent.GradientDescentOptimizer(0.001)
loss = 'mse'
- metrics = ['mae']
+ metrics = ['mae', keras.metrics.CategoricalAccuracy()]
model.compile(optimizer, loss, metrics=metrics, distribute=distribution)
- inputs = np.zeros((10, 3), dtype=np.float32)
- targets = np.zeros((10, 4), dtype=np.float32)
- dataset = dataset_ops.Dataset.from_tensor_slices((inputs, targets))
- dataset = dataset.repeat(100)
- dataset = batch_wrapper(dataset, 10, distribution)
+ dataset = get_dataset(distribution)
# Call fit with validation data
model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=0,
@@ -432,7 +455,7 @@
optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=0.001)
loss = 'mse'
- metrics = ['mae']
+ metrics = ['mae', keras.metrics.CategoricalAccuracy()]
strategy = mirrored_strategy.MirroredStrategy(['/device:GPU:0',
'/device:CPU:0'])
model.compile(optimizer, loss, metrics=metrics, distribute=strategy)
@@ -459,23 +482,17 @@
model.fit(dataset_dict, epochs=1, steps_per_epoch=2, verbose=1)
- @combinations.generate(all_combinations())
+ @combinations.generate(strategy_combinations())
def test_fit_eval_and_predict_methods_on_dataset(self, distribution):
with self.cached_session():
- x = keras.layers.Input(shape=(3,), name='input')
- y = keras.layers.Dense(4, name='dense')(x)
- model = keras.Model(x, y)
+ model = get_model()
optimizer = gradient_descent.GradientDescentOptimizer(0.001)
loss = 'mse'
- metrics = ['mae']
+ metrics = ['mae', keras.metrics.CategoricalAccuracy()]
model.compile(optimizer, loss, metrics=metrics, distribute=distribution)
- inputs = np.zeros((10, 3), dtype=np.float32)
- targets = np.zeros((10, 4), dtype=np.float32)
- dataset = dataset_ops.Dataset.from_tensor_slices((inputs, targets))
- dataset = dataset.repeat(100)
- dataset = batch_wrapper(dataset, 10, distribution)
+ dataset = get_dataset(distribution)
model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=1)
model.evaluate(dataset, steps=2, verbose=1)
@@ -484,37 +501,23 @@
model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=0,
validation_data=dataset, validation_steps=2)
- def test_raise_error_for_stateful_metrics(self):
-
- class ExampleStatefulMetric(keras.layers.Layer):
-
- def __init__(self, name='true_positives', **kwargs):
- super(ExampleStatefulMetric, self).__init__(name=name, **kwargs)
- self.stateful = True
-
- def __call__(self, y_true, y_pred):
- return y_pred - y_true
-
+ @combinations.generate(strategy_and_optimizer_combinations())
+ def test_fit_eval_and_predict_with_optimizer(self, distribution, optimizer):
with self.cached_session():
- x = keras.layers.Input(shape=(3,), name='input')
- y = keras.layers.Dense(4, name='dense')(x)
- model = keras.Model(x, y)
+ model = get_model()
- optimizer = gradient_descent.GradientDescentOptimizer(0.001)
loss = 'mse'
- metrics = ['mae', ExampleStatefulMetric()]
- strategy = mirrored_strategy.MirroredStrategy(['/device:GPU:1',
- '/device:GPU:0'])
- with self.assertRaisesRegexp(
- NotImplementedError, 'Stateful metrics are not supported with '
- 'DistributionStrategy.'):
- model.compile(optimizer, loss, metrics=metrics, distribute=strategy)
+ model.compile(optimizer(), loss, distribute=distribution)
+
+ dataset = get_dataset(distribution)
+
+ model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=1)
+ model.evaluate(dataset, steps=2, verbose=1)
+ model.predict(dataset, steps=2)
def test_unsupported_features(self):
with self.cached_session():
- x = keras.layers.Input(shape=(3,), name='input')
- y = keras.layers.Dense(4, name='dense')(x)
- model = keras.Model(x, y)
+ model = get_model()
optimizer = gradient_descent.GradientDescentOptimizer(0.001)
loss = 'mse'
@@ -524,11 +527,7 @@
model.compile(optimizer, loss, metrics=metrics, distribute=strategy)
- inputs = np.zeros((10, 3), dtype=np.float32)
- targets = np.zeros((10, 4), dtype=np.float32)
- dataset = dataset_ops.Dataset.from_tensor_slices((inputs, targets))
- dataset = dataset.repeat(100)
- dataset = dataset.batch(10)
+ dataset = get_dataset(strategy)
# Test with validation split
with self.assertRaisesRegexp(
@@ -565,9 +564,7 @@
def test_calling_with_unsupported_predefined_callbacks(self):
with self.cached_session():
- x = keras.layers.Input(shape=(3,), name='input')
- y = keras.layers.Dense(4, name='dense')(x)
- model = keras.Model(x, y)
+ model = get_model()
optimizer = gradient_descent.GradientDescentOptimizer(0.001)
loss = 'mse'
@@ -576,11 +573,7 @@
'/device:GPU:0'])
model.compile(optimizer, loss, metrics=metrics, distribute=strategy)
- inputs = np.zeros((10, 3), dtype=np.float32)
- targets = np.zeros((10, 4), dtype=np.float32)
- dataset = dataset_ops.Dataset.from_tensor_slices((inputs, targets))
- dataset = dataset.repeat(100)
- dataset = dataset.batch(10)
+ dataset = get_dataset(strategy)
def schedule(_):
return 0.001
@@ -604,9 +597,7 @@
def test_dataset_input_shape_validation(self):
with self.cached_session():
- x = keras.layers.Input(shape=(3,), name='input')
- y = keras.layers.Dense(4, name='dense')(x)
- model = keras.Model(x, y)
+ model = get_model()
optimizer = rmsprop.RMSPropOptimizer(learning_rate=0.001)
loss = 'mse'
@@ -640,17 +631,13 @@
mode=['graph']))
def test_dataset_input_shape_fully_defined(self, distribution):
with self.cached_session():
- x = keras.layers.Input(shape=(3,), name='input')
- y = keras.layers.Dense(4, name='dense')(x)
- model = keras.Model(x, y)
+ model = get_model()
optimizer = rmsprop.RMSPropOptimizer(learning_rate=0.001)
loss = 'mse'
model.compile(optimizer, loss, distribute=distribution)
- inputs = np.zeros((10, 3), dtype=np.float32)
- targets = np.zeros((10, 4), dtype=np.float32)
- dataset = dataset_ops.Dataset.from_tensor_slices((inputs, targets))
+ dataset = get_dataset(distribution)
# Input shapes are not fully known. Batch dimension is unknown as we are
# not using the drop_remainder argument.
dataset = dataset.repeat(100).batch(10)
@@ -722,7 +709,7 @@
class NormalizationLayerWithDistributionStrategyTest(
test.TestCase, parameterized.TestCase):
- @combinations.generate(all_combinations())
+ @combinations.generate(strategy_combinations())
def test_batchnorm_correctness(self, distribution):
with self.cached_session():
model = keras.models.Sequential()
@@ -750,7 +737,37 @@
class CorrectnessWithDistributionStrategyTest(test.TestCase,
parameterized.TestCase):
- @combinations.generate(all_combinations())
+ @combinations.generate(strategy_combinations())
+ def test_metric_correctness(self, distribution):
+ with self.cached_session():
+ keras.backend.set_image_data_format('channels_last')
+ num_samples = 10000
+
+ x_train = np.random.randint(0, 2, num_samples)
+ x_train = np.reshape(x_train, (num_samples, 1))
+ y_train = x_train
+ x_train = x_train.astype('float32')
+ y_train = y_train.astype('float32')
+
+ # Create identity model.
+ model = keras.Sequential()
+ model.add(
+ keras.layers.Dense(1, input_shape=(1,), kernel_initializer='ones'))
+ model.compile(
+ loss=keras.losses.mean_squared_error,
+ optimizer=gradient_descent.GradientDescentOptimizer(0.5),
+ metrics=[keras.metrics.BinaryAccuracy()],
+ distribute=distribution)
+
+ batch_size = 64
+ batch_size //= distribution.num_towers
+ 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])
+
+ @combinations.generate(strategy_combinations())
def test_correctness(self, distribution):
with self.cached_session():
keras.backend.set_image_data_format('channels_last')
diff --git a/tensorflow/contrib/distribute/python/metrics_v1_test.py b/tensorflow/contrib/distribute/python/metrics_v1_test.py
index f7773af..8163494 100644
--- a/tensorflow/contrib/distribute/python/metrics_v1_test.py
+++ b/tensorflow/contrib/distribute/python/metrics_v1_test.py
@@ -86,11 +86,10 @@
def _test_metric(self, distribution, dataset_fn, metric_fn, expected_fn):
with ops.Graph().as_default(), distribution.scope():
iterator = distribution.distribute_dataset(
- dataset_fn).make_initializable_iterator()
+ dataset_fn).make_one_shot_iterator()
value, update = distribution.call_for_each_tower(
metric_fn, iterator.get_next())
update = distribution.group(update)
- self.evaluate(iterator.initializer)
self.evaluate(variables.local_variables_initializer())
# TODO(josh11b): Once we switch to using a global batch size for input,
# replace "distribution.num_towers" with "1".
diff --git a/tensorflow/contrib/distribute/python/minimize_loss_test.py b/tensorflow/contrib/distribute/python/minimize_loss_test.py
index d082d5c..ba147e7 100644
--- a/tensorflow/contrib/distribute/python/minimize_loss_test.py
+++ b/tensorflow/contrib/distribute/python/minimize_loss_test.py
@@ -41,14 +41,6 @@
class MinimizeLossStepTest(test.TestCase, parameterized.TestCase):
- def _get_iterator(self, ds):
- if context.executing_eagerly():
- iterator = ds.make_one_shot_iterator()
- else:
- iterator = ds.make_initializable_iterator()
- self.evaluate(iterator.initializer)
- return iterator
-
@combinations.generate(
combinations.times(
combinations.distributions_and_v1_optimizers(),
@@ -70,7 +62,8 @@
distribution.call_for_each_tower(
model_fn, *inputs, run_concurrently=layer.built))
- iterator = self._get_iterator(distribution.distribute_dataset(dataset_fn))
+ iterator = distribution.distribute_dataset(
+ dataset_fn).make_one_shot_iterator()
def run_step():
return distribution.run_steps_on_dataset(
@@ -106,7 +99,8 @@
model_fn, dataset_fn, layer = minimize_loss_example(
optimizer_fn, use_bias=True, use_callable_loss=use_callable_loss)
- iterator = self._get_iterator(distribution.distribute_dataset(dataset_fn))
+ iterator = distribution.distribute_dataset(
+ dataset_fn).make_one_shot_iterator()
def run_step():
return distribution.group(
@@ -165,7 +159,8 @@
distribution.call_for_each_tower(
model_fn, *inputs, run_concurrently=layer.built))
- iterator = self._get_iterator(distribution.distribute_dataset(dataset_fn))
+ iterator = distribution.distribute_dataset(
+ dataset_fn).make_one_shot_iterator()
def run_step():
return distribution.run_steps_on_dataset(
@@ -249,7 +244,8 @@
fetches += ops.get_collection(ops.GraphKeys.UPDATE_OPS)
return control_flow_ops.group(fetches)
- iterator = self._get_iterator(distribution.distribute_dataset(dataset_fn))
+ iterator = distribution.distribute_dataset(
+ dataset_fn).make_one_shot_iterator()
def run_step():
return distribution.run_steps_on_dataset(
@@ -342,7 +338,8 @@
distribution.call_for_each_tower(
model_fn, x, y, run_concurrently=False))
- iterator = self._get_iterator(distribution.distribute_dataset(dataset_fn))
+ iterator = distribution.distribute_dataset(
+ dataset_fn).make_one_shot_iterator()
def run_step():
return distribution.run_steps_on_dataset(
@@ -435,7 +432,8 @@
output=loss)
return distribution.group(train_op)
- iterator = self._get_iterator(distribution.distribute_dataset(dataset_fn))
+ iterator = distribution.distribute_dataset(
+ dataset_fn).make_one_shot_iterator()
def run_step():
initial_loss = lambda: constant_op.constant(1e7)
diff --git a/tensorflow/contrib/distribute/python/mirrored_strategy.py b/tensorflow/contrib/distribute/python/mirrored_strategy.py
index 945f450..6bd380a 100644
--- a/tensorflow/contrib/distribute/python/mirrored_strategy.py
+++ b/tensorflow/contrib/distribute/python/mirrored_strategy.py
@@ -347,6 +347,8 @@
set, the `configure` method will try to find the best one.
prefetch_on_device: optional boolean to specify whether to prefetch input
data to devices.
+ auto_shard_dataset: whether to auto-shard the dataset when there are
+ multiple workers.
"""
def __init__(self,
@@ -354,11 +356,13 @@
num_gpus=None,
num_gpus_per_worker=None,
cross_tower_ops=None,
- prefetch_on_device=None):
+ prefetch_on_device=None,
+ auto_shard_dataset=False):
super(MirroredStrategy, self).__init__()
self._cross_tower_ops = cross_tower_ops
self._prefetch_on_device = prefetch_on_device
+ self._auto_shard_dataset = auto_shard_dataset
# Rememeber num GPUs which might be needed by `configure` method.
if num_gpus is not None and num_gpus_per_worker is not None:
raise ValueError(
@@ -477,13 +481,11 @@
if self._cluster_spec:
return values.MultiWorkerDataset(
partial(self._call_dataset_fn, dataset_fn), self._worker_device_map,
- self._prefetch_on_device)
+ self._prefetch_on_device, self._auto_shard_dataset)
else:
return values.PerDeviceDataset(
- self._call_dataset_fn(dataset_fn),
- self._devices,
- self._prefetch_on_device,
- source_device=device_util.resolve("/device:CPU:0"))
+ self._call_dataset_fn(dataset_fn), self._devices,
+ self._prefetch_on_device)
# TODO(priyag): Deal with OutOfRange errors once b/111349762 is fixed.
def _run_steps_on_dataset(self, fn, iterator, iterations,
@@ -625,9 +627,11 @@
return self._get_cross_tower_ops().batch_reduce(aggregation,
value_destination_pairs)
- def _update(self, var, fn, *args, **kwargs):
+ def _update(self, var, options, fn, *args, **kwargs):
# TODO(josh11b): In eager mode, use one thread per device.
assert isinstance(var, values.DistributedVariable)
+ should_group = options.pop("grouped")
+ assert not options # Validate that we are processing all of the options.
updates = {}
for d, v in var._index.items(): # pylint: disable=protected-access
name = "update_%d" % self._device_index.get(d)
@@ -636,10 +640,12 @@
updates[d] = fn(v,
*values.select_device_mirrored(d, args),
**values.select_device_mirrored(d, kwargs))
- return values.regroup(updates, values.Mirrored)
+ return values.update_regroup(self, updates, should_group)
- def _update_non_slot(self, colocate_with, fn, *args, **kwargs):
+ def _update_non_slot(self, colocate_with, options, fn, *args, **kwargs):
assert isinstance(colocate_with, list)
+ should_group = options.pop("grouped")
+ assert not options # Validate that we are processing all of the options.
# TODO(josh11b): In eager mode, use one thread per device.
updates = {}
for d in colocate_with:
@@ -647,7 +653,7 @@
with ops.device(d), distribute_lib.UpdateContext(d), ops.name_scope(name):
updates[d] = fn(*values.select_device_mirrored(d, args),
**values.select_device_mirrored(d, kwargs))
- return values.regroup(updates, values.Mirrored)
+ return values.update_regroup(self, updates, should_group)
def read_var(self, tower_local_var):
"""Read the aggregate value of a tower-local variable."""
diff --git a/tensorflow/contrib/distribute/python/mirrored_strategy_multigpu_test.py b/tensorflow/contrib/distribute/python/mirrored_strategy_multigpu_test.py
index 04c712c..eeac528 100644
--- a/tensorflow/contrib/distribute/python/mirrored_strategy_multigpu_test.py
+++ b/tensorflow/contrib/distribute/python/mirrored_strategy_multigpu_test.py
@@ -300,15 +300,9 @@
dist = mirrored_strategy.MirroredStrategy(
["/device:GPU:0", "/device:CPU:0"])
- ds = dist.distribute_dataset(
- lambda: dataset_ops.Dataset.from_tensors([[1.]]).repeat(10))
- if context.executing_eagerly():
- iterator = ds.make_one_shot_iterator()
- else:
- iterator = ds.make_initializable_iterator()
- self.evaluate([iterator.initializer])
-
- features = iterator.get_next()
+ features = dist.distribute_dataset(
+ lambda: dataset_ops.Dataset.from_tensors([[1.]]).repeat(10)
+ ).make_one_shot_iterator().get_next()
with dist.scope():
result = dist.call_for_each_tower(
@@ -832,7 +826,7 @@
with dist.scope():
ret_v_sum = dist.call_for_each_tower(model_fn, run_concurrently=False)
- update_ops = dist.unwrap(dist.update(ret_v_sum, update, 5.0))
+ update_ops = dist.update(ret_v_sum, update, 5.0, grouped=False)
# Initialize variables.
self.evaluate(variables.global_variables_initializer())
diff --git a/tensorflow/contrib/distribute/python/monitor.py b/tensorflow/contrib/distribute/python/monitor.py
index 17b7ab7..7644ace 100644
--- a/tensorflow/contrib/distribute/python/monitor.py
+++ b/tensorflow/contrib/distribute/python/monitor.py
@@ -51,7 +51,6 @@
else:
if session is None:
raise ValueError("Should provide a `session` in Graph mode.")
- session.run(step_callable._iterator.initializer) # pylint: disable=protected-access
self._run_step = session.make_callable(step_callable())
session.run(variables.global_variables_initializer())
diff --git a/tensorflow/contrib/distribute/python/one_device_strategy.py b/tensorflow/contrib/distribute/python/one_device_strategy.py
index 23b220f..f525919 100644
--- a/tensorflow/contrib/distribute/python/one_device_strategy.py
+++ b/tensorflow/contrib/distribute/python/one_device_strategy.py
@@ -141,14 +141,21 @@
else:
assert False
- def _update(self, var, fn, *args, **kwargs):
- with ops.device(self._device), distribute_lib.UpdateContext(self._device):
- return fn(var, *args, **kwargs)
+ def _update(self, var, options, fn, *args, **kwargs):
+ # The implementations of _update() and _update_non_slot() are identical
+ # except _update() passes `var` as the first argument to `fn()`.
+ return self._update_non_slot(var, options, fn, var, *args, **kwargs)
- def _update_non_slot(self, colocate_with, fn, *args, **kwargs):
+ def _update_non_slot(self, colocate_with, options, fn, *args, **kwargs):
del colocate_with
+ should_group = options.pop("grouped")
+ assert not options # Validate that we are processing all of the options.
with ops.device(self._device), distribute_lib.UpdateContext(self._device):
- return fn(*args, **kwargs)
+ result = fn(*args, **kwargs)
+ if should_group:
+ return result
+ else:
+ return nest.map_structure(self._unwrap, result)
def read_var(self, tower_local_var):
"""Read the aggregate value of a tower-local variable."""
diff --git a/tensorflow/contrib/distribute/python/optimizer_v2_test.py b/tensorflow/contrib/distribute/python/optimizer_v2_test.py
index 3064433..6e9ba37 100644
--- a/tensorflow/contrib/distribute/python/optimizer_v2_test.py
+++ b/tensorflow/contrib/distribute/python/optimizer_v2_test.py
@@ -42,11 +42,8 @@
model_fn, dataset_fn, layer = minimize_loss_example(
optimizer_fn, use_bias=True, use_callable_loss=use_callable_loss)
- ds = distribution.distribute_dataset(dataset_fn)
- if context.executing_eagerly():
- iterator = ds.make_one_shot_iterator()
- else:
- iterator = ds.make_initializable_iterator()
+ iterator = distribution.distribute_dataset(
+ dataset_fn).make_one_shot_iterator()
def run_step():
return control_flow_ops.group(distribution.unwrap(
@@ -55,7 +52,6 @@
if not context.executing_eagerly():
with self.cached_session() as sess:
- sess.run(iterator.initializer)
run_step = sess.make_callable(run_step())
self.evaluate(variables.global_variables_initializer())
diff --git a/tensorflow/contrib/distribute/python/parameter_server_strategy.py b/tensorflow/contrib/distribute/python/parameter_server_strategy.py
index 1125d02..6ddd915 100644
--- a/tensorflow/contrib/distribute/python/parameter_server_strategy.py
+++ b/tensorflow/contrib/distribute/python/parameter_server_strategy.py
@@ -343,21 +343,33 @@
return nest.map_structure(_select_fn, structured)
- def _update(self, var, fn, *args, **kwargs):
+ def _update(self, var, options, fn, *args, **kwargs):
if isinstance(var, values.AggregatingVariable):
var = var.get()
if not isinstance(var, resource_variable_ops.ResourceVariable):
raise ValueError(
"You can not update `var` %r. It must be a Variable." % var)
+ should_group = options.pop("grouped")
+ assert not options # Validate that we are processing all of the options.
with ops.colocate_with(var), distribute_lib.UpdateContext(var.device):
- return fn(var, *self._select_single_value(args),
- **self._select_single_value(kwargs))
+ result = fn(var, *self._select_single_value(args),
+ **self._select_single_value(kwargs))
+ if should_group:
+ return result
+ else:
+ return nest.map_structure(self._unwrap, result)
# TODO(yuefengz): does it need to call _select_single_value?
- def _update_non_slot(self, colocate_with, fn, *args, **kwargs):
+ def _update_non_slot(self, colocate_with, options, fn, *args, **kwargs):
+ should_group = options.pop("grouped")
+ assert not options # Validate that we are processing all of the options.
with ops.device(
colocate_with.device), distribute_lib.UpdateContext(colocate_with):
- return fn(*args, **kwargs)
+ result = fn(*args, **kwargs)
+ if should_group:
+ return result
+ else:
+ return nest.map_structure(self._unwrap, result)
def _unwrap(self, val):
if isinstance(val, values.DistributedValues):
diff --git a/tensorflow/contrib/distribute/python/parameter_server_strategy_test.py b/tensorflow/contrib/distribute/python/parameter_server_strategy_test.py
index 12789e0..353d11a 100644
--- a/tensorflow/contrib/distribute/python/parameter_server_strategy_test.py
+++ b/tensorflow/contrib/distribute/python/parameter_server_strategy_test.py
@@ -395,7 +395,8 @@
# TODO(yuefengz): support non-Mirrored variable as destinations.
g = d.reduce(
variable_scope.VariableAggregation.SUM, g, destinations=v)
- with ops.control_dependencies(d.unwrap(d.update(v, update, g))):
+ with ops.control_dependencies(
+ d.update(v, update, g, grouped=False)):
after_list.append(d.read_var(v))
return before_list, after_list
diff --git a/tensorflow/contrib/distribute/python/prefetching_ops_v2.py b/tensorflow/contrib/distribute/python/prefetching_ops_v2.py
new file mode 100644
index 0000000..d48aa9c
--- /dev/null
+++ b/tensorflow/contrib/distribute/python/prefetching_ops_v2.py
@@ -0,0 +1,232 @@
+# 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.
+# ==============================================================================
+"""Extension of prefetching_ops to support more than one device."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import warnings
+
+from tensorflow.python.data.experimental.ops import prefetching_ops
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import iterator_ops
+from tensorflow.python.data.util import nest as data_nest
+from tensorflow.python.data.util import sparse
+from tensorflow.python.eager import context
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import function
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
+from tensorflow.python.util import nest
+
+
+# pylint: disable=protected-access
+class _PrefetchToDeviceIterator(object):
+ """A replacement for `tf.data.Iterator` that prefetches to another device.
+
+ Args:
+ input_dataset: The input dataset.
+ one_shot: If true, we make a one shot iterator that's already initialized.
+ devices: Devices on which to prefetch.
+ buffer_size: Size of the prefetching buffer.
+ shared_name: (Optional.) If non-empty, the returned iterator will be shared
+ under the given name across multiple sessions that share the same devices
+ (e.g. when using a remote server). Only used if one_shot is False.
+
+ Returns:
+ An Iterator type object.
+ """
+
+ def __init__(self,
+ input_dataset,
+ one_shot,
+ devices,
+ buffer_size,
+ shared_name=None):
+ self._input_dataset = input_dataset
+ self._get_next_call_count = 0
+ self._one_shot = one_shot
+ if shared_name is None:
+ shared_name = ""
+ self._devices = devices
+
+ if self._one_shot:
+ self._input_iterator = input_dataset.make_one_shot_iterator()
+ else:
+ self._input_iterator = iterator_ops.Iterator.from_structure(
+ self._input_dataset.output_types, self._input_dataset.output_shapes,
+ shared_name, self._input_dataset.output_classes)
+ input_iterator_handle = self._input_iterator.string_handle()
+
+ @function.Defun(dtypes.string)
+ def _prefetch_fn(handle):
+ """Prefetches one element from `input_iterator`."""
+ remote_iterator = iterator_ops.Iterator.from_string_handle(
+ handle, self._input_iterator.output_types,
+ self._input_iterator.output_shapes,
+ self._input_iterator.output_classes)
+ ret = remote_iterator.get_next()
+ return nest.flatten(sparse.serialize_sparse_tensors(ret))
+
+ target_device = ged_ops.experimental_iterator_get_device(
+ self._input_iterator._iterator_resource)
+ self._buffering_resources = []
+ for device in nest.flatten(self._devices):
+ with ops.device(device):
+ buffer_resource_handle = prefetching_ops.function_buffering_resource(
+ f=_prefetch_fn,
+ output_types=data_nest.flatten(
+ sparse.as_dense_types(self._input_dataset.output_types,
+ self._input_dataset.output_classes)),
+ target_device=target_device,
+ string_arg=input_iterator_handle,
+ buffer_size=buffer_size,
+ shared_name=shared_name)
+ self._buffering_resources.append(buffer_resource_handle)
+
+ if not self._one_shot:
+ reset_ops = []
+ for buffer_resource in self._buffering_resources:
+ reset_ops.append(
+ ged_ops.experimental_function_buffering_resource_reset(
+ buffer_resource))
+ with ops.control_dependencies(reset_ops):
+ self._initializer = self._input_iterator.make_initializer(
+ self._input_dataset)
+
+ def get_next(self, name=None):
+ """See `tf.data.Iterator.get_next`."""
+ self._get_next_call_count += 1
+ if self._get_next_call_count > iterator_ops.GET_NEXT_CALL_WARNING_THRESHOLD:
+ warnings.warn(iterator_ops.GET_NEXT_CALL_WARNING_MESSAGE)
+
+ flat_result = []
+ # TODO(priyag): This will fail if the input size (typically number of
+ # batches) is not divisible by number of devices.
+ # How do we handle that more gracefully / let the user know?
+ for buffer_resource in self._buffering_resources:
+ flat_ret = ged_ops.experimental_function_buffering_resource_get_next(
+ buffer_resource,
+ output_types=data_nest.flatten(
+ sparse.as_dense_types(self.output_types, self.output_classes)),
+ name=name)
+
+ ret = sparse.deserialize_sparse_tensors(
+ data_nest.pack_sequence_as(self.output_types, flat_ret),
+ self.output_types, self.output_shapes, self.output_classes)
+
+ for tensor, shape in zip(
+ data_nest.flatten(ret), data_nest.flatten(self.output_shapes)):
+ if isinstance(tensor, ops.Tensor):
+ tensor.set_shape(shape)
+ flat_result.append(ret)
+
+ return nest.pack_sequence_as(self._devices, flat_result)
+
+ @property
+ def initializer(self):
+ if self._one_shot:
+ raise NotImplementedError("Can't initialize a one_shot_iterator")
+ return self._initializer
+
+ @property
+ def output_classes(self):
+ return self._input_dataset.output_classes
+
+ @property
+ def output_shapes(self):
+ return self._input_dataset.output_shapes
+
+ @property
+ def output_types(self):
+ return self._input_dataset.output_types
+
+
+# pylint: enable=protected-access
+
+
+class _PrefetchToDeviceDataset(dataset_ops.UnaryDataset):
+ """A `Dataset` whose iterator prefetches elements to other device(s)."""
+
+ def __init__(self, input_dataset, devices, buffer_size):
+ super(_PrefetchToDeviceDataset, self).__init__(input_dataset)
+ self._input_dataset = input_dataset
+ self._devices = devices
+ self._buffer_size = buffer_size if buffer_size is not None else 1
+
+ def make_one_shot_iterator(self):
+ return _PrefetchToDeviceIterator(
+ self._input_dataset,
+ one_shot=True,
+ devices=self._devices,
+ buffer_size=self._buffer_size)
+
+ def make_initializable_iterator(self, shared_name=None):
+ if context.executing_eagerly():
+ raise RuntimeError(
+ "make_initializable_iterator is not supported when eager "
+ "execution is enabled.")
+
+ return _PrefetchToDeviceIterator(
+ self._input_dataset,
+ one_shot=False,
+ devices=self._devices,
+ buffer_size=self._buffer_size,
+ shared_name=shared_name)
+
+ def _as_variant_tensor(self):
+ # TODO(mrry): Raise this error earlier (e.g. when one of the Dataset
+ # transformation methods is called.
+ # TODO(mrry): Investigate support for chaining further transformations after
+ # the prefetch, including GPU support.
+ raise NotImplementedError("`prefetch_to_devices()` must be the last "
+ "transformation in a dataset pipeline.")
+
+ # TODO(priyag): Fix the output types, shapes and classes to match the result
+ # of get_next (which has the additional nesting layer of devices now).
+ @property
+ def output_types(self):
+ return self._input_dataset.output_types
+
+ @property
+ def output_shapes(self):
+ return self._input_dataset.output_shapes
+
+ @property
+ def output_classes(self):
+ return self._input_dataset.output_classes
+
+
+def prefetch_to_devices(devices, buffer_size=None):
+ """A transformation that prefetches dataset values to the given `devices`.
+
+ NOTE: Although the transformation creates a `tf.data.Dataset`, the
+ transformation must be the final `Dataset` in the input pipeline.
+
+ Args:
+ devices: A nested structure of devices on which to prefetch the data. It can
+ be a single device name, or a tuple or list of device names.
+ buffer_size: (Optional.) The number of elements to buffer on each device.
+ Defaults to an automatically chosen value.
+
+ Returns:
+ A `Dataset` transformation function, which can be passed to
+ `tf.data.Dataset.apply`.
+ """
+
+ def _apply_fn(dataset):
+ return _PrefetchToDeviceDataset(dataset, devices, buffer_size)
+
+ return _apply_fn
diff --git a/tensorflow/contrib/distribute/python/prefetching_ops_v2_test.py b/tensorflow/contrib/distribute/python/prefetching_ops_v2_test.py
new file mode 100644
index 0000000..1679910
--- /dev/null
+++ b/tensorflow/contrib/distribute/python/prefetching_ops_v2_test.py
@@ -0,0 +1,90 @@
+# 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 prefetching_ops_v2."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.contrib.distribute.python import prefetching_ops_v2
+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
+
+
+class PrefetchingOpsV2Test(test.TestCase):
+
+ def testPrefetchToOneDevice(self):
+ if not test_util.is_gpu_available():
+ self.skipTest("No GPU available")
+
+ host_dataset = dataset_ops.Dataset.range(10)
+ device_dataset = host_dataset.apply(
+ prefetching_ops_v2.prefetch_to_devices("/gpu:0"))
+
+ iterator = device_dataset.make_one_shot_iterator()
+ next_element = iterator.get_next()
+
+ with self.cached_session() as sess:
+ for i in range(10):
+ self.assertEqual(i, sess.run(next_element))
+ with self.assertRaises(errors.OutOfRangeError):
+ sess.run(next_element)
+
+ def testPrefetchToTwoDevicesInAList(self):
+ if not test_util.is_gpu_available():
+ self.skipTest("No GPU available")
+
+ host_dataset = dataset_ops.Dataset.range(10)
+ device_dataset = host_dataset.apply(
+ prefetching_ops_v2.prefetch_to_devices(["/cpu:0", "/gpu:0"]))
+
+ iterator = device_dataset.make_one_shot_iterator()
+ next_element = iterator.get_next()
+
+ output = []
+ # TODO(rohanj): Modify test to go till the end of the dataset when we
+ # switch to MultiDeviceIterator.
+ with self.cached_session() as sess:
+ for _ in range(4):
+ result = sess.run(next_element)
+ self.assertEqual(2, len(result))
+ output.extend(result)
+ self.assertEquals(set(range(8)), set(output))
+
+ def testPrefetchToTwoDevicesWithReinit(self):
+ if not test_util.is_gpu_available():
+ self.skipTest("No GPU available")
+
+ host_dataset = dataset_ops.Dataset.range(10)
+ device_dataset = host_dataset.apply(
+ prefetching_ops_v2.prefetch_to_devices(["/cpu:0", "/gpu:0"]))
+
+ iterator = device_dataset.make_initializable_iterator()
+ next_element = iterator.get_next()
+
+ # TODO(rohanj): Modify test to go till the end of the dataset when we
+ # switch to MultiDeviceIterator.
+ with self.cached_session() as sess:
+ sess.run(iterator.initializer)
+ for _ in range(4):
+ sess.run(next_element)
+ sess.run(iterator.initializer)
+ for _ in range(4):
+ sess.run(next_element)
+
+
+if __name__ == "__main__":
+ test.main()
diff --git a/tensorflow/contrib/distribute/python/step_fn.py b/tensorflow/contrib/distribute/python/step_fn.py
index 23bf361..1b5a4f6 100644
--- a/tensorflow/contrib/distribute/python/step_fn.py
+++ b/tensorflow/contrib/distribute/python/step_fn.py
@@ -19,7 +19,6 @@
from __future__ import print_function
from tensorflow.python.eager import backprop
-from tensorflow.python.eager import context
from tensorflow.python.training import optimizer as optimizer_lib
@@ -51,11 +50,7 @@
def __init__(self, dataset_fn, distribution):
super(StandardInputStep, self).__init__(distribution)
self._distributed_input = distribution.distribute_dataset(dataset_fn)
- if context.executing_eagerly():
- self._iterator = self._distributed_input.make_one_shot_iterator()
- else:
- # TODO(priyag): Expose initializer via some initializer property.
- self._iterator = self._distributed_input.make_initializable_iterator()
+ self._iterator = self._distributed_input.make_one_shot_iterator()
class StandardSingleLossStep(StandardInputStep):
diff --git a/tensorflow/contrib/distribute/python/step_fn_test.py b/tensorflow/contrib/distribute/python/step_fn_test.py
index 1ff9b9c..f1ada49 100644
--- a/tensorflow/contrib/distribute/python/step_fn_test.py
+++ b/tensorflow/contrib/distribute/python/step_fn_test.py
@@ -50,7 +50,6 @@
run_step = single_loss_step
else:
with self.cached_session() as sess:
- sess.run(single_loss_step._iterator.initializer)
run_step = sess.make_callable(single_loss_step())
self.evaluate(variables.global_variables_initializer())
diff --git a/tensorflow/contrib/distribute/python/strategy_test_lib.py b/tensorflow/contrib/distribute/python/strategy_test_lib.py
index 5d498fb..fd280f5 100644
--- a/tensorflow/contrib/distribute/python/strategy_test_lib.py
+++ b/tensorflow/contrib/distribute/python/strategy_test_lib.py
@@ -115,7 +115,8 @@
with ops.control_dependencies([fetched]):
g = d.reduce(
variable_scope.VariableAggregation.SUM, g, destinations=v)
- with ops.control_dependencies(d.unwrap(d.update(v, update, g))):
+ with ops.control_dependencies(d.update(
+ v, update, g, grouped=False)):
after_list.append(d.read_var(v))
return before_list, after_list
@@ -169,7 +170,8 @@
with ops.control_dependencies([fetched]):
g = d.reduce(
variable_scope.VariableAggregation.SUM, g, destinations=v)
- with ops.control_dependencies(d.unwrap(d.update(v, update, g))):
+ with ops.control_dependencies(d.update(
+ v, update, g, grouped=False)):
after_list.append(d.read_var(v))
return before_list, after_list
diff --git a/tensorflow/contrib/distribute/python/tpu_strategy.py b/tensorflow/contrib/distribute/python/tpu_strategy.py
index a6762e5..c3c7df3 100644
--- a/tensorflow/contrib/distribute/python/tpu_strategy.py
+++ b/tensorflow/contrib/distribute/python/tpu_strategy.py
@@ -29,6 +29,7 @@
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.python.eager import context
+from tensorflow.python.eager import tape
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
@@ -37,9 +38,13 @@
from tensorflow.python.ops import variable_scope as vs
from tensorflow.python.ops import variables as variables_lib
from tensorflow.python.training import device_util
+from tensorflow.python.training import distribute as distribute_lib
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()
@@ -56,6 +61,58 @@
return tpu_system_metadata
+# TODO(jhseu): Deduplicate with MirroredStrategy?
+def _create_tpu_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 TPUMirroredVariable to those collections instead.
+ collections = kwargs.pop("collections", None)
+ if collections is None:
+ collections = [ops.GraphKeys.GLOBAL_VARIABLES]
+ kwargs["collections"] = []
+
+ # TODO(jhseu): Should we have different behavior for different
+ # synchronization settings?
+
+ # Get aggregation value
+ # TODO(jhseu): Support aggregation in a tower context.
+ aggregation = kwargs.pop("aggregation", vs.VariableAggregation.NONE)
+ if aggregation not in [
+ vs.VariableAggregation.NONE,
+ vs.VariableAggregation.SUM,
+ vs.VariableAggregation.MEAN,
+ vs.VariableAggregation.ONLY_FIRST_TOWER,
+ ]:
+ raise ValueError("Invalid variable aggregation mode: {} for variable: {}"
+ .format(aggregation, kwargs["name"]))
+
+ # Ignore user-specified caching device, not needed for mirrored variables.
+ kwargs.pop("caching_device", None)
+
+ # TODO(josh11b,apassos): It would be better if variable initialization
+ # was never recorded on the tape instead of having to do this manually
+ # here.
+ with tape.stop_recording():
+ index = real_mirrored_creator(devices, *args, **kwargs)
+ result = values.TPUMirroredVariable(index, index[devices[0]], aggregation)
+
+ if not context.executing_eagerly():
+ g = ops.get_default_graph()
+ # If "trainable" is True, next_creator() will add the member variables
+ # to the TRAINABLE_VARIABLES collection, so we manually remove
+ # them and replace with the MirroredVariable. We can't set
+ # "trainable" to False for next_creator() since that causes functions
+ # like implicit_gradients to skip those variables.
+ if kwargs.get("trainable", True):
+ collections.append(ops.GraphKeys.TRAINABLE_VARIABLES)
+ l = g.get_collection_ref(ops.GraphKeys.TRAINABLE_VARIABLES)
+ for v in index.values():
+ l.remove(v)
+ g.add_to_collections(collections, result)
+ return result
+
+
+# TODO(jhseu): Stop inheriting from OneDeviceStrategy.
class TPUStrategy(one_device_strategy.OneDeviceStrategy):
"""Experimental TPU distribution strategy implementation."""
@@ -82,6 +139,15 @@
# TODO(sourabhbajaj): Change this from num_cores to metadata_override
self._num_cores_override = num_cores
+ # TODO(jhseu): Switch to DeviceAssignment to support pods and model
+ # parallelism.
+ device_map = {d.name: i for i, d in enumerate(self._tpu_metadata.devices)
+ if "device:TPU:" in d.name}
+ self._device_index = values.PerDevice(device_map)
+ self._tpu_devices = sorted(device_map.keys())
+ # Only create variables for the number of towers we're running.
+ self._tpu_devices = self._tpu_devices[:self.num_towers]
+
# 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
@@ -231,6 +297,7 @@
# For outputs that have already been aggregated, take the first value
# from the list as each value should be the same. Else return the full
# list of values.
+ # TODO(josh11b): If aggregation is NONE, we should return a PerDevice value.
if aggregation is not variables_lib.VariableAggregation.NONE:
# TODO(priyag): Should this return the element or a list with 1 element
last_step_tensor_outputs_dict[name] = output[0]
@@ -239,6 +306,8 @@
return ctx
def _call_for_each_tower(self, fn, *args, **kwargs):
+ # TODO(jhseu): Consider making it so call_for_each_tower implies that we're
+ # in a tpu.rewrite(), and update TPUMirroredVariable accordingly.
kwargs.pop('run_concurrently', None)
with one_device_strategy._OneDeviceTowerContext(self): # pylint: disable=protected-access
return fn(*args, **kwargs)
@@ -248,7 +317,15 @@
# TODO(priyag): Add appopriate call here when eager is supported for TPUs.
raise NotImplementedError('Eager mode not supported in TPUStrategy.')
else:
- return [tpu.initialize_system()]
+ # 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)
def finalize(self):
if context.executing_eagerly():
@@ -257,21 +334,53 @@
else:
return [tpu.shutdown_system()]
+ def _get_devices_from(self, colocate_with=None):
+ # TODO(jhseu): Change this when we support model parallelism.
+ return self._tpu_devices
+
+ def _create_variable(self, next_creator, *args, **kwargs):
+ """Create a TPUMirroredVariable. See `DistributionStrategy.scope`."""
+ colocate_with = kwargs.pop("colocate_with", None)
+ devices = self._get_devices_from(colocate_with)
+
+ def _real_mirrored_creator(devices, *args, **kwargs): # pylint: disable=g-missing-docstring
+ index = {}
+ for i, d in enumerate(devices):
+ with ops.device(d):
+ if i > 0:
+ # Give replicas meaningful distinct names:
+ var0name = index[devices[0]].name.split(":")[0]
+ # We append a / to variable names created on towers with id > 0 to
+ # ensure that we ignore the name scope and instead use the given
+ # name as the absolute name of the variable.
+ kwargs["name"] = "%s/replica_%d/" % (var0name, i)
+ # Initialize replicas with the same value:
+ if context.executing_eagerly():
+ kwargs["initial_value"] = array_ops.identity(
+ index[devices[0]].value())
+ else:
+ def initial_value_fn(device=d):
+ with ops.device(device):
+ return array_ops.identity(index[devices[0]].initial_value)
+ kwargs["initial_value"] = initial_value_fn
+ with context.context().device_policy(context.DEVICE_PLACEMENT_SILENT):
+ v = next_creator(*args, **kwargs)
+ assert not isinstance(v, values.TPUMirroredVariable)
+ index[d] = v
+ return index
+
+ return _create_tpu_mirrored_variable(devices, _real_mirrored_creator, *args,
+ **kwargs)
+
def _reduce(self, aggregation, value, destinations):
- graph = ops.get_default_graph()
- cf_context = graph._get_control_flow_context() # pylint: disable=protected-access
- # If we're inside the ReplicateContext, reduction should be done using
- # CrossReplicaSum while outside we can directly use an add_n op.
- while cf_context:
- if isinstance(cf_context, tpu.TPUReplicateContext):
- if aggregation == vs.VariableAggregation.MEAN:
- # TODO(jhseu): Revisit once we support model-parallelism.
- value *= (1. / self.num_towers)
- elif aggregation != vs.VariableAggregation.SUM:
- raise NotImplementedError(
- 'Currently only support sum & mean in TPUStrategy.')
- return tpu_ops.cross_replica_sum(value)
- cf_context = cf_context.outer_context
+ if values._enclosing_tpu_context() is not None: # pylint: disable=protected-access
+ if aggregation == vs.VariableAggregation.MEAN:
+ # TODO(jhseu): Revisit once we support model-parallelism.
+ value *= (1. / self.num_towers)
+ elif aggregation != vs.VariableAggregation.SUM:
+ raise NotImplementedError(
+ "Currently only support sum & mean in TPUStrategy.")
+ return tpu_ops.cross_replica_sum(value)
# Validate that the destination is same as the host device
# Note we don't do this when in replicate context as the reduction is
@@ -290,10 +399,46 @@
return output * (1. / len(value))
return output
- def _unwrap(self, value):
- if isinstance(value, list):
- return value
- return [value]
+ def _update(self, var, options, fn, *args, **kwargs):
+ assert isinstance(var, values.TPUMirroredVariable)
+ should_group = options.pop("grouped")
+ assert not options # Validate that we are processing all of the options.
+
+ if values._enclosing_tpu_context() is not None: # pylint: disable=protected-access
+ if should_group:
+ return fn(var, *args, **kwargs)
+ else:
+ return [fn(var, *args, **kwargs)]
+
+ # Otherwise, we revert to MirroredStrategy behavior and update each variable
+ # directly.
+ updates = {}
+ for d, v in var._index.items(): # pylint: disable=protected-access
+ name = "update_%d" % self._device_index.get(d)
+ with ops.device(d), distribute_lib.UpdateContext(d), ops.name_scope(name):
+ # If args and kwargs are not mirrored, the value is returned as is.
+ updates[d] = fn(v,
+ *values.select_device_mirrored(d, args),
+ **values.select_device_mirrored(d, kwargs))
+ return values.update_regroup(self, updates, should_group)
+
+ # TODO(josh11b): Need to implement _update_non_slot()!
+
+ def read_var(self, var):
+ assert isinstance(var, values.TPUMirroredVariable)
+ return var.read_value()
+
+ 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)]
+ elif isinstance(val, list):
+ # TODO(josh11b): We need to remove this case; per device values should
+ # be represented using a PerDevice wrapper instead of a list with
+ # one entry per device.
+ return val
+ return [val]
+
@property
def num_towers(self):
@@ -323,6 +468,14 @@
def should_save_summary(self):
return True
+ @property
+ def worker_devices(self):
+ return self._tpu_devices
+
+ @property
+ def parameter_devices(self):
+ return self._tpu_devices
+
def get_host_cpu_device(self, host_id):
if self._tpu_cluster_resolver.get_master() in ('', 'local'):
return '/replica:0/task:0/device:CPU:0'
diff --git a/tensorflow/contrib/distribute/python/values.py b/tensorflow/contrib/distribute/python/values.py
index a0cd029..18ceba4 100644
--- a/tensorflow/contrib/distribute/python/values.py
+++ b/tensorflow/contrib/distribute/python/values.py
@@ -22,17 +22,20 @@
from __future__ import print_function
import collections
+import contextlib
import weakref
import six
from tensorflow.contrib.distribute.python import input_ops
-from tensorflow.python.data.ops import multi_device_iterator_ops
+from tensorflow.contrib.distribute.python import prefetching_ops_v2
from tensorflow.python.eager import context
+from tensorflow.python.eager import tape
from tensorflow.python.framework import device as tf_device
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_ops
+from tensorflow.python.ops import gen_resource_variable_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import variable_scope as vs
from tensorflow.python.ops import variables as variables_lib
@@ -363,18 +366,7 @@
# We are calling assign on the mirrored variable in cross tower context,
# use update to update the variable.
strategy = distribution_strategy_context.get_distribution_strategy()
- updates = strategy.update(self, f, *args, **kwargs)
- grouped = strategy.group(updates)
- if isinstance(updates, DistributedValues) and updates.is_tensor_like:
- # Make sure we run all updates. Without this, something like
- # session.run(mirrored_var.assign*(...)) may only update one tower.
- index = {}
- for d in updates.devices:
- with ops.device(d), ops.control_dependencies([grouped]):
- index[d] = array_ops.identity(updates.get(d))
- return Mirrored(index)
- else:
- return grouped
+ return strategy.update(self, f, *args, **kwargs)
else:
_assert_tower_context()
# We are calling an assign function on the mirrored variable in tower
@@ -453,6 +445,384 @@
_tensor_conversion_mirrored)
+def _enclosing_tpu_context():
+ # pylint: disable=protected-access
+ tpu_context = ops.get_default_graph()._get_control_flow_context()
+ # pylint: enable=protected-access
+ while tpu_context is not None and not isinstance(
+ tpu_context, control_flow_ops.XLAControlFlowContext):
+ tpu_context = tpu_context.outer_context
+ return tpu_context
+
+
+# TODO(jhseu): Deduplicate code. We copy code because we don't want to
+# inherit from DistributedDelegate. DistributedDelegate will not work in a
+# tpu.replicate() because it assumes that you're in a device context where you
+# can operate on a single version of the variable, but a tpu.replicate()
+# operates on all variables and is replicated during a rewrite pass.
+class TPUMirroredVariable(checkpointable.CheckpointableBase):
+ """Holds a map from device to TPU variables whose values are kept in sync."""
+
+ def __init__(self, index, primary_var, aggregation):
+ # Use a weakref to make it easy to map from the contained values
+ # to the container without introducing a reference cycle.
+ for v in six.itervalues(index):
+ v._mirrored_container = weakref.ref(self) # pylint: disable=protected-access
+ self._index = {device_util.canonicalize(key): value
+ for key, value in six.iteritems(index)}
+ self._primary_var = primary_var
+ self._common_name = self._primary_var.name.split(":")[0]
+ self._aggregation = aggregation
+ # Needed for GradientTape
+ self._trainable = self._primary_var.trainable
+
+ def _get(self, device=None):
+ """Returns the value for the current device or raises a ValueError."""
+ if device is None:
+ tower_context = distribution_strategy_context.get_tower_context()
+ if tower_context:
+ device = tower_context.device
+ else:
+ device = distribute_lib.get_update_device()
+ if device is None:
+ return self._get_cross_tower()
+ device = device_util.canonicalize(device)
+ try:
+ return self._index[device]
+ except KeyError as e:
+ six.raise_from(
+ ValueError("Device %s not found in %s (current device %s)" %
+ (device, self._index.keys(), device_util.current())), e)
+
+ # pylint: disable=multiple-statements
+ def __add__(self, o): return self.read_value() + o
+ def __radd__(self, o): return o + self.read_value()
+ def __sub__(self, o): return self.read_value() - o
+ def __rsub__(self, o): return o - self.read_value()
+ def __mul__(self, o): return self.read_value() * o
+ def __rmul__(self, o): return o * self.read_value()
+ def __truediv__(self, o): return self.read_value() / o
+ def __rtruediv__(self, o): return o / self.read_value()
+ def __floordiv__(self, o): return self.read_value() // o
+ def __rfloordiv__(self, o): return o // self.read_value()
+ def __mod__(self, o): return self.read_value() % o
+ def __rmod__(self, o): return o % self.read_value()
+ def __lt__(self, o): return self.read_value() < o
+ def __le__(self, o): return self.read_value() <= o
+ def __gt__(self, o): return self.read_value() > o
+ def __ge__(self, o): return self.read_value() >= o
+ def __and__(self, o): return self.read_value() & o
+ def __rand__(self, o): return o & self.read_value()
+ def __or__(self, o): return self.read_value() | o
+ def __ror__(self, o): return o | self.read_value()
+ def __xor__(self, o): return self.read_value() ^ o
+ def __rxor__(self, o): return o ^ self.read_value()
+ def __getitem__(self, o): return self.read_value()[o]
+ def __pow__(self, o, modulo=None): return pow(self.read_value(), o, modulo)
+ def __rpow__(self, o): return pow(o, self.read_value())
+ def __invert__(self): return ~self.read_value()
+ def __neg__(self): return -self.read_value()
+ def __abs__(self): return abs(self.read_value())
+
+ def __div__(self, o):
+ try:
+ return self.read_value().__div__(o)
+ except AttributeError:
+ # See https://docs.python.org/3/library/constants.html#NotImplemented
+ return NotImplemented
+
+ def __rdiv__(self, o):
+ try:
+ return self.read_value().__rdiv__(o)
+ except AttributeError:
+ # See https://docs.python.org/3/library/constants.html#NotImplemented
+ return NotImplemented
+
+ def __matmul__(self, o):
+ try:
+ return self.read_value().__matmul__(o)
+ except AttributeError:
+ # See https://docs.python.org/3/library/constants.html#NotImplemented
+ return NotImplemented
+
+ def __rmatmul__(self, o):
+ try:
+ return self.read_value().__rmatmul__(o)
+ except AttributeError:
+ # See https://docs.python.org/3/library/constants.html#NotImplemented
+ return NotImplemented
+
+ @property
+ def handle(self):
+ # If we're in a tpu.rewrite(), return the replicated handle.
+ tpu_context = _enclosing_tpu_context()
+ if tpu_context is not None:
+ return tpu_context.get_replicated_var_handle(
+ self._common_name, nest.flatten(self._index))
+
+ device = distribute_lib.get_update_device()
+ if device is None:
+ return self._primary_var.handle
+ device = device_util.canonicalize(device)
+ try:
+ return self._index[device].handle
+ except KeyError as e:
+ six.raise_from(
+ ValueError("Device %s not found in %s (current device %s)" %
+ (device, self._index.keys(), device_util.current())), e)
+
+ # The arguments to update() are automatically unwrapped so the update()
+ # function would normally see regular variables, not MirroredVariables.
+ # However, the update function can still operate on wrapped MirroredVariables
+ # through object members, captured arguments, etc. This is more likely in an
+ # update_non_slot() function (like OptimizerV2._finish), which can
+ # update several non-slot variables in one call.
+ def _assign_func(self, *args, **kwargs):
+ if distribution_strategy_context.get_distribution_strategy().__class__.__name__ != "TPUStrategy":
+ raise ValueError("You may only assign to a TPUMirroredVariable within a "
+ "TPUStrategy.")
+ f = kwargs.pop("f")
+ if distribution_strategy_context.get_cross_tower_context():
+ if _enclosing_tpu_context() is not None:
+ return distribution_strategy_context.get_distribution_strategy().update(
+ self, f, *args, **kwargs)
+
+ update_device = distribute_lib.get_update_device()
+ # We are calling update on the mirrored variable in cross tower context.
+ if update_device is not None:
+ # We are calling an assign function on the mirrored variable in cross
+ # tower context.
+ v = self._get(device=update_device)
+ return f(v, *args, **kwargs)
+
+ return distribution_strategy_context.get_distribution_strategy().update(
+ self, f, *args, **kwargs)
+ else:
+ _assert_tower_context()
+ # We are calling an assign function on the mirrored variable in tower
+ # context.
+ # We reduce the value we want to assign/add/sub. More details about how we
+ # handle the different use cases can be found in the _reduce method.
+ # We call the function on each of the mirrored variables with the reduced
+ # value.
+ if self._aggregation == vs.VariableAggregation.NONE:
+ raise ValueError("You must specify an aggregation method to update a "
+ "TPUMirroredVariable in Tower Context.")
+
+ def merge_fn(strategy, value, *other_args, **other_kwargs):
+ return strategy.update(
+ self, f,
+ strategy.reduce(
+ aggregation=self._aggregation, value=value, destinations=self),
+ *other_args, **other_kwargs)
+
+ return distribution_strategy_context.get_tower_context().merge_call(
+ merge_fn, *args, **kwargs)
+
+ @contextlib.contextmanager
+ def _handle_graph(self, handle):
+ # Note: might have an eager tensor but not be executing eagerly when
+ # building functions.
+ if (context.executing_eagerly() or isinstance(handle, ops.EagerTensor)
+ or ops.has_default_graph()):
+ yield
+ else:
+ with handle.graph.as_default():
+ yield
+
+ @property
+ def trainable(self):
+ return self._trainable
+
+ def _read_variable_op(self, parent_op=None):
+ if self.trainable:
+ tape.variable_accessed(self)
+ if parent_op is not None:
+ with ops.control_dependencies([parent_op]):
+ return gen_resource_variable_ops.read_variable_op(
+ self.handle, self.dtype)
+
+ return gen_resource_variable_ops.read_variable_op(
+ self.handle, self.dtype)
+
+ def read_value(self):
+ return self._read_variable_op()
+
+ def assign_sub(self, *args, **kwargs):
+ def assign_sub_fn(var, delta, **kw):
+ name = kw.pop("name", None)
+ read_value = kw.pop("read_value", True)
+ with self._handle_graph(var.handle):
+ op = gen_resource_variable_ops.assign_sub_variable_op(
+ var.handle, ops.convert_to_tensor(delta, dtype=self.dtype),
+ name=name)
+ if read_value:
+ return self._read_variable_op(parent_op=op)
+ return op
+
+ return self._assign_func(f=assign_sub_fn, *args, **kwargs)
+
+ def assign_add(self, *args, **kwargs):
+ def assign_add_fn(var, delta, **kw):
+ name = kw.pop("name", None)
+ read_value = kw.pop("read_value", True)
+ with self._handle_graph(var.handle):
+ op = gen_resource_variable_ops.assign_add_variable_op(
+ var.handle, ops.convert_to_tensor(delta, dtype=self.dtype),
+ name=name)
+ if read_value:
+ return self._read_variable_op(parent_op=op)
+ return op
+
+ return self._assign_func(f=assign_add_fn, *args, **kwargs)
+
+ def assign(self, *args, **kwargs):
+ def assign_fn(var, value, **kw):
+ name = kw.pop("name", None)
+ read_value = kw.pop("read_value", True)
+ with self._handle_graph(var.handle):
+ op = gen_resource_variable_ops.assign_variable_op(
+ var.handle, ops.convert_to_tensor(value, dtype=self.dtype),
+ name=name)
+ if read_value:
+ return self._read_variable_op(parent_op=op)
+ return op
+
+ return self._assign_func(f=assign_fn, *args, **kwargs)
+
+ @property
+ def aggregation(self):
+ return self._aggregation
+
+ @property
+ def constraint(self):
+ return None
+
+ @property
+ def initializer(self):
+ return control_flow_ops.group(
+ [v.initializer for v in nest.flatten(self._index)])
+
+ @property
+ def graph(self):
+ return self._primary_var.graph
+
+ @property
+ def _shared_name(self):
+ return self._common_name
+
+ @property
+ def _unique_id(self):
+ return self._primary_var._unique_id # pylint: disable=protected-access
+
+ @property
+ def name(self):
+ return self._primary_var.name
+
+ @property
+ def dtype(self):
+ return self._primary_var.dtype
+
+ @property
+ def shape(self):
+ return self._primary_var.shape
+
+ def get_shape(self):
+ return self._primary_var.get_shape()
+
+ def to_proto(self, export_scope=None):
+ return self._primary_var.to_proto(export_scope=export_scope)
+
+ def _get_cross_tower(self):
+ device = device_util.canonicalize(device_util.current())
+ if device in self._index:
+ return self._index[device]
+ return self._primary_var
+
+ def _as_graph_element(self):
+ # pylint: disable=protected-access
+ if distribution_strategy_context.get_cross_tower_context():
+ return self._primary_var._as_graph_element()
+ return self._read_variable_op()
+
+ def _gather_saveables_for_checkpoint(self):
+ """Overrides CheckpointableBase method.
+
+ This allows both name-based and object-based save and restore of
+ MirroredVariables.
+
+ Returns:
+ A dictionary mapping attribute names to `SaveableObject` factories.
+ """
+ def _saveable_factory(name=self._common_name):
+ return _MirroredSaveable(self, self._primary_var, name)
+ return {checkpointable.VARIABLE_VALUE_KEY: _saveable_factory}
+
+ def _should_act_as_resource_variable(self):
+ """Pass resource_variable_ops.is_resource_variable check."""
+ pass
+
+ # Needed to pass ResourceVariable checks.
+ @property
+ def op(self):
+ return self._primary_var.op
+
+ @property
+ def _in_graph_mode(self):
+ return self._primary_var._in_graph_mode # pylint: disable=protected-access
+
+ def _dense_var_to_tensor(self, dtype=None, name=None, as_ref=False):
+ """Converts a variable to a tensor."""
+ # pylint: disable=protected-access
+ if _enclosing_tpu_context() is None:
+ return self._get()._dense_var_to_tensor(dtype, name, as_ref)
+ # pylint: enable=protected-access
+ if dtype is not None and dtype != self.dtype:
+ raise NotImplementedError
+ if as_ref:
+ return self.handle
+ else:
+ return self.read_value()
+
+ def is_initialized(self, name=None):
+ """Identifies if all the component variables are initialized.
+
+ Args:
+ name: Name of the final `logical_and` op.
+
+ Returns:
+ The op that evaluates to True or False depending on if all the
+ component variables are initialized.
+ """
+ # TODO(jhseu): Do we need TPU context implementation?
+
+ # We have to cast the self._index.values() to a `list` because when we
+ # use `model_to_estimator` to run tf.keras models, self._index.values() is
+ # of type `dict_values` and not `list`.
+ values_list = nest.flatten(self._index)
+ result = values_list[0].is_initialized()
+ # We iterate through the list of values except the last one to allow us to
+ # name the final `logical_and` op the same name that is passed by the user
+ # to the `is_initialized` op. For distributed variables, the
+ # `is_initialized` op is a `logical_and` op.
+ for v in values_list[1:-1]:
+ result = math_ops.logical_and(result, v.is_initialized())
+ result = math_ops.logical_and(result, values_list[-1].is_initialized(),
+ name=name)
+ return result
+
+
+# Register a conversion function which reads the value of the variable,
+# allowing instances of the class to be used as tensors.
+def _tensor_conversion_tpu_mirrored(var, dtype=None, name=None, as_ref=False):
+ return var._dense_var_to_tensor(dtype=dtype, name=name, as_ref=as_ref) # pylint: disable=protected-access
+
+
+ops.register_tensor_conversion_function(TPUMirroredVariable,
+ _tensor_conversion_tpu_mirrored)
+ops.register_dense_tensor_like_type(TPUMirroredVariable)
+
+
class _TowerLocalSaveable(saver.BaseSaverBuilder.SaveableObject):
"""Class for defining how to restore a TowerLocalVariable."""
@@ -668,6 +1038,29 @@
return nest.map_structure(_get_mirrored, structured)
+def update_regroup(strategy, updates, should_group):
+ """Regroup for an update, with dependencies to ensure all updates execute."""
+ regrouped = regroup(updates, Mirrored)
+ if not should_group:
+ return nest.map_structure(strategy.unwrap, regrouped)
+ grouped_flat = []
+ for u in nest.flatten(regrouped):
+ if isinstance(u, DistributedValues):
+ g = strategy.group(u)
+ if u.is_tensor_like:
+ # Make sure we run all updates. Without this, something like
+ # session.run(strategy.update(...)) may only update one tower.
+ index = {}
+ for d in u.devices:
+ with ops.device(d), ops.control_dependencies([g]):
+ index[d] = array_ops.identity(u.get(d))
+ g = Mirrored(index)
+ else:
+ g = u
+ grouped_flat.append(g)
+ return nest.pack_sequence_as(regrouped, grouped_flat)
+
+
class PerDeviceDataIterator(object):
"""An iterator (like `tf.data.Iterator`) into a `PerDeviceDataset`."""
@@ -683,7 +1076,7 @@
def get_next(self, name=None):
"""Scatter the input across devices."""
if self._prefetch_on_device:
- data_list = self._iterator.get_next()
+ data_list = self._iterator.get_next(name=name)
index = dict(zip(self._devices, data_list))
else:
batch = self._iterator.get_next(name=name)
@@ -703,26 +1096,21 @@
class PerDeviceDataset(object):
"""Like `tf.data.Dataset` split devices, producing `PerDevice` data."""
- def __init__(
- self,
- dataset,
- devices,
- prefetch_on_device=None,
- source_device="/cpu:0",
- ):
+ def __init__(self, dataset, devices, prefetch_on_device=None):
self._devices = devices
- self._source_device = source_device if source_device is not None else "/cpu:0"
# Default to using prefetching in graph mode, unless specified.
- # TODO(rohanj): Enable prefetching in eager mode.
+ # TODO(priyag): Enable prefetching in eager mode.
self._prefetch_on_device = prefetch_on_device
if self._prefetch_on_device is None:
self._prefetch_on_device = not context.executing_eagerly()
assert not (self._prefetch_on_device and context.executing_eagerly()), (
"Prefetching is only supported in graph mode currently")
- self._dataset = dataset
- if not self._prefetch_on_device:
+ if self._prefetch_on_device:
+ self._dataset = dataset.apply(
+ prefetching_ops_v2.prefetch_to_devices(self._devices))
+ else:
# TODO(priyag): If dropping remainder is not appropriate, find another
# approach to distributing the dataset when not possible to divide evenly.
# Possibly not an issue when we start using PartitionedDataset.
@@ -730,33 +1118,15 @@
def make_one_shot_iterator(self):
"""Get a one time use iterator for the distributed PerDeviceDataset."""
- # Graph mode prefetching with one shot iterator is disabled.
- if not context.executing_eagerly():
- raise ValueError("Cannot create a one shot iterator. Please use "
- "`make_initializable_iterator()` instead.")
- # Eager mode prefetching would error out in constructor. Only remaining
- # cases are non-prefetching eager / graph mode. We delegate to
- # PerDeviceDataIterator to handle them.
dataset_iterator = self._dataset.make_one_shot_iterator()
- return PerDeviceDataIterator(
- dataset_iterator, self._devices, prefetch_on_device=False)
+ return PerDeviceDataIterator(dataset_iterator, self._devices,
+ self._prefetch_on_device)
def make_initializable_iterator(self):
"""Get an initializable iterator for the distributed PerDeviceDataset."""
- # Eager mode generates already initialized iterators. Hence we cannot create
- # an initializable iterator.
- if context.executing_eagerly():
- raise ValueError("Cannot create initializable iterator in Eager mode. "
- "Please use `make_one_shot_iterator` instead.")
- if self._prefetch_on_device:
- dataset_iterator = multi_device_iterator_ops.MultiDeviceIterator(
- self._dataset, self._devices, source_device=self._source_device)
- else:
- dataset_iterator = self._dataset.make_initializable_iterator()
- return PerDeviceDataIterator(
- dataset_iterator,
- self._devices,
- prefetch_on_device=self._prefetch_on_device)
+ dataset_iterator = self._dataset.make_initializable_iterator()
+ return PerDeviceDataIterator(dataset_iterator, self._devices,
+ self._prefetch_on_device)
class MultiWorkerDataIterator(object):
@@ -816,7 +1186,8 @@
eager mode.
"""
- def __init__(self, dataset_fn, worker_device_map, prefetch_on_device=None):
+ def __init__(self, dataset_fn, worker_device_map, prefetch_on_device=None,
+ auto_shard=False):
"""Initialize the MultiWorkerDataset object.
Args:
@@ -824,6 +1195,7 @@
worker_device_map: a dict mapping from each worker to a list of devices
that belong to this worker.
prefetch_on_device: whether to prefetch to devices.
+ auto_shard: whether to auto-shard the dataset.
"""
self._worker_device_map = worker_device_map
self._datasets = {}
@@ -833,13 +1205,11 @@
six.iteritems(worker_device_map)):
with ops.device(worker):
worker_input = dataset_fn()
- worker_input = input_ops.auto_shard_dataset(
- worker_input, len(worker_device_map), i)
+ if auto_shard:
+ worker_input = input_ops.auto_shard_dataset(
+ worker_input, len(worker_device_map), i)
self._datasets[worker] = PerDeviceDataset(
- worker_input,
- worker_devices,
- source_device=worker,
- prefetch_on_device=prefetch_on_device)
+ worker_input, worker_devices, prefetch_on_device=prefetch_on_device)
def make_one_shot_iterator(self):
iterators = {}
diff --git a/tensorflow/contrib/distribute/python/values_test.py b/tensorflow/contrib/distribute/python/values_test.py
index 002d61f..121d2fb 100644
--- a/tensorflow/contrib/distribute/python/values_test.py
+++ b/tensorflow/contrib/distribute/python/values_test.py
@@ -349,11 +349,7 @@
def _test_iterator_no_prefetch(self, devices, dataset, expected_values):
per_device_dataset = values.PerDeviceDataset(
dataset, devices, prefetch_on_device=False)
- if context.executing_eagerly():
- iterator = per_device_dataset.make_one_shot_iterator()
- else:
- iterator = per_device_dataset.make_initializable_iterator()
- self.evaluate([iterator.initializer])
+ iterator = per_device_dataset.make_one_shot_iterator()
for expected_value in expected_values:
next_element = iterator.get_next()
@@ -370,14 +366,21 @@
if not context.executing_eagerly():
per_device_dataset = values.PerDeviceDataset(
dataset, devices, prefetch_on_device=True)
- iterator = per_device_dataset.make_initializable_iterator()
- self.evaluate([iterator.initializer])
+ iterator = per_device_dataset.make_one_shot_iterator()
+ # With prefetching, we cannot guarantee which input ends up on which
+ # device, so we verify that the complete set seen on all devices is
+ # correct, and equal numbers are distributed to each device.
+ combined_actual = []
+ combined_expected = []
for expected_value in expected_values:
next_element = iterator.get_next()
- computed_value = self.evaluate(
- [values.select_device(d, next_element) for d in devices])
- self.assertEqual(expected_value, computed_value)
+ combined_actual.extend(
+ self.evaluate(
+ [values.select_device(d, next_element) for d in devices]))
+ combined_expected.extend(expected_value)
+
+ self.assertEqual(set(combined_expected), set(combined_actual))
with self.assertRaises(errors.OutOfRangeError):
next_element = iterator.get_next()
@@ -638,7 +641,7 @@
if context.num_gpus() < 1 and context.executing_eagerly():
self.skipTest("A GPU is not available for this test in eager mode.")
- with self.test_session() as sess:
+ with self.cached_session(config=self.config) as sess:
v, devices, mirrored = _make_mirrored()
# Overwrite the initial values.
@@ -741,7 +744,7 @@
if context.num_gpus() < 1 or context.executing_eagerly():
self.skipTest("A GPU is not available for this test or it's eager mode.")
- with self.test_session(
+ with self.session(
graph=ops.Graph()) as sess, mirrored_strategy.MirroredStrategy(
["/device:GPU:0"]).scope():
with ops.device("/device:GPU:0"):
@@ -824,7 +827,7 @@
if context.num_gpus() < 1 and context.executing_eagerly():
self.skipTest("A GPU is not available for this test in eager mode.")
- with self.test_session() as sess:
+ with self.cached_session(config=self.config) as sess:
v, tower_local = _make_tower_local(variable_scope.VariableAggregation.SUM)
# Overwrite the initial values.
@@ -847,7 +850,7 @@
if context.num_gpus() < 1 and context.executing_eagerly():
self.skipTest("A GPU is not available for this test in eager mode.")
- with self.test_session() as sess:
+ with self.cached_session(config=self.config) as sess:
v, tower_local = _make_tower_local(
variable_scope.VariableAggregation.MEAN)
diff --git a/tensorflow/contrib/distributions/python/kernel_tests/moving_stats_test.py b/tensorflow/contrib/distributions/python/kernel_tests/moving_stats_test.py
index 3c988da..be7c756 100644
--- a/tensorflow/contrib/distributions/python/kernel_tests/moving_stats_test.py
+++ b/tensorflow/contrib/distributions/python/kernel_tests/moving_stats_test.py
@@ -38,8 +38,8 @@
true_stddev = np.array([[1.1, 0.5]])
with self.cached_session() as sess:
# Start "x" out with this mean.
- mean_var = variables.Variable(array_ops.zeros_like(true_mean))
- variance_var = variables.Variable(array_ops.ones_like(true_stddev))
+ mean_var = variables.VariableV1(array_ops.zeros_like(true_mean))
+ variance_var = variables.VariableV1(array_ops.ones_like(true_stddev))
x = random_ops.random_normal(shape, dtype=np.float64, seed=0)
x = true_stddev * x + true_mean
ema, emv = moving_stats.assign_moving_mean_variance(
@@ -115,7 +115,7 @@
# Start "x" out with this mean.
x = random_ops.random_normal(shape, dtype=np.float64, seed=0)
x = true_stddev * x + true_mean
- log_mean_exp_var = variables.Variable(array_ops.zeros_like(true_mean))
+ log_mean_exp_var = variables.VariableV1(array_ops.zeros_like(true_mean))
variables.global_variables_initializer().run()
log_mean_exp = moving_stats.assign_log_moving_mean_exp(
log_mean_exp_var, x, decay=decay)
diff --git a/tensorflow/contrib/eager/python/datasets.py b/tensorflow/contrib/eager/python/datasets.py
index 135095a..3aed121 100644
--- a/tensorflow/contrib/eager/python/datasets.py
+++ b/tensorflow/contrib/eager/python/datasets.py
@@ -18,7 +18,7 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.contrib.data.python.ops import prefetching_ops
+from tensorflow.python.data.experimental.ops import prefetching_ops
from tensorflow.python.data.ops import iterator_ops
from tensorflow.python.eager import context
from tensorflow.python.framework import ops
@@ -54,7 +54,7 @@
"""
if isinstance(dataset, prefetching_ops._PrefetchToDeviceDataset): # pylint: disable=protected-access
raise TypeError(
- "`tf.contrib.data.prefetch_to_device()` is not compatible with "
+ "`tf.data.experimental.prefetch_to_device()` is not compatible with "
"`tf.contrib.eager.Iterator`. Use `for ... in dataset:` to iterate "
"over the dataset instead.")
diff --git a/tensorflow/contrib/eager/python/datasets_test.py b/tensorflow/contrib/eager/python/datasets_test.py
index a753d77..6a508fc 100644
--- a/tensorflow/contrib/eager/python/datasets_test.py
+++ b/tensorflow/contrib/eager/python/datasets_test.py
@@ -24,11 +24,11 @@
import numpy as np
from tensorflow.contrib import lookup
-from tensorflow.contrib.data.python.ops import prefetching_ops
-from tensorflow.contrib.data.python.ops import threadpool
-from tensorflow.contrib.data.python.ops import unique
from tensorflow.contrib.eager.python import datasets
from tensorflow.python.data import Dataset
+from tensorflow.python.data.experimental.ops import prefetching_ops
+from tensorflow.python.data.experimental.ops import threadpool
+from tensorflow.python.data.experimental.ops import unique
from tensorflow.python.eager import test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
diff --git a/tensorflow/contrib/eager/python/examples/BUILD b/tensorflow/contrib/eager/python/examples/BUILD
index 6f02c90..97c299a 100644
--- a/tensorflow/contrib/eager/python/examples/BUILD
+++ b/tensorflow/contrib/eager/python/examples/BUILD
@@ -6,6 +6,7 @@
py_library(
name = "examples_pip",
deps = [
+ "//tensorflow/contrib/eager/python/examples/densenet",
"//tensorflow/contrib/eager/python/examples/gan:mnist",
"//tensorflow/contrib/eager/python/examples/l2hmc",
"//tensorflow/contrib/eager/python/examples/l2hmc:neural_nets",
diff --git a/tensorflow/contrib/eager/python/examples/gan/BUILD b/tensorflow/contrib/eager/python/examples/gan/BUILD
index c61ec2d..d64c8eb 100644
--- a/tensorflow/contrib/eager/python/examples/gan/BUILD
+++ b/tensorflow/contrib/eager/python/examples/gan/BUILD
@@ -3,6 +3,7 @@
package(default_visibility = ["//tensorflow:internal"])
load("//tensorflow:tensorflow.bzl", "cuda_py_test")
+load("//tensorflow:tensorflow.bzl", "py_binary")
py_binary(
name = "mnist",
diff --git a/tensorflow/contrib/eager/python/examples/l2hmc/l2hmc_test.py b/tensorflow/contrib/eager/python/examples/l2hmc/l2hmc_test.py
index c38a159..1c925e4 100644
--- a/tensorflow/contrib/eager/python/examples/l2hmc/l2hmc_test.py
+++ b/tensorflow/contrib/eager/python/examples/l2hmc/l2hmc_test.py
@@ -45,6 +45,17 @@
return loss, samples
+# To be defunnable, the function cannot return an Operation, so the above
+# function is used for defun or eager, and this function is used in graph to be
+# able to run the gradient updates.
+def graph_step(dynamics, optimizer, samples):
+ loss, grads, samples, _ = l2hmc.loss_and_grads(
+ dynamics, samples, loss_fn=l2hmc.compute_loss)
+ train_op = optimizer.apply_gradients(zip(grads, dynamics.variables))
+
+ return train_op, loss, samples
+
+
def warmup(dynamics,
optimizer,
n_iters=1,
@@ -134,51 +145,48 @@
"""Benchmark Graph performance."""
hparams = get_default_hparams()
- tf.reset_default_graph()
- with tf.Graph().as_default():
- energy_fn, _, _ = l2hmc.get_scg_energy_fn()
- dynamics = l2hmc.Dynamics(
- x_dim=hparams.x_dim,
- minus_loglikelihood_fn=energy_fn,
- n_steps=hparams.n_steps,
- eps=hparams.eps)
- x = tf.placeholder(tf.float32, shape=[None, hparams.x_dim])
- loss, x_out, _ = l2hmc.compute_loss(dynamics, x)
+ tf.enable_resource_variables()
+ for sample_size in [10, 25, 50, 100, 200]:
+ hparams.n_samples = sample_size
+ tf.reset_default_graph()
+ with tf.Graph().as_default():
+ energy_fn, _, _ = l2hmc.get_scg_energy_fn()
+ x = tf.random_normal([hparams.n_samples, hparams.x_dim],
+ dtype=tf.float32)
+ dynamics = l2hmc.Dynamics(
+ x_dim=hparams.x_dim,
+ minus_loglikelihood_fn=energy_fn,
+ n_steps=hparams.n_steps,
+ eps=hparams.eps)
+ loss, _, _ = l2hmc.compute_loss(dynamics, x)
- global_step = tf.Variable(0., name="global_step", trainable=False)
- learning_rate = tf.train.exponential_decay(
- hparams.learning_rate, global_step, 1000, 0.96, staircase=True)
- optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
- train_op = optimizer.minimize(loss, global_step=global_step)
+ optimizer = tf.train.AdamOptimizer(learning_rate=hparams.learning_rate)
+ train_op, loss, _ = graph_step(dynamics, optimizer, x)
- # Single thread; fairer comparison against eager
- session_conf = tf.ConfigProto(
- intra_op_parallelism_threads=1, inter_op_parallelism_threads=1)
+ # Single thread; fairer comparison against eager
+ session_conf = tf.ConfigProto(inter_op_parallelism_threads=1)
- with tf.Session(config=session_conf) as sess:
- sess.run(tf.global_variables_initializer())
+ with tf.Session(config=session_conf) as sess:
+ sess.run(tf.global_variables_initializer())
- # Warmup to reduce initialization effect when timing
- samples = npr.normal(size=[hparams.n_samples, hparams.x_dim])
- for _ in range(hparams.n_warmup_iters):
- _, _, _, _ = sess.run(
- [x_out, loss, train_op, learning_rate], feed_dict={x: samples})
+ # Warmup to reduce initialization effect when timing
+ for _ in range(hparams.n_warmup_iters):
+ _, _ = sess.run([train_op, loss])
- # Training
- start_time = time.time()
- for i in range(hparams.n_iters):
- samples, loss_np, _, _ = sess.run(
- [x_out, loss, train_op, learning_rate], feed_dict={x: samples})
- print("Iteration %d: loss %.4f" % (i, loss_np))
- wall_time = time.time() - start_time
- examples_per_sec = hparams.n_samples / wall_time
+ # Training
+ start_time = time.time()
+ for i in range(hparams.n_iters):
+ _, loss_np = sess.run([train_op, loss])
+ print("Iteration %d: loss %.4f" % (i, loss_np))
+ wall_time = (time.time() - start_time) / hparams.n_iters
+ examples_per_sec = hparams.n_samples / wall_time
- self.report_benchmark(
- name="graph_train_%s" % ("gpu"
- if tf.test.is_gpu_available() else "cpu"),
- iters=hparams.n_iters,
- extras={"examples_per_sec": examples_per_sec},
- wall_time=wall_time)
+ self.report_benchmark(
+ name="graph_train_%s_%d" %
+ ("gpu" if tf.test.is_gpu_available() else "cpu", sample_size),
+ iters=hparams.n_iters,
+ extras={"examples_per_sec": examples_per_sec},
+ wall_time=wall_time)
def benchmark_eager(self):
self._benchmark_eager()
@@ -190,32 +198,44 @@
"""Benchmark Eager performance."""
hparams = get_default_hparams()
- energy_fn, _, _ = l2hmc.get_scg_energy_fn()
- dynamics = l2hmc.Dynamics(
- x_dim=hparams.x_dim,
- minus_loglikelihood_fn=energy_fn,
- n_steps=hparams.n_steps,
- eps=hparams.eps)
- optimizer = tf.train.AdamOptimizer(learning_rate=hparams.learning_rate)
- step_fn = tfe.defun(step) if defun else step
+ for sample_size in [10, 25, 50, 100, 200]:
+ hparams.n_samples = sample_size
+ energy_fn, _, _ = l2hmc.get_scg_energy_fn()
+ dynamics = l2hmc.Dynamics(
+ x_dim=hparams.x_dim,
+ minus_loglikelihood_fn=energy_fn,
+ n_steps=hparams.n_steps,
+ eps=hparams.eps)
+ optimizer = tf.train.AdamOptimizer(learning_rate=hparams.learning_rate)
+ step_fn = tfe.defun(step) if defun else step
- # Warmup to reduce initialization effect when timing
- warmup(dynamics, optimizer, n_iters=hparams.n_warmup_iters, step_fn=step_fn)
+ # Warmup to reduce initialization effect when timing
+ warmup(
+ dynamics,
+ optimizer,
+ n_iters=hparams.n_warmup_iters,
+ n_samples=hparams.n_samples,
+ step_fn=step_fn)
- # Training
- samples = tf.random_normal(
- shape=[hparams.n_samples, hparams.x_dim], dtype=tf.float32)
- start_time = time.time()
- fit(dynamics, samples, optimizer, step_fn=step_fn, n_iters=hparams.n_iters)
- wall_time = time.time() - start_time
- examples_per_sec = hparams.n_samples / wall_time
+ # Training
+ samples = tf.random_normal(
+ shape=[hparams.n_samples, hparams.x_dim], dtype=tf.float32)
+ start_time = time.time()
+ fit(dynamics,
+ samples,
+ optimizer,
+ step_fn=step_fn,
+ n_iters=hparams.n_iters)
+ wall_time = (time.time() - start_time) / hparams.n_iters
+ examples_per_sec = hparams.n_samples / wall_time
- self.report_benchmark(
- name="eager_train_%s%s" % ("gpu" if tf.test.is_gpu_available() else
- "cpu", "_defun" if defun else ""),
- iters=hparams.n_iters,
- extras={"examples_per_sec": examples_per_sec},
- wall_time=wall_time)
+ self.report_benchmark(
+ name="eager_train_%s%s_%d" %
+ ("gpu" if tf.test.is_gpu_available() else "cpu",
+ "_defun" if defun else "", sample_size),
+ iters=hparams.n_iters,
+ extras={"examples_per_sec": examples_per_sec},
+ wall_time=wall_time)
del dynamics
diff --git a/tensorflow/contrib/eager/python/examples/linear_regression/BUILD b/tensorflow/contrib/eager/python/examples/linear_regression/BUILD
index 2f6cfdf..74ce9e8 100644
--- a/tensorflow/contrib/eager/python/examples/linear_regression/BUILD
+++ b/tensorflow/contrib/eager/python/examples/linear_regression/BUILD
@@ -3,6 +3,7 @@
package(default_visibility = ["//tensorflow:internal"])
load("//tensorflow:tensorflow.bzl", "cuda_py_test")
+load("//tensorflow:tensorflow.bzl", "py_binary")
py_binary(
name = "linear_regression",
diff --git a/tensorflow/contrib/eager/python/examples/revnet/imagenet_input.py b/tensorflow/contrib/eager/python/examples/revnet/imagenet_input.py
index 34a9984..d85188d 100644
--- a/tensorflow/contrib/eager/python/examples/revnet/imagenet_input.py
+++ b/tensorflow/contrib/eager/python/examples/revnet/imagenet_input.py
@@ -169,11 +169,11 @@
# Read the data from disk in parallel
dataset = dataset.apply(
- tf.contrib.data.parallel_interleave(
+ tf.data.experimental.parallel_interleave(
fetch_dataset, cycle_length=self.num_parallel_calls, sloppy=True))
if self.cache:
dataset = dataset.cache().apply(
- tf.contrib.data.shuffle_and_repeat(1024 * 16))
+ tf.data.experimental.shuffle_and_repeat(1024 * 16))
else:
dataset = dataset.shuffle(1024)
@@ -188,9 +188,11 @@
# batch size. As long as this validation is done with consistent batch size,
# exactly the same images will be used.
dataset = dataset.apply(
- tf.contrib.data.map_and_batch(
- self.dataset_parser, batch_size=batch_size,
- num_parallel_batches=self.num_cores, drop_remainder=True))
+ tf.data.experimental.map_and_batch(
+ self.dataset_parser,
+ batch_size=batch_size,
+ num_parallel_batches=self.num_cores,
+ drop_remainder=True))
# Transpose for performance on TPU
if self.transpose_input:
diff --git a/tensorflow/contrib/eager/python/examples/rnn_colorbot/BUILD b/tensorflow/contrib/eager/python/examples/rnn_colorbot/BUILD
index f83eb5c..d500b63 100644
--- a/tensorflow/contrib/eager/python/examples/rnn_colorbot/BUILD
+++ b/tensorflow/contrib/eager/python/examples/rnn_colorbot/BUILD
@@ -3,6 +3,7 @@
package(default_visibility = ["//tensorflow:internal"])
load("//tensorflow:tensorflow.bzl", "cuda_py_test")
+load("//tensorflow:tensorflow.bzl", "py_binary")
py_binary(
name = "rnn_colorbot",
diff --git a/tensorflow/contrib/eager/python/examples/rnn_ptb/BUILD b/tensorflow/contrib/eager/python/examples/rnn_ptb/BUILD
index 4b4792c..2cc2fcb 100644
--- a/tensorflow/contrib/eager/python/examples/rnn_ptb/BUILD
+++ b/tensorflow/contrib/eager/python/examples/rnn_ptb/BUILD
@@ -3,6 +3,7 @@
package(default_visibility = ["//tensorflow:internal"])
load("//tensorflow:tensorflow.bzl", "cuda_py_test")
+load("//tensorflow:tensorflow.bzl", "py_binary")
py_binary(
name = "rnn_ptb",
diff --git a/tensorflow/contrib/eager/python/remote_test.py b/tensorflow/contrib/eager/python/remote_test.py
index ba6fe97..7aa4b59 100644
--- a/tensorflow/contrib/eager/python/remote_test.py
+++ b/tensorflow/contrib/eager/python/remote_test.py
@@ -47,8 +47,9 @@
@functools.wraps(f)
def decorator(self, *args, **kwargs):
- with context.execution_mode(context.ASYNC):
- f(self, *args, **kwargs)
+ # TODO(b/117110239): Re-enable.
+ # with context.execution_mode(context.ASYNC):
+ # f(self, *args, **kwargs)
with context.execution_mode(context.SYNC):
f(self, *args, **kwargs)
diff --git a/tensorflow/contrib/estimator/__init__.py b/tensorflow/contrib/estimator/__init__.py
index 78914ec..419609b 100644
--- a/tensorflow/contrib/estimator/__init__.py
+++ b/tensorflow/contrib/estimator/__init__.py
@@ -76,7 +76,7 @@
'stop_if_no_decrease_hook',
'build_raw_supervised_input_receiver_fn',
'build_supervised_input_receiver_fn_from_input_fn',
- 'SavedModelEstimator'
+ 'SavedModelEstimator',
'DNNClassifierWithLayerAnnotations',
'DNNRegressorWithLayerAnnotations',
]
diff --git a/tensorflow/contrib/estimator/python/estimator/dnn_with_layer_annotations.py b/tensorflow/contrib/estimator/python/estimator/dnn_with_layer_annotations.py
index 3fd9f12..5faf0aa 100644
--- a/tensorflow/contrib/estimator/python/estimator/dnn_with_layer_annotations.py
+++ b/tensorflow/contrib/estimator/python/estimator/dnn_with_layer_annotations.py
@@ -75,7 +75,9 @@
weight_collections=None,
trainable=True,
cols_to_vars=None,
- cols_to_output_tensors=None):
+ scope=None,
+ cols_to_output_tensors=None,
+ from_template=False):
"""Returns a dense `Tensor` as input layer based on given `feature_columns`.
Generally a single example in training data is described with
@@ -111,9 +113,12 @@
'some_variable:0' shape=(5, 10), <tf.Variable 'some_variable:1'
shape=(5, 10)]} If a column creates no variables, its value will be an
empty list.
+ scope: A name or variable scope to use
cols_to_output_tensors: If not `None`, must be a dictionary that will be
filled with a mapping from '_FeatureColumn' to the associated output
`Tensor`s.
+ from_template: True if the method is being instantiated from a
+ `make_template`.
Returns:
A `Tensor` which represents input layer of a model. Its shape
@@ -131,7 +136,9 @@
weight_collections=weight_collections,
trainable=trainable,
cols_to_vars=cols_to_vars,
- cols_to_output_tensors=local_cols_to_output_tensors)
+ scope=scope,
+ cols_to_output_tensors=local_cols_to_output_tensors,
+ from_template=from_template)
if cols_to_output_tensors is not None:
cols_to_output_tensors = local_cols_to_output_tensors
@@ -296,9 +303,9 @@
def _model_fn(features, labels, mode, config):
with _monkey_patch(
- feature_column_lib, 'input_layer',
+ feature_column_lib, '_internal_input_layer',
make_input_layer_with_layer_annotations(
- feature_column_lib.input_layer)):
+ feature_column_lib._internal_input_layer)): # pylint: disable=protected-access
return original.model_fn(features, labels, mode, config)
return estimator.Estimator(
@@ -417,9 +424,9 @@
def _model_fn(features, labels, mode, config):
with _monkey_patch(
- feature_column_lib, 'input_layer',
+ feature_column_lib, '_internal_input_layer',
make_input_layer_with_layer_annotations(
- feature_column_lib.input_layer)):
+ feature_column_lib._internal_input_layer)): # pylint: disable=protected-access
return original.model_fn(features, labels, mode, config)
return estimator.Estimator(
diff --git a/tensorflow/contrib/estimator/python/estimator/hooks_test.py b/tensorflow/contrib/estimator/python/estimator/hooks_test.py
index c6c6cad..62ffad5 100644
--- a/tensorflow/contrib/estimator/python/estimator/hooks_test.py
+++ b/tensorflow/contrib/estimator/python/estimator/hooks_test.py
@@ -294,7 +294,7 @@
def model_fn(features, labels, mode):
_, _ = features, labels
- w = variables.Variable(
+ w = variables.VariableV1(
initial_value=[0.],
trainable=False,
collections=[ops.GraphKeys.SAVEABLE_OBJECTS])
diff --git a/tensorflow/contrib/estimator/python/estimator/rnn_test.py b/tensorflow/contrib/estimator/python/estimator/rnn_test.py
index 1aebed3..89506ee 100644
--- a/tensorflow/contrib/estimator/python/estimator/rnn_test.py
+++ b/tensorflow/contrib/estimator/python/estimator/rnn_test.py
@@ -25,12 +25,12 @@
import numpy as np
import six
-from tensorflow.contrib.data.python.ops import readers
from tensorflow.contrib.estimator.python.estimator import head as head_lib
from tensorflow.contrib.estimator.python.estimator import rnn
from tensorflow.contrib.feature_column.python.feature_column import sequence_feature_column as seq_fc
from tensorflow.core.example import example_pb2
from tensorflow.core.example import feature_pb2
+from tensorflow.python.data.experimental.ops import readers
from tensorflow.python.estimator import model_fn
from tensorflow.python.estimator.canned import metric_keys
from tensorflow.python.estimator.canned import parsing_utils
diff --git a/tensorflow/contrib/factorization/python/ops/gmm_ops.py b/tensorflow/contrib/factorization/python/ops/gmm_ops.py
index e076631..d365ad1 100644
--- a/tensorflow/contrib/factorization/python/ops/gmm_ops.py
+++ b/tensorflow/contrib/factorization/python/ops/gmm_ops.py
@@ -154,10 +154,10 @@
def _create_variables(self):
"""Initializes GMM algorithm."""
init_value = array_ops.constant([], dtype=dtypes.float32)
- self._means = variables.Variable(init_value,
- name=self.CLUSTERS_VARIABLE,
- validate_shape=False)
- self._covs = variables.Variable(
+ self._means = variables.VariableV1(init_value,
+ name=self.CLUSTERS_VARIABLE,
+ validate_shape=False)
+ self._covs = variables.VariableV1(
init_value, name=self.CLUSTERS_COVS_VARIABLE, validate_shape=False)
# Mixture weights, representing the probability that a randomly
# selected unobservable data (in EM terms) was generated by component k.
@@ -165,9 +165,9 @@
array_ops.tile([1.0 / self._num_classes], [self._num_classes]),
name=self.CLUSTERS_WEIGHT,
validate_shape=False)
- self._cluster_centers_initialized = variables.Variable(False,
- dtype=dtypes.bool,
- name='initialized')
+ self._cluster_centers_initialized = variables.VariableV1(False,
+ dtype=dtypes.bool,
+ name='initialized')
def _initialize_variables(self, data, initial_means=None):
"""Initializes variables.
diff --git a/tensorflow/contrib/factorization/python/ops/wals_test.py b/tensorflow/contrib/factorization/python/ops/wals_test.py
index 9bdbd05..75d577f 100644
--- a/tensorflow/contrib/factorization/python/ops/wals_test.py
+++ b/tensorflow/contrib/factorization/python/ops/wals_test.py
@@ -420,13 +420,13 @@
class SweepHookTest(test.TestCase):
def test_sweeps(self):
- is_row_sweep_var = variables.Variable(True)
- is_sweep_done_var = variables.Variable(False)
- init_done = variables.Variable(False)
- row_prep_done = variables.Variable(False)
- col_prep_done = variables.Variable(False)
- row_train_done = variables.Variable(False)
- col_train_done = variables.Variable(False)
+ is_row_sweep_var = variables.VariableV1(True)
+ is_sweep_done_var = variables.VariableV1(False)
+ init_done = variables.VariableV1(False)
+ row_prep_done = variables.VariableV1(False)
+ col_prep_done = variables.VariableV1(False)
+ row_train_done = variables.VariableV1(False)
+ col_train_done = variables.VariableV1(False)
init_op = state_ops.assign(init_done, True)
row_prep_op = state_ops.assign(row_prep_done, True)
@@ -486,7 +486,7 @@
def test_stop(self):
hook = wals_lib._StopAtSweepHook(last_sweep=10)
- completed_sweeps = variables.Variable(
+ completed_sweeps = variables.VariableV1(
8, name=wals_lib.WALSMatrixFactorization.COMPLETED_SWEEPS)
train_op = state_ops.assign_add(completed_sweeps, 1)
hook.begin()
diff --git a/tensorflow/contrib/framework/python/ops/variables_test.py b/tensorflow/contrib/framework/python/ops/variables_test.py
index f9b0efd..c223df5 100644
--- a/tensorflow/contrib/framework/python/ops/variables_test.py
+++ b/tensorflow/contrib/framework/python/ops/variables_test.py
@@ -192,7 +192,7 @@
def test_invalid_dtype(self):
with ops.Graph().as_default() as g:
self.assertEquals(None, variables_lib2.get_global_step())
- variables_lib.Variable(
+ variables_lib.VariableV1(
0.0,
trainable=False,
dtype=dtypes.float32,
@@ -205,7 +205,7 @@
def test_invalid_shape(self):
with ops.Graph().as_default() as g:
self.assertEquals(None, variables_lib2.get_global_step())
- variables_lib.Variable(
+ variables_lib.VariableV1(
[0],
trainable=False,
dtype=dtypes.int32,
@@ -229,7 +229,7 @@
def test_get_global_step(self):
with ops.Graph().as_default() as g:
self.assertEquals(None, variables_lib2.get_global_step())
- variables_lib.Variable(
+ variables_lib.VariableV1(
0,
trainable=False,
dtype=dtypes.int32,
@@ -607,10 +607,10 @@
with self.cached_session():
with variable_scope.variable_scope('A'):
variables_lib2.local_variable([5])
- a = variables_lib.Variable([5])
+ a = variables_lib.VariableV1([5])
with variable_scope.variable_scope('B'):
variables_lib2.local_variable([5])
- b = variables_lib.Variable([5])
+ b = variables_lib.VariableV1([5])
self.assertEquals([a], variables_lib2.get_trainable_variables('A'))
self.assertEquals([b], variables_lib2.get_trainable_variables('B'))
@@ -953,7 +953,7 @@
# Create a set of variables to save in the checkpoint.
for var_name in var_names_to_values:
var_value = var_names_to_values[var_name]
- var_list.append(variables_lib.Variable(var_value, name=var_name))
+ var_list.append(variables_lib.VariableV1(var_value, name=var_name))
saver = saver_lib.Saver(var_list)
init_op = variables_lib.variables_initializer(var_list)
sess.run(init_op)
@@ -1106,7 +1106,7 @@
# Create a set of variables to save in the checkpoint.
for var_name in var_names_to_values:
var_value = var_names_to_values[var_name]
- var_list.append(variables_lib.Variable(var_value, name=var_name))
+ var_list.append(variables_lib.VariableV1(var_value, name=var_name))
saver = saver_lib.Saver(var_list)
init_op = variables_lib.variables_initializer(var_list)
sess.run(init_op)
@@ -1297,7 +1297,7 @@
class ZeroInitializerOpTest(test.TestCase):
def _testZeroInitializer(self, shape, initializer, use_init):
- var = variables_lib.Variable(initializer)
+ var = variables_lib.VariableV1(initializer)
var_zero = variables_lib2.zero_initializer(var)
with self.cached_session() as sess:
with self.assertRaisesOpError('Attempting to use uninitialized value'):
@@ -1350,12 +1350,12 @@
g = ops.Graph()
with g.as_default():
var_list = []
- var_list.append(variables_lib.Variable(0, name='conv1/weights'))
- var_list.append(variables_lib.Variable(0, name='conv1/biases'))
- var_list.append(variables_lib.Variable(0, name='conv2/weights'))
- var_list.append(variables_lib.Variable(0, name='conv2/biases'))
- var_list.append(variables_lib.Variable(0, name='clfs/weights'))
- var_list.append(variables_lib.Variable(0, name='clfs/biases'))
+ var_list.append(variables_lib.VariableV1(0, name='conv1/weights'))
+ var_list.append(variables_lib.VariableV1(0, name='conv1/biases'))
+ var_list.append(variables_lib.VariableV1(0, name='conv2/weights'))
+ var_list.append(variables_lib.VariableV1(0, name='conv2/biases'))
+ var_list.append(variables_lib.VariableV1(0, name='clfs/weights'))
+ var_list.append(variables_lib.VariableV1(0, name='clfs/biases'))
self._var_list = var_list
def _test_filter_variables(self,
diff --git a/tensorflow/contrib/fused_conv/BUILD b/tensorflow/contrib/fused_conv/BUILD
index 9725233..490da9b 100644
--- a/tensorflow/contrib/fused_conv/BUILD
+++ b/tensorflow/contrib/fused_conv/BUILD
@@ -17,11 +17,14 @@
exports_files(["LICENSE"])
-load("//tensorflow:tensorflow.bzl", "tf_kernel_library")
+load(
+ "//tensorflow:tensorflow.bzl",
+ "tf_kernel_library",
+ "tf_custom_op_library",
+ "tf_gen_op_libs",
+ "tf_gen_op_wrapper_py",
+)
load("//tensorflow:tensorflow.bzl", "cuda_py_test")
-load("//tensorflow:tensorflow.bzl", "tf_custom_op_library")
-load("//tensorflow:tensorflow.bzl", "tf_gen_op_libs")
-load("//tensorflow:tensorflow.bzl", "tf_gen_op_wrapper_py")
load("//tensorflow:tensorflow.bzl", "tf_custom_op_py_library")
tf_custom_op_py_library(
@@ -109,12 +112,13 @@
deps = [":fused_conv2d_bias_activation_op_op_lib"],
)
-cuda_py_test(
- name = "fused_conv2d_bias_activation_op_test",
- srcs = ["python/ops/fused_conv2d_bias_activation_op_test.py"],
- additional_deps = [
+py_library(
+ name = "fused_conv2d_bias_activation_op_test_base",
+ testonly = 1,
+ srcs = ["python/ops/fused_conv2d_bias_activation_op_test_base.py"],
+ visibility = ["//tensorflow/compiler/tf2xla:internal"],
+ deps = [
":fused_conv_py",
- "//third_party/py/numpy",
"//tensorflow/python:array_ops",
"//tensorflow/python:client",
"//tensorflow/python:client_testlib",
@@ -127,8 +131,21 @@
"//tensorflow/python:random_ops",
"//tensorflow/python:training",
"//tensorflow/python:variables",
+ "//third_party/py/numpy",
+ "@absl_py//absl/testing:parameterized",
+ ],
+)
+
+cuda_py_test(
+ name = "fused_conv2d_bias_activation_op_test",
+ size = "large",
+ srcs = ["python/ops/fused_conv2d_bias_activation_op_test.py"],
+ additional_deps = [
+ ":fused_conv2d_bias_activation_op_test_base",
+ "//tensorflow/python:client_testlib",
],
tags = [
+ "no_pip",
"requires-gpu-sm70",
],
)
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 e9e6464..93b1aaa 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
@@ -111,8 +111,8 @@
context,
(GetTensorDim(strides, data_format_, 'N') == 1 &&
GetTensorDim(strides, data_format_, 'C') == 1),
- errors::InvalidArgument("Convolutional strides are not supported in "
- "the batch or depth dimensions."));
+ errors::Unimplemented("Convolutional strides are not supported in "
+ "the batch and depth dimensions."));
// Assuming qint8 <--> NCHW_VECT_C, OIHW_VECT_I (int8x4) here.
constexpr bool is_int8x4 = std::is_same<T, qint8>::value;
diff --git a/tensorflow/contrib/fused_conv/python/ops/fused_conv2d_bias_activation_op_test.py b/tensorflow/contrib/fused_conv/python/ops/fused_conv2d_bias_activation_op_test.py
index 4894298..e5c8a34 100644
--- a/tensorflow/contrib/fused_conv/python/ops/fused_conv2d_bias_activation_op_test.py
+++ b/tensorflow/contrib/fused_conv/python/ops/fused_conv2d_bias_activation_op_test.py
@@ -12,896 +12,27 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Functional tests for fused conv2d bias and activation operation."""
+
+"""Tests for fused convolutions."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
-import numpy as np
-
-from tensorflow.contrib.fused_conv.python.ops import fused_conv2d_bias_activation_op
-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 gen_array_ops
-from tensorflow.python.ops import nn_ops
-from tensorflow.python.ops import random_ops
+from tensorflow.contrib.fused_conv.python.ops import fused_conv2d_bias_activation_op_test_base as test_base
from tensorflow.python.platform import test
-from tensorflow.python.platform import tf_logging
-def GetShrunkInceptionShapes(shrink=10):
- """Iterator for smaller versions of convolution shapes in 2015 Inception.
+# Instantiate the two test suites from test_base, mixing in test.TestCase as
+# the test framework.
+class FusedConv2DBiasActivationTest(test_base.FusedConv2DBiasActivationTest,
+ test.TestCase):
+ pass
- Relative to inception, each depth value is `depth // shrink`.
- Args:
- shrink: Factor to shrink each depth value by relative to Inception.
+class FusedConvInt8Tests(test_base.FusedConvInt8Tests, test.TestCase):
+ pass
- Yields:
- Tuple (input_size, filter_size, out_size, stride, padding), the convolution
- parameters of Inception layers.
- """
- input_sizes = [[4, 5, 5, 1248], [4, 8, 8, 384], [4, 8, 8, 384], [
- 4, 8, 8, 2048
- ], [4, 8, 8, 448], [4, 8, 8, 2048], [4, 8, 8, 2048], [4, 8, 8, 2048], [
- 4, 8, 8, 1760
- ], [4, 8, 8, 1760], [4, 8, 8, 1760], [4, 8, 8, 1760], [4, 17, 17, 192], [
- 4, 17, 17, 192
- ], [4, 17, 17, 1248], [4, 17, 17, 128], [4, 17, 17, 1248], [4, 17, 17, 224], [
- 4, 17, 17, 192
- ], [4, 17, 17, 192], [4, 17, 17, 1216], [4, 17, 17, 1216], [4, 17, 17, 224], [
- 4, 17, 17, 192
- ], [4, 17, 17, 192], [4, 17, 17, 1152], [4, 17, 17, 1152], [4, 17, 17, 192], [
- 4, 17, 17, 160
- ], [4, 17, 17, 1152], [4, 17, 17, 1024], [4, 17, 17, 128], [4, 17, 17, 1024],
- [4, 17, 17, 128], [4, 17, 17, 1024], [4, 17, 17, 128], [
- 4, 17, 17, 768
- ], [4, 17, 17, 128], [4, 17, 17, 128], [4, 17, 17, 768],
- [4, 17, 17, 768], [4, 35, 35, 96], [4, 35, 35, 288], [
- 4, 35, 35, 64
- ], [4, 35, 35, 288], [4, 35, 35, 256], [4, 35, 35, 48], [
- 4, 35, 35, 256
- ], [4, 35, 35, 96], [4, 35, 35, 192], [4, 35, 35, 192], [
- 4, 35, 35, 192
- ], [4, 73, 73, 64], [4, 73, 73, 64], [4, 147, 147, 24]]
- filter_sizes = [[1, 1, 1248, 128], [1, 3, 384, 384], [3, 1, 384, 384], [
- 1, 1, 2048, 192
- ], [3, 3, 448, 384], [1, 1, 2048, 320], [1, 1, 2048, 448], [1, 1, 2048, 384],
- [1, 1, 1760, 384], [1, 1, 1760, 192], [1, 1, 1760, 448], [
- 1, 1, 1760, 320
- ], [3, 3, 192, 192], [3, 3, 192, 192], [1, 1, 1248, 192], [
- 3, 3, 128, 320
- ], [1, 1, 1248, 128], [1, 3, 224, 224], [3, 1, 192, 256], [
- 1, 3, 192, 256
- ], [1, 1, 1216, 192], [1, 1, 1216, 96], [3, 1, 224, 224], [
- 3, 3, 192, 224
- ], [1, 3, 192, 192], [1, 1, 1152, 192], [1, 1, 1152, 128], [
- 3, 1, 192, 192
- ], [3, 3, 160, 192], [1, 1, 1152, 160], [1, 1, 1024, 128], [
- 1, 3, 128, 192
- ], [1, 1, 1024, 160], [3, 1, 128, 192], [1, 1, 1024, 256], [
- 3, 1, 128, 128
- ], [1, 1, 768, 192], [1, 3, 128, 128], [3, 3, 128, 128], [
- 1, 1, 768, 128
- ], [1, 1, 768, 320], [3, 3, 96, 96], [3, 3, 288, 384], [
- 3, 3, 64, 96
- ], [1, 1, 288, 64], [1, 1, 256, 64], [5, 5, 48, 64],
- [1, 1, 256, 48], [3, 3, 96, 96], [1, 1, 192, 32], [
- 1, 1, 192, 64
- ], [1, 1, 192, 48], [3, 3, 64, 192], [1, 1, 64,
- 64], [1, 1, 24, 64]]
- out_sizes = [[4, 5, 5, 128], [4, 8, 8, 384], [4, 8, 8, 384], [4, 8, 8, 192], [
- 4, 8, 8, 384
- ], [4, 8, 8, 320], [4, 8, 8, 448], [4, 8, 8, 384], [4, 8, 8, 384], [
- 4, 8, 8, 192
- ], [4, 8, 8, 448], [4, 8, 8, 320], [4, 8, 8, 192], [4, 17, 17, 192], [
- 4, 17, 17, 192
- ], [4, 8, 8, 320], [4, 17, 17, 128], [4, 17, 17, 224], [4, 17, 17, 256], [
- 4, 17, 17, 256
- ], [4, 17, 17, 192], [4, 17, 17, 96], [4, 17, 17, 224], [4, 17, 17, 224], [
- 4, 17, 17, 192
- ], [4, 17, 17, 192], [4, 17, 17, 128], [4, 17, 17, 192], [4, 17, 17, 192], [
- 4, 17, 17, 160
- ], [4, 17, 17, 128], [4, 17, 17, 192], [4, 17, 17, 160], [4, 17, 17, 192], [
- 4, 17, 17, 256
- ], [4, 17, 17, 128], [4, 17, 17, 192], [4, 17, 17, 128], [4, 17, 17, 128], [
- 4, 17, 17, 128
- ], [4, 17, 17, 320], [4, 17, 17, 96], [4, 17, 17, 384], [4, 35, 35, 96], [
- 4, 35, 35, 64
- ], [4, 35, 35, 64], [4, 35, 35, 64], [4, 35, 35, 48], [4, 35, 35, 96],
- [4, 35, 35, 32], [4, 35, 35, 64], [4, 35, 35, 48],
- [4, 71, 71, 192], [4, 73, 73, 64], [4, 147, 147, 64]]
- strides = [
- 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1,
- 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 1, 1, 1, 1, 1, 1, 1,
- 1, 1, 1, 1, 1
- ]
- # Shrink sizes to make the test faster
- for i in input_sizes:
- i[3] //= shrink
- for f in filter_sizes:
- f[2] //= shrink
- f[3] //= shrink
- for o in out_sizes:
- o[3] //= shrink
- # pylint: disable=invalid-name
- VALID = "VALID"
- SAME = "SAME"
- # pylint: enable=invalid-name
- paddings = [
- SAME, SAME, SAME, SAME, SAME, SAME, SAME, SAME, SAME, SAME, SAME, SAME,
- VALID, SAME, SAME, VALID, SAME, SAME, SAME, SAME, SAME, SAME, SAME, SAME,
- SAME, SAME, SAME, SAME, SAME, SAME, SAME, SAME, SAME, SAME, SAME, SAME,
- SAME, SAME, SAME, SAME, SAME, VALID, VALID, SAME, SAME, SAME, SAME, SAME,
- SAME, SAME, SAME, SAME, VALID, VALID, VALID
- ]
- for i, f, o, s, p in zip(input_sizes, filter_sizes, out_sizes, strides,
- paddings):
- yield i, f, o, s, p
-
-def GetTestConfigs():
- """Get all the valid tests configs to run.
-
- Returns:
- all the valid test configs as tuples of data_format and use_gpu.
- """
- test_configs = [("NCHW", True), ("NHWC", True)]
- return test_configs
-
-
-class FusedConv2DBiasActivationTest(test.TestCase):
-
- def _DtypesToTest(self, use_gpu):
- return [dtypes.float32]
-
- def _FilterFormatsToTest(self, use_gpu):
- return ["HWIO", "OIHW"]
-
- def _SetupValuesForDevice(self, tensor_in_sizes, filter_in_sizes, bias,
- strides, padding, activation_mode, data_format,
- filter_format, dtype):
- """Verifies the output values of the convolution function.
-
- Args:
- tensor_in_sizes: Input tensor dimensions in
- [batch, input_rows, input_cols, input_depth].
- filter_in_sizes: Filter tensor dimensions in
- [kernel_rows, kernel_cols, input_depth, output_depth].
- bias: 1-D bias tensor of length output_depth.
- strides: Stride: [col_stride, row_stride]
- padding: Padding type.
- activation_mode: Activation mode.
- data_format: Format of the data tensors.
- filter_format: Filter format to use for the fused convolution.
- dtype: Data type for inputs and outputs.
- Returns:
- Symbolic tensor value and reference value that can be used to
- execute the computation and verify the results.
- """
- input_size = np.prod(tensor_in_sizes)
- filter_size = np.prod(filter_in_sizes)
- bias_size = filter_in_sizes[-1] # equals to output depth
- # Initializes the input tensor with array containing incrementing
- # numbers from 1.
- x1 = [f * 1.0 for f in range(1, input_size + 1)]
- x2 = [f * 1.0 for f in range(1, filter_size + 1)]
- # This is to guarantee that there is always negative values after
- # bias add so that we can test whether relu works correctly.
- x3 = bias
- with self.test_session(use_gpu=True):
- t1 = constant_op.constant(x1, shape=tensor_in_sizes, dtype=dtype)
- t2 = constant_op.constant(x2, shape=filter_in_sizes, dtype=dtype)
- fused_t2 = t2
- if filter_format == "OIHW":
- fused_t2 = HwioToOihw(t2)
- t3 = constant_op.constant(x3, shape=[bias_size], dtype=dtype)
- strides = [1] + strides + [1]
- if data_format == "NCHW":
- t1 = test_util.NHWCToNCHW(t1)
- strides = test_util.NHWCToNCHW(strides)
- output = fused_conv2d_bias_activation_op.fused_conv2d_bias_activation(
- t1,
- fused_t2,
- t3,
- strides=strides,
- padding=padding,
- data_format=data_format,
- filter_format=filter_format,
- activation_mode=activation_mode)
- ref_conv_output = nn_ops.conv2d(
- t1, t2, strides=strides, padding=padding, data_format=data_format)
- ref_bias_output = nn_ops.bias_add(
- ref_conv_output, t3, data_format=data_format)
- ref_output = nn_ops.relu(ref_bias_output)
- if data_format == "NCHW":
- output = test_util.NCHWToNHWC(output)
- ref_output = test_util.NCHWToNHWC(ref_output)
-
- return output, ref_output
-
- def _CompareFwdValues(self, tensor_in_sizes, filter_in_sizes, conv_strides,
- padding):
- """Verifies that CPU and GPU produce the same values.
-
- Args:
- tensor_in_sizes: Input tensor dimensions in
- [batch, input_rows, input_cols, input_depth].
- filter_in_sizes: Filter tensor dimensions in
- [kernel_rows, kernel_cols, input_depth, output_depth].
- conv_strides: [row_stride, col_stride] for the convolution;
- padding: Padding type.
- """
- x1 = np.random.rand(*tensor_in_sizes).astype(np.float32)
- x2 = np.random.rand(*filter_in_sizes).astype(np.float32)
- x3 = np.random.rand(*[filter_in_sizes[-1]]).astype(np.float32)
-
- def _SetupVal(data_format, use_gpu):
- with self.test_session(use_gpu=use_gpu):
- t1 = constant_op.constant(x1, shape=tensor_in_sizes)
- t2 = constant_op.constant(x2, shape=filter_in_sizes)
- t3 = constant_op.constant(x3, shape=[filter_in_sizes[-1]])
- strides = [1] + conv_strides + [1]
- if data_format == "NCHW":
- t1 = test_util.NHWCToNCHW(t1)
- strides = test_util.NHWCToNCHW(strides)
- output = fused_conv2d_bias_activation_op.fused_conv2d_bias_activation(
- t1,
- t2,
- t3,
- strides=strides,
- padding=padding,
- data_format=data_format,
- activation_mode="Relu")
-
- if data_format == "NCHW":
- output = test_util.NCHWToNHWC(output)
- return output
-
- tensors = []
- for (data_format, use_gpu) in GetTestConfigs():
- tensors.append(_SetupVal(data_format, use_gpu))
- with self.cached_session() as sess:
- values = sess.run(tensors)
- for i in range(1, len(values)):
- self.assertAllClose(values[0], values[i], rtol=1e-3, atol=1e-3)
-
- def _VerifyValues(self, tensor_in_sizes, filter_in_sizes, bias, strides,
- padding):
- tensors = []
- ref_tensors = []
- for (data_format, use_gpu) in GetTestConfigs():
- for dtype in self._DtypesToTest(use_gpu):
- for filter_format in self._FilterFormatsToTest(use_gpu):
- result, expected = self._SetupValuesForDevice(
- tensor_in_sizes, filter_in_sizes, bias, strides, padding, "Relu",
- data_format, filter_format, dtype)
- tensors.append(result)
- ref_tensors.append(expected)
- with self.cached_session() as sess:
- values = sess.run(tensors)
- ref_values = sess.run(ref_tensors)
- for i in range(len(tensors)):
- conv = tensors[i]
- value = values[i]
- ref_value = ref_values[i]
- tf_logging.info("expected = ", ref_value)
- tf_logging.info("actual = ", value)
- tol = 1e-5
- if value.dtype == np.float16:
- tol = 1e-3
- self.assertAllClose(
- np.ravel(ref_value), np.ravel(value), atol=tol, rtol=tol)
- self.assertShapeEqual(value, conv)
-
- def testConv2D1x1Filter(self, gpu_only=True):
- if gpu_only and not test.is_gpu_available():
- tf_logging.info("Skipping Conv2D1x1Filter test.")
- return
- # expected_output = [
- # 0.0, 0.0, 0.0, 21.0, 0.0, 0.0, 57.0, 0.0, 0.0, 93.0, 41.0, 0.0, 129.0,
- # 86.0, 43.0, 165.0, 131.0, 97.0
- # ]
- medians = [-45.0, -130.0, -215.0]
- self._VerifyValues(
- tensor_in_sizes=[1, 2, 3, 3],
- filter_in_sizes=[1, 1, 3, 3],
- bias=medians,
- strides=[1, 1],
- padding="VALID")
-
- def testConv2DEmpty(self, gpu_only=True):
- if gpu_only and not test.is_gpu_available():
- tf_logging.info("Skipping Conv2DEmpty test.")
- return
- # expected_output = []
- self._VerifyValues(
- tensor_in_sizes=[0, 2, 3, 3],
- filter_in_sizes=[1, 1, 3, 3],
- bias=[0.0, 0.0, 0.0],
- strides=[1, 1],
- padding="VALID")
-
- def testConv2D2x2Filter(self, gpu_only=True):
- if gpu_only and not test.is_gpu_available():
- tf_logging.info("Skipping Conv2D2x2Filter test.")
- return
- # expected_output = [0.0, 0.0, 0.0, 401.0, 533.0, 665.0]
- self._VerifyValues(
- tensor_in_sizes=[1, 2, 3, 3],
- filter_in_sizes=[2, 2, 3, 3],
- bias=[-2500.0, -2500.0, -2500.0],
- strides=[1, 1],
- padding="VALID")
-
- def testConv2D1x2Filter(self, gpu_only=True):
- if gpu_only and not test.is_gpu_available():
- tf_logging.info("Skipping Conv2D1x2Filter test.")
- return
- # expected_output = [
- # 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 190.0, 265.0, 340.0, 343.0, 436.0, 529.0
- # ]
- self._VerifyValues(
- tensor_in_sizes=[1, 2, 3, 3],
- filter_in_sizes=[1, 2, 3, 3],
- bias=[-500.0, -500.0, -500.0],
- strides=[1, 1],
- padding="VALID")
-
- def testConv2D2x2FilterStride2(self, gpu_only=True):
- if gpu_only and not test.is_gpu_available():
- tf_logging.info("Skipping Conv2D2x2FilterStride2 test.")
- return
- # expected_output = [0.0, 67.0, 163.0]
- self._VerifyValues(
- tensor_in_sizes=[1, 2, 3, 3],
- filter_in_sizes=[2, 2, 3, 3],
- bias=[-2300.0, -2300.0, -2300.0],
- strides=[2, 2],
- padding="VALID")
-
- def testConv2D2x2FilterStride2Same(self, gpu_only=True):
- if gpu_only and not test.is_gpu_available():
- tf_logging.info("Skipping Conv2D2x2FilterStride2Same test.")
- return
- # expected_output = [0.0, 2367.0, 2463.0, 1230.0, 1305.0, 1380.0]
- self._VerifyValues(
- tensor_in_sizes=[1, 2, 3, 3],
- filter_in_sizes=[2, 2, 3, 3],
- bias=[-2300.0, -1000.0, -1000.0],
- strides=[2, 2],
- padding="SAME")
-
- def testConv2D2x2FilterStride1x2(self, gpu_only=True):
- if gpu_only and not test.is_gpu_available():
- tf_logging.info("Skipping Conv2D2x2FilterStride1x2 test.")
- return
- # expected_output = [0.0, 0.0, 8.0, 28.0, 48.0, 68.0]
- self._VerifyValues(
- tensor_in_sizes=[1, 3, 6, 1],
- filter_in_sizes=[2, 2, 1, 1],
- bias=[-90.0],
- strides=[1, 2],
- padding="VALID")
-
- def testConv2DKernelSmallerThanStrideValid(self, gpu_only=True):
- if gpu_only and not test.is_gpu_available():
- tf_logging.info("Skipping Conv2DKernelSmallerThanStrideValid test.")
- return
- # expected_output = [0, 0, 175, 205]
- self._VerifyValues(
- tensor_in_sizes=[1, 7, 7, 1],
- filter_in_sizes=[2, 2, 1, 1],
- bias=[-100.0],
- strides=[3, 3],
- padding="VALID")
-
- def testConv2DKernelSmallerThanStrideSame(self, gpu_only=True):
- if gpu_only and not test.is_gpu_available():
- tf_logging.info("Skipping Conv2DKernelSmallerThanStrideSame test.")
- return
- # expected = [0, 0, 2, 4]
- self._VerifyValues(
- tensor_in_sizes=[1, 3, 3, 1],
- filter_in_sizes=[1, 1, 1, 1],
- bias=[-5.0],
- strides=[2, 2],
- padding="SAME")
-
- # expected = [0, 0, 4, 6]
- self._VerifyValues(
- tensor_in_sizes=[1, 4, 4, 1],
- filter_in_sizes=[1, 1, 1, 1],
- bias=[-5.0],
- strides=[2, 2],
- padding="SAME")
-
- # expected = [4, 0, 1, 0]
- self._VerifyValues(
- tensor_in_sizes=[1, 4, 4, 1],
- filter_in_sizes=[2, 2, 1, 1],
- bias=[-40.0],
- strides=[3, 3],
- padding="SAME")
-
- def testConv2DKernelSizeMatchesInputSize(self, gpu_only=True):
- if gpu_only and not test.is_gpu_available():
- tf_logging.info("Skipping Conv2DKernelSizeMatchesInputSize test.")
- return
- # expected = [0, 5]
- self._VerifyValues(
- tensor_in_sizes=[1, 2, 2, 1],
- filter_in_sizes=[2, 2, 1, 2],
- bias=[-50.0, -55.0],
- strides=[1, 1],
- padding="VALID")
-
- # expected = [0, 2, 282, 322]
- self._VerifyValues(
- tensor_in_sizes=[1, 8, 8, 1],
- filter_in_sizes=[2, 2, 1, 1],
- bias=[-200.0],
- strides=[4, 4],
- padding="SAME")
-
- def testShapeFunctionEdgeCases(self):
- # All shapes unknown.
- c1 = fused_conv2d_bias_activation_op.fused_conv2d_bias_activation(
- array_ops.placeholder(dtypes.float32),
- array_ops.placeholder(dtypes.float32),
- array_ops.placeholder(dtypes.float32),
- strides=[1, 1, 1, 1],
- padding="SAME",
- activation_mode="Relu")
- self.assertEqual([None, None, None, None], c1.get_shape().as_list())
-
- # Incorrect input shape.
- with self.assertRaises(ValueError):
- fused_conv2d_bias_activation_op.fused_conv2d_bias_activation(
- array_ops.placeholder(dtypes.float32, shape=[1, 3]),
- array_ops.placeholder(dtypes.float32),
- array_ops.placeholder(dtypes.float32),
- strides=[1, 1, 1, 1],
- padding="SAME",
- activation_mode="Relu")
-
- # Incorrect filter shape.
- with self.assertRaises(ValueError):
- fused_conv2d_bias_activation_op.fused_conv2d_bias_activation(
- array_ops.placeholder(dtypes.float32),
- array_ops.placeholder(dtypes.float32, shape=[1, 3]),
- array_ops.placeholder(dtypes.float32),
- strides=[1, 1, 1, 1],
- padding="SAME",
- activation_mode="Relu")
-
- # Depth mismatch.
- with self.assertRaises(ValueError):
- fused_conv2d_bias_activation_op.fused_conv2d_bias_activation(
- array_ops.placeholder(dtypes.float32, shape=[32, 20, 20, 3]),
- array_ops.placeholder(dtypes.float32, shape=[4, 4, 2, 2]),
- array_ops.placeholder(dtypes.float32),
- strides=[1, 1, 1, 1],
- padding="SAME",
- activation_mode="Relu")
-
- def testOpEdgeCases(self, gpu_only=True):
- if gpu_only and not test.is_gpu_available():
- tf_logging.info("Skipping OpEdgeCases tests.")
- return
- with self.cached_session() as sess:
- # Illegal strides.
- with self.assertRaisesRegexp(errors_impl.InvalidArgumentError,
- "Convolutional strides are not supported in "
- "the batch or depth dimensions."):
- sess.run(
- fused_conv2d_bias_activation_op.fused_conv2d_bias_activation(
- array_ops.placeholder(dtypes.float32),
- array_ops.placeholder(dtypes.float32),
- array_ops.placeholder(dtypes.float32),
- strides=[2, 1, 1, 1],
- padding="SAME",
- activation_mode="Relu"))
- with self.assertRaisesRegexp(errors_impl.InvalidArgumentError,
- "Convolutional strides are not supported in "
- "the batch or depth dimensions."):
- sess.run(
- fused_conv2d_bias_activation_op.fused_conv2d_bias_activation(
- array_ops.placeholder(dtypes.float32),
- array_ops.placeholder(dtypes.float32),
- array_ops.placeholder(dtypes.float32),
- strides=[1, 1, 1, 2],
- padding="SAME",
- activation_mode="Relu"))
-
- # Illegal activation mode.
- with self.assertRaisesRegexp(ValueError,
- "Op passed string 'Tanh' not in:"):
- sess.run(
- fused_conv2d_bias_activation_op.fused_conv2d_bias_activation(
- array_ops.placeholder(dtypes.float32),
- array_ops.placeholder(dtypes.float32),
- array_ops.placeholder(dtypes.float32),
- strides=[1, 1, 1, 1],
- padding="SAME",
- activation_mode="Tanh"))
-
- # Filter larger than input.
- with self.assertRaisesRegexp(ValueError, "Negative dimension size"):
- sess.run(
- fused_conv2d_bias_activation_op.fused_conv2d_bias_activation(
- array_ops.placeholder(dtypes.float32, shape=[32, 20, 20, 3]),
- array_ops.placeholder(dtypes.float32, shape=[20, 21, 3, 2]),
- array_ops.placeholder(dtypes.float32, shape=[2]),
- strides=[1, 1, 1, 1],
- padding="VALID",
- activation_mode="Relu"))
- with self.assertRaisesRegexp(ValueError, "Negative dimension size"):
- sess.run(
- fused_conv2d_bias_activation_op.fused_conv2d_bias_activation(
- array_ops.placeholder(dtypes.float32, shape=[32, 20, 20, 3]),
- array_ops.placeholder(dtypes.float32, shape=[21, 20, 3, 2]),
- array_ops.placeholder(dtypes.float32, shape=[2]),
- strides=[1, 1, 1, 1],
- padding="VALID",
- activation_mode="Relu"))
-
-
-def GetInceptionFwdTest(input_size, filter_size, stride, padding,
- gpu_only=True):
-
- def Test(self):
- if gpu_only and not test.is_gpu_available():
- tf_logging.info("Skipping InceptionFwd %s", (input_size, filter_size,
- stride, padding))
- return
- tf_logging.info("Testing InceptionFwd %s", (input_size, filter_size, stride,
- padding))
- self._CompareFwdValues(input_size, filter_size, [stride, stride], padding)
-
- return Test
-
-
-def CalculateConvolvedOutputDim(input_dim, filter_dim, stride, padding_type):
- """Calculates the size of an output dimension of a strided convolution.
-
- Given the sizes of the corresponding dimension of the input and filter shapes,
- and the stride and padding_types, calculates the size of the output dimension.
- This function can be called separately for each input dimension.
-
- Args:
- input_dim: An `int` specifying the size of the input dimension.
- filter_dim: An `int` specifying the size of the filter dimension.
- stride: An `int` specifying the step size of the convolution along the
- input dimension.
- padding_type: either 'VALID' or 'SAME'.
-
- Returns:
- The size of the output dimension.
- """
- if padding_type == "VALID":
- return (input_dim - filter_dim + stride) // stride
- else: # padding_type == 'SAME'
- return (input_dim + stride - 1) // stride
-
-
-def NchwVectCToNchw(in_tensor):
- # [N, C / 4, H, W, 4] => [N, C / 4, 4, H, W] == [N, C, H, W]
- t = array_ops.transpose(in_tensor, [0, 1, 4, 2, 3])
- n = in_tensor.shape.dims[0].value
- c = in_tensor.shape.dims[1].value * in_tensor.shape.dims[4].value
- h = in_tensor.shape.dims[2].value
- w = in_tensor.shape.dims[3].value
- return array_ops.reshape(t, [n, c, h, w])
-
-
-def OihwVectIToHwio(in_tensor):
- # [O, I / 4, H, W, 4] => [O, I / 4, 4, H, W] == [O, I, H, W]
- t = array_ops.transpose(in_tensor, [2, 3, 1, 4, 0])
- o = in_tensor.shape.dims[0].value
- i = in_tensor.shape.dims[1].value * in_tensor.shape.dims[4].value
- h = in_tensor.shape.dims[2].value
- w = in_tensor.shape.dims[3].value
- return array_ops.reshape(t, [h, w, i, o])
-
-
-def NchwToNchwVectC(in_tensor):
- n, c, h, w = in_tensor.shape.as_list()
- assert c % 4 == 0
- t = array_ops.reshape(in_tensor, [n, c // 4, 4, h, w])
- return array_ops.transpose(t, [0, 1, 3, 4, 2])
-
-
-def HwioToOihw(in_tensor):
- return array_ops.transpose(in_tensor, [3, 2, 0, 1])
-
-
-def SimulateFusedConv2dBiasActivationInt8(conv_input_scale, conv_input, kernel,
- padding, strides, side_input_scale,
- side_input, biases, apply_relu):
- """Simulates the int8 fused 2-D convolution op using separate float ops.
-
- The arguments and return values have the same format, meanings and
- restrictions as the actual op.
- Args:
- conv_input_scale: A scalar 'float'.
- conv_input: A `Tensor` of type `qint8` in NCHW_VECT_C layout.
- kernel: A `Tensor` of type `qint8` in OIHW_VECT_I layout.
- padding: A `string` from: `"SAME", "VALID"`.
- strides: A list of `ints`.
- side_input_scale: A scalar 'float'.
- side_input: A `Tensor` of type `qint8` in NCHW_VECT_C layout.
- biases: A `Tensor` of type `float32` in NCHW layout.
- apply_relu: A boolean to specify whether to apply "Relu" activation function
- that clips outputs to the range [0, 127], or "None" activation that clips
- to the range [-128, 127].
- Returns:
- A `Tensor` of type `qint8` in NCHW_VECT_C layout.
- """
- conv_result = nn_ops.conv2d(
- NchwVectCToNchw(gen_array_ops.dequantize(conv_input, -128, 127)),
- OihwVectIToHwio(gen_array_ops.dequantize(kernel, -128, 127)),
- strides=strides,
- padding=padding,
- data_format="NCHW") * conv_input_scale
-
- conv_and_side_inputs = conv_result + side_input_scale * NchwVectCToNchw(
- gen_array_ops.dequantize(side_input, -128, 127))
-
- output = nn_ops.bias_add(conv_and_side_inputs, biases, data_format="NCHW")
- if apply_relu:
- output = nn_ops.relu(output)
-
- result, _, _ = gen_array_ops.quantize_v2(
- NchwToNchwVectC(output), -128, 127, dtypes.qint8)
- return result
-
-
-class FusedConvInt8Tests(test.TestCase):
- _test_params = [
- {
- "batch_size": 1,
- "input_channels": 4,
- "output_channels": 4,
- "input_height": 8,
- "input_width": 8,
- "filter_height": 6,
- "filter_width": 6,
- "vertical_stride": 2,
- "horizontal_stride": 2,
- "conv_input_scale": 0.002,
- "side_input_scale": 0.0,
- "bias_scale": 1,
- "padding_type": "SAME"
- },
- {
- "batch_size": 1,
- "input_channels": 4,
- "output_channels": 4,
- "input_height": 6,
- "input_width": 6,
- "filter_height": 6,
- "filter_width": 6,
- "vertical_stride": 2,
- "horizontal_stride": 2,
- "conv_input_scale": 0.002,
- "side_input_scale": 0.0,
- "bias_scale": 1,
- "padding_type": "SAME"
- },
- {
- "batch_size": 2,
- "input_channels": 8,
- "output_channels": 16,
- "input_height": 8,
- "input_width": 8,
- "filter_height": 3,
- "filter_width": 3,
- "vertical_stride": 2,
- "horizontal_stride": 2,
- "conv_input_scale": 0.002,
- "side_input_scale": 0.0,
- "bias_scale": 1,
- "padding_type": "VALID"
- },
- {
- "batch_size": 2,
- "input_channels": 8,
- "output_channels": 16,
- "input_height": 8,
- "input_width": 8,
- "filter_height": 3,
- "filter_width": 3,
- "vertical_stride": 2,
- "horizontal_stride": 2,
- "conv_input_scale": 0.002,
- "side_input_scale": 0.0,
- "bias_scale": 1,
- "padding_type": "SAME"
- },
- {
- "batch_size": 2,
- "input_channels": 8,
- "output_channels": 16,
- "input_height": 8,
- "input_width": 8,
- "filter_height": 3,
- "filter_width": 3,
- "vertical_stride": 2,
- "horizontal_stride": 2,
- "conv_input_scale": 0.002,
- "side_input_scale": 0.5,
- "bias_scale": 1,
- "padding_type": "VALID"
- },
- {
- "batch_size": 2,
- "input_channels": 16,
- "output_channels": 16,
- "input_height": 9,
- "input_width": 9,
- "filter_height": 3,
- "filter_width": 3,
- "vertical_stride": 1,
- "horizontal_stride": 1,
- "conv_input_scale": 0.001,
- "side_input_scale": 0.5,
- "bias_scale": 1,
- "padding_type": "SAME"
- },
- {
- "batch_size": 3,
- "input_channels": 8,
- "output_channels": 8,
- "input_height": 9,
- "input_width": 9,
- "filter_height": 5,
- "filter_width": 5,
- "vertical_stride": 1,
- "horizontal_stride": 1,
- "conv_input_scale": 0.001,
- "side_input_scale": 0.5,
- "bias_scale": 1,
- "padding_type": "SAME"
- },
- {
- "batch_size": 3,
- "input_channels": 8,
- "output_channels": 8,
- "input_height": 9,
- "input_width": 9,
- "filter_height": 7,
- "filter_width": 1,
- "vertical_stride": 2,
- "horizontal_stride": 1,
- "conv_input_scale": 0.002,
- "side_input_scale": 0.5,
- "bias_scale": 1,
- "padding_type": "SAME"
- },
- {
- "batch_size": 3,
- "input_channels": 8,
- "output_channels": 8,
- "input_height": 9,
- "input_width": 9,
- "filter_height": 1,
- "filter_width": 7,
- "vertical_stride": 1,
- "horizontal_stride": 1,
- "conv_input_scale": 0.002,
- "side_input_scale": 0.5,
- "bias_scale": 1,
- "padding_type": "SAME"
- },
- ]
-
- def runTest(self, test_param, apply_relu):
- batch_size = test_param["batch_size"]
- input_channels = test_param["input_channels"]
- output_channels = test_param["output_channels"]
- input_height = test_param["input_height"]
- input_width = test_param["input_width"]
- filter_height = test_param["filter_height"]
- filter_width = test_param["filter_width"]
- vertical_stride = test_param["vertical_stride"]
- horizontal_stride = test_param["horizontal_stride"]
- conv_input_scale = test_param["conv_input_scale"]
- side_input_scale = test_param["side_input_scale"]
- bias_scale = test_param["bias_scale"]
- padding_type = test_param["padding_type"]
-
- conv_input, _, _ = gen_array_ops.quantize_v2(
- random_ops.random_uniform(
- [batch_size, input_channels // 4, input_height, input_width, 4],
- minval=-0.0,
- maxval=1.0,
- dtype=dtypes.float32), -1.0, 1.0, dtypes.qint8)
-
- kernel, _, _ = gen_array_ops.quantize_v2(
- random_ops.random_uniform(
- [
- output_channels, input_channels // 4, filter_height,
- filter_width, 4
- ],
- minval=-1.0,
- maxval=1.0,
- dtype=dtypes.float32), -1.0, 1.0, dtypes.qint8)
-
- output_height = CalculateConvolvedOutputDim(input_height, filter_height,
- vertical_stride, padding_type)
- output_width = CalculateConvolvedOutputDim(input_width, filter_width,
- horizontal_stride, padding_type)
- tf_logging.info("output_height=", output_height, ", output_width=",
- output_width)
-
- side_input, _, _ = gen_array_ops.quantize_v2(
- random_ops.random_uniform(
- [batch_size, output_channels // 4, output_height, output_width, 4],
- minval=0.0,
- maxval=1.0,
- dtype=dtypes.float32), -1.0, 1.0, dtypes.qint8)
-
- biases = random_ops.random_uniform(
- [output_channels],
- minval=-10 * bias_scale,
- maxval=20 * bias_scale,
- dtype=dtypes.float32)
-
- strides = [1, 1, vertical_stride, horizontal_stride]
-
- actual = fused_conv2d_bias_activation_op.fused_conv2d_bias_activation(
- conv_input,
- kernel,
- biases,
- strides=strides,
- padding=padding_type,
- conv_input_scale=conv_input_scale,
- side_input_scale=side_input_scale,
- side_input=side_input,
- activation_mode="Relu" if apply_relu else "None",
- data_format="NCHW_VECT_C",
- filter_format="OIHW_VECT_I")
-
- expected = SimulateFusedConv2dBiasActivationInt8(
- conv_input_scale, conv_input, kernel, padding_type, strides,
- side_input_scale, side_input, biases, apply_relu)
-
- with self.test_session(use_gpu=True) as sess:
- actual_y, expected_y = sess.run([actual, expected])
- self.assertAllClose(actual_y, expected_y, rtol=0, atol=1)
-
- def testFusedConvInt8(self):
- if not test.is_gpu_available(
- cuda_only=True, min_cuda_compute_capability=(6, 1)):
- tf_logging.info("int8 test skipped because not run with --config=cuda or "
- "no GPUs with compute capability >= 6.1 are available.")
- return
- for apply_relu in [True, False]:
- for test_param in self._test_params:
- self.runTest(test_param, apply_relu)
-
-
-if __name__ == "__main__":
- for index, (input_size_, filter_size_, output_size_, stride_,
- padding_) in enumerate(GetShrunkInceptionShapes()):
- setattr(FusedConv2DBiasActivationTest, "testInceptionFwd_" + str(index),
- GetInceptionFwdTest(input_size_, filter_size_, stride_, padding_))
-
- # TODO(b/35359731)
- # Fwd, BckInput, and BackFilter to test that for certain input parameter
- # set, winograd nonfused algorithm will be excluded from conv autotune. If
- # in such case, winograd nonfused algorithm is added as one option of the
- # conv autotune, and cuDNN version is smaller than 7, the following tests
- # will fail.
- ishape = [1, 400, 400, 1]
- fshape = [1, 1, 1, 256]
- oshape = [1, 400, 400, 256]
- setattr(FusedConv2DBiasActivationTest,
- "testInceptionFwd_No_Winograd_Nonfused",
- GetInceptionFwdTest(ishape, fshape, 1, "SAME", gpu_only=True))
+if __name__ == '__main__':
test.main()
diff --git a/tensorflow/contrib/fused_conv/python/ops/fused_conv2d_bias_activation_op_test_base.py b/tensorflow/contrib/fused_conv/python/ops/fused_conv2d_bias_activation_op_test_base.py
new file mode 100644
index 0000000..35fc65e
--- /dev/null
+++ b/tensorflow/contrib/fused_conv/python/ops/fused_conv2d_bias_activation_op_test_base.py
@@ -0,0 +1,945 @@
+# 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.
+# ==============================================================================
+
+"""Provides test suites that can be run to test fused convolutions.
+
+Each of the two test suites in this module, FusedConv2DBiasActivationTest and
+FusedConvInt8Tests, should be "instantiated" by declaring a class which inherits
+from the FusedConv test and a class that provides the standard test.TestCase
+API.
+
+See e.g. fused_conv2d_bias_activation_op_test.py in this folder.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import contextlib
+import numpy as np
+
+from tensorflow.contrib.fused_conv.python.ops import fused_conv2d_bias_activation_op
+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 gen_array_ops
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.platform import test
+from tensorflow.python.platform import tf_logging
+
+
+def _GetShrunkInceptionShapes(shrink=10):
+ """Iterator for smaller versions of convolution shapes in 2015 Inception.
+
+ Relative to inception, each depth value is `depth // shrink`.
+
+ Args:
+ shrink: Factor to shrink each depth value by relative to Inception.
+
+ Yields:
+ Tuple (input_size, filter_size, out_size, stride, padding), the convolution
+ parameters of Inception layers.
+ """
+ input_sizes = [[4, 5, 5, 1248], [4, 8, 8, 384], [4, 8, 8, 384], [
+ 4, 8, 8, 2048
+ ], [4, 8, 8, 448], [4, 8, 8, 2048], [4, 8, 8, 2048], [4, 8, 8, 2048], [
+ 4, 8, 8, 1760
+ ], [4, 8, 8, 1760], [4, 8, 8, 1760], [4, 8, 8, 1760], [4, 17, 17, 192], [
+ 4, 17, 17, 192
+ ], [4, 17, 17, 1248], [4, 17, 17, 128], [4, 17, 17, 1248], [4, 17, 17, 224], [
+ 4, 17, 17, 192
+ ], [4, 17, 17, 192], [4, 17, 17, 1216], [4, 17, 17, 1216], [4, 17, 17, 224], [
+ 4, 17, 17, 192
+ ], [4, 17, 17, 192], [4, 17, 17, 1152], [4, 17, 17, 1152], [4, 17, 17, 192], [
+ 4, 17, 17, 160
+ ], [4, 17, 17, 1152], [4, 17, 17, 1024], [4, 17, 17, 128], [4, 17, 17, 1024],
+ [4, 17, 17, 128], [4, 17, 17, 1024], [4, 17, 17, 128], [
+ 4, 17, 17, 768
+ ], [4, 17, 17, 128], [4, 17, 17, 128], [4, 17, 17, 768],
+ [4, 17, 17, 768], [4, 35, 35, 96], [4, 35, 35, 288], [
+ 4, 35, 35, 64
+ ], [4, 35, 35, 288], [4, 35, 35, 256], [4, 35, 35, 48], [
+ 4, 35, 35, 256
+ ], [4, 35, 35, 96], [4, 35, 35, 192], [4, 35, 35, 192], [
+ 4, 35, 35, 192
+ ], [4, 73, 73, 64], [4, 73, 73, 64], [4, 147, 147, 24]]
+ filter_sizes = [[1, 1, 1248, 128], [1, 3, 384, 384], [3, 1, 384, 384], [
+ 1, 1, 2048, 192
+ ], [3, 3, 448, 384], [1, 1, 2048, 320], [1, 1, 2048, 448], [1, 1, 2048, 384],
+ [1, 1, 1760, 384], [1, 1, 1760, 192], [1, 1, 1760, 448], [
+ 1, 1, 1760, 320
+ ], [3, 3, 192, 192], [3, 3, 192, 192], [1, 1, 1248, 192], [
+ 3, 3, 128, 320
+ ], [1, 1, 1248, 128], [1, 3, 224, 224], [3, 1, 192, 256], [
+ 1, 3, 192, 256
+ ], [1, 1, 1216, 192], [1, 1, 1216, 96], [3, 1, 224, 224], [
+ 3, 3, 192, 224
+ ], [1, 3, 192, 192], [1, 1, 1152, 192], [1, 1, 1152, 128], [
+ 3, 1, 192, 192
+ ], [3, 3, 160, 192], [1, 1, 1152, 160], [1, 1, 1024, 128], [
+ 1, 3, 128, 192
+ ], [1, 1, 1024, 160], [3, 1, 128, 192], [1, 1, 1024, 256], [
+ 3, 1, 128, 128
+ ], [1, 1, 768, 192], [1, 3, 128, 128], [3, 3, 128, 128], [
+ 1, 1, 768, 128
+ ], [1, 1, 768, 320], [3, 3, 96, 96], [3, 3, 288, 384], [
+ 3, 3, 64, 96
+ ], [1, 1, 288, 64], [1, 1, 256, 64], [5, 5, 48, 64],
+ [1, 1, 256, 48], [3, 3, 96, 96], [1, 1, 192, 32], [
+ 1, 1, 192, 64
+ ], [1, 1, 192, 48], [3, 3, 64, 192], [1, 1, 64,
+ 64], [1, 1, 24, 64]]
+ out_sizes = [[4, 5, 5, 128], [4, 8, 8, 384], [4, 8, 8, 384], [4, 8, 8, 192], [
+ 4, 8, 8, 384
+ ], [4, 8, 8, 320], [4, 8, 8, 448], [4, 8, 8, 384], [4, 8, 8, 384], [
+ 4, 8, 8, 192
+ ], [4, 8, 8, 448], [4, 8, 8, 320], [4, 8, 8, 192], [4, 17, 17, 192], [
+ 4, 17, 17, 192
+ ], [4, 8, 8, 320], [4, 17, 17, 128], [4, 17, 17, 224], [4, 17, 17, 256], [
+ 4, 17, 17, 256
+ ], [4, 17, 17, 192], [4, 17, 17, 96], [4, 17, 17, 224], [4, 17, 17, 224], [
+ 4, 17, 17, 192
+ ], [4, 17, 17, 192], [4, 17, 17, 128], [4, 17, 17, 192], [4, 17, 17, 192], [
+ 4, 17, 17, 160
+ ], [4, 17, 17, 128], [4, 17, 17, 192], [4, 17, 17, 160], [4, 17, 17, 192], [
+ 4, 17, 17, 256
+ ], [4, 17, 17, 128], [4, 17, 17, 192], [4, 17, 17, 128], [4, 17, 17, 128], [
+ 4, 17, 17, 128
+ ], [4, 17, 17, 320], [4, 17, 17, 96], [4, 17, 17, 384], [4, 35, 35, 96], [
+ 4, 35, 35, 64
+ ], [4, 35, 35, 64], [4, 35, 35, 64], [4, 35, 35, 48], [4, 35, 35, 96],
+ [4, 35, 35, 32], [4, 35, 35, 64], [4, 35, 35, 48],
+ [4, 71, 71, 192], [4, 73, 73, 64], [4, 147, 147, 64]]
+ strides = [
+ 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1,
+ 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 1, 1, 1, 1, 1, 1, 1,
+ 1, 1, 1, 1, 1
+ ]
+ # Shrink sizes to make the test faster
+ for i in input_sizes:
+ i[3] //= shrink
+ for f in filter_sizes:
+ f[2] //= shrink
+ f[3] //= shrink
+ for o in out_sizes:
+ o[3] //= shrink
+ # pylint: disable=invalid-name
+ VALID = "VALID"
+ SAME = "SAME"
+ # pylint: enable=invalid-name
+ paddings = [
+ SAME, SAME, SAME, SAME, SAME, SAME, SAME, SAME, SAME, SAME, SAME, SAME,
+ VALID, SAME, SAME, VALID, SAME, SAME, SAME, SAME, SAME, SAME, SAME, SAME,
+ SAME, SAME, SAME, SAME, SAME, SAME, SAME, SAME, SAME, SAME, SAME, SAME,
+ SAME, SAME, SAME, SAME, SAME, VALID, VALID, SAME, SAME, SAME, SAME, SAME,
+ SAME, SAME, SAME, SAME, VALID, VALID, VALID
+ ]
+ for i, f, o, s, p in zip(input_sizes, filter_sizes, out_sizes, strides,
+ paddings):
+ yield i, f, o, s, p
+
+
+def _GetTestConfigs():
+ """Get all the valid tests configs to run.
+
+ Returns:
+ all the valid test configs as tuples of data_format and use_gpu.
+ """
+ test_configs = [("NCHW", True), ("NHWC", True)]
+ return test_configs
+
+
+def _IotaNdF32Constant(dim_sizes):
+
+ def MakeList(dims):
+ if len(dims) == 1:
+ return [float(1 + f) for f in range(dims[0])]
+ return [MakeList(dims[1:]) for _ in range(dims[0])]
+
+ return constant_op.constant(MakeList(dim_sizes), dtype=dtypes.float32)
+
+
+def _GetInceptionFwdTest(input_size,
+ filter_size,
+ stride,
+ padding,
+ gpu_only=True):
+
+ def Test(self):
+ if gpu_only and not test.is_gpu_available():
+ tf_logging.info("Skipping InceptionFwd %s",
+ (input_size, filter_size, stride, padding))
+ return
+ tf_logging.info("Testing InceptionFwd %s",
+ (input_size, filter_size, stride, padding))
+ self.CompareFwdValues(input_size, filter_size, [stride, stride], padding)
+
+ return Test
+
+
+class FusedConv2DBiasActivationTest(object):
+
+ @contextlib.contextmanager
+ def test_scope(self): # pylint: disable=invalid-name
+ """Can be overridden in base classes to provide a test scope."""
+ yield
+
+ def _DtypesToTest(self, use_gpu):
+ return [dtypes.float32]
+
+ def _FilterFormatsToTest(self, use_gpu):
+ return ["HWIO", "OIHW"]
+
+ def _SetupValuesForDevice(self, tensor_in_sizes, filter_in_sizes, bias,
+ strides, padding, activation_mode, data_format,
+ filter_format, dtype):
+ """Verifies the output values of the convolution function.
+
+ Args:
+ tensor_in_sizes: Input tensor dimensions in
+ [batch, input_rows, input_cols, input_depth].
+ filter_in_sizes: Filter tensor dimensions in
+ [kernel_rows, kernel_cols, input_depth, output_depth].
+ bias: 1-D bias tensor of length output_depth.
+ strides: Stride: [col_stride, row_stride]
+ padding: Padding type.
+ activation_mode: Activation mode.
+ data_format: Format of the data tensors.
+ filter_format: Filter format to use for the fused convolution.
+ dtype: Data type for inputs and outputs.
+ Returns:
+ Symbolic tensor value and reference value that can be used to
+ execute the computation and verify the results.
+ """
+ input_size = np.prod(tensor_in_sizes)
+ filter_size = np.prod(filter_in_sizes)
+ bias_size = filter_in_sizes[-1] # equals to output depth
+ # Initializes the input tensor with array containing incrementing
+ # numbers from 1.
+ x1 = [f * 1.0 for f in range(1, input_size + 1)]
+ x2 = [f * 1.0 for f in range(1, filter_size + 1)]
+ # This is to guarantee that there are always negative values after
+ # bias add so that we can test whether relu works correctly.
+ x3 = bias
+ with self.cached_session(use_gpu=True), self.test_scope():
+ t1 = constant_op.constant(x1, shape=tensor_in_sizes, dtype=dtype)
+ t2 = constant_op.constant(x2, shape=filter_in_sizes, dtype=dtype)
+ fused_t2 = t2
+ if filter_format == "OIHW":
+ fused_t2 = _HwioToOihw(t2)
+ t3 = constant_op.constant(x3, shape=[bias_size], dtype=dtype)
+ strides = [1] + strides + [1]
+ if data_format == "NCHW":
+ t1 = test_util.NHWCToNCHW(t1)
+ strides = test_util.NHWCToNCHW(strides)
+ output = fused_conv2d_bias_activation_op.fused_conv2d_bias_activation(
+ t1,
+ fused_t2,
+ t3,
+ strides=strides,
+ padding=padding,
+ data_format=data_format,
+ filter_format=filter_format,
+ activation_mode=activation_mode)
+ ref_conv_output = nn_ops.conv2d(
+ t1, t2, strides=strides, padding=padding, data_format=data_format)
+ ref_bias_output = nn_ops.bias_add(
+ ref_conv_output, t3, data_format=data_format)
+ ref_output = nn_ops.relu(ref_bias_output)
+ if data_format == "NCHW":
+ output = test_util.NCHWToNHWC(output)
+ ref_output = test_util.NCHWToNHWC(ref_output)
+
+ return output, ref_output
+
+ def CompareFwdValues(self, tensor_in_sizes, filter_in_sizes, conv_strides,
+ padding):
+ """Verifies that CPU and GPU produce the same values.
+
+ Args:
+ tensor_in_sizes: Input tensor dimensions in
+ [batch, input_rows, input_cols, input_depth].
+ filter_in_sizes: Filter tensor dimensions in
+ [kernel_rows, kernel_cols, input_depth, output_depth].
+ conv_strides: [row_stride, col_stride] for the convolution;
+ padding: Padding type.
+ """
+ x1 = np.random.rand(*tensor_in_sizes).astype(np.float32)
+ x2 = np.random.rand(*filter_in_sizes).astype(np.float32)
+ x3 = np.random.rand(*[filter_in_sizes[-1]]).astype(np.float32)
+
+ def _SetupVal(data_format, use_gpu):
+ with self.cached_session(use_gpu=use_gpu), self.test_scope():
+ t1 = constant_op.constant(x1, shape=tensor_in_sizes)
+ t2 = constant_op.constant(x2, shape=filter_in_sizes)
+ t3 = constant_op.constant(x3, shape=[filter_in_sizes[-1]])
+ strides = [1] + conv_strides + [1]
+ if data_format == "NCHW":
+ t1 = test_util.NHWCToNCHW(t1)
+ strides = test_util.NHWCToNCHW(strides)
+ output = fused_conv2d_bias_activation_op.fused_conv2d_bias_activation(
+ t1,
+ t2,
+ t3,
+ strides=strides,
+ padding=padding,
+ data_format=data_format,
+ activation_mode="Relu")
+
+ if data_format == "NCHW":
+ output = test_util.NCHWToNHWC(output)
+ return output
+
+ tensors = []
+ for (data_format, use_gpu) in _GetTestConfigs():
+ tensors.append(_SetupVal(data_format, use_gpu))
+ with self.cached_session() as sess, self.test_scope():
+ values = sess.run(tensors)
+ for i in range(1, len(values)):
+ self.assertAllClose(values[0], values[i], rtol=1e-3, atol=1e-3)
+
+ def _VerifyValues(self, tensor_in_sizes, filter_in_sizes, bias, strides,
+ padding):
+ tensors = []
+ ref_tensors = []
+ for (data_format, use_gpu) in _GetTestConfigs():
+ for dtype in self._DtypesToTest(use_gpu):
+ for filter_format in self._FilterFormatsToTest(use_gpu):
+ result, expected = self._SetupValuesForDevice(
+ tensor_in_sizes, filter_in_sizes, bias, strides, padding, "Relu",
+ data_format, filter_format, dtype)
+ tensors.append(result)
+ ref_tensors.append(expected)
+ with self.cached_session() as sess, self.test_scope():
+ values = sess.run(tensors)
+ ref_values = sess.run(ref_tensors)
+ for i in range(len(tensors)):
+ conv = tensors[i]
+ value = values[i]
+ ref_value = ref_values[i]
+ tf_logging.info("expected = %s", ref_value)
+ tf_logging.info("actual = %s", value)
+ tol = 1e-5
+ if value.dtype == np.float16:
+ tol = 1e-3
+ self.assertAllClose(
+ np.ravel(ref_value), np.ravel(value), atol=tol, rtol=tol)
+ self.assertShapeEqual(value, conv)
+
+ def testConv2D1x1Filter(self, gpu_only=True):
+ if gpu_only and not test.is_gpu_available():
+ tf_logging.info("Skipping Conv2D1x1Filter test.")
+ return
+ # expected_output = [
+ # 0.0, 0.0, 0.0, 21.0, 0.0, 0.0, 57.0, 0.0, 0.0, 93.0, 41.0, 0.0, 129.0,
+ # 86.0, 43.0, 165.0, 131.0, 97.0
+ # ]
+ medians = [-45.0, -130.0, -215.0]
+ self._VerifyValues(
+ tensor_in_sizes=[1, 2, 3, 3],
+ filter_in_sizes=[1, 1, 3, 3],
+ bias=medians,
+ strides=[1, 1],
+ padding="VALID")
+
+ def testConv2DEmpty(self, gpu_only=True):
+ if gpu_only and not test.is_gpu_available():
+ tf_logging.info("Skipping Conv2DEmpty test.")
+ return
+ # expected_output = []
+ self._VerifyValues(
+ tensor_in_sizes=[0, 2, 3, 3],
+ filter_in_sizes=[1, 1, 3, 3],
+ bias=[0.0, 0.0, 0.0],
+ strides=[1, 1],
+ padding="VALID")
+
+ def testConv2D2x2Filter(self, gpu_only=True):
+ if gpu_only and not test.is_gpu_available():
+ tf_logging.info("Skipping Conv2D2x2Filter test.")
+ return
+ # expected_output = [0.0, 0.0, 0.0, 401.0, 533.0, 665.0]
+ self._VerifyValues(
+ tensor_in_sizes=[1, 2, 3, 3],
+ filter_in_sizes=[2, 2, 3, 3],
+ bias=[-2500.0, -2500.0, -2500.0],
+ strides=[1, 1],
+ padding="VALID")
+
+ def testConv2D1x2Filter(self, gpu_only=True):
+ if gpu_only and not test.is_gpu_available():
+ tf_logging.info("Skipping Conv2D1x2Filter test.")
+ return
+ # expected_output = [
+ # 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 190.0, 265.0, 340.0, 343.0, 436.0, 529.0
+ # ]
+ self._VerifyValues(
+ tensor_in_sizes=[1, 2, 3, 3],
+ filter_in_sizes=[1, 2, 3, 3],
+ bias=[-500.0, -500.0, -500.0],
+ strides=[1, 1],
+ padding="VALID")
+
+ def testConv2D2x2FilterStride2(self, gpu_only=True):
+ if gpu_only and not test.is_gpu_available():
+ tf_logging.info("Skipping Conv2D2x2FilterStride2 test.")
+ return
+ # expected_output = [0.0, 67.0, 163.0]
+ self._VerifyValues(
+ tensor_in_sizes=[1, 2, 3, 3],
+ filter_in_sizes=[2, 2, 3, 3],
+ bias=[-2300.0, -2300.0, -2300.0],
+ strides=[2, 2],
+ padding="VALID")
+
+ def testConv2D2x2FilterStride2Same(self, gpu_only=True):
+ if gpu_only and not test.is_gpu_available():
+ tf_logging.info("Skipping Conv2D2x2FilterStride2Same test.")
+ return
+ # expected_output = [0.0, 2367.0, 2463.0, 1230.0, 1305.0, 1380.0]
+ self._VerifyValues(
+ tensor_in_sizes=[1, 2, 3, 3],
+ filter_in_sizes=[2, 2, 3, 3],
+ bias=[-2300.0, -1000.0, -1000.0],
+ strides=[2, 2],
+ padding="SAME")
+
+ def testConv2D2x2FilterStride1x2(self, gpu_only=True):
+ if gpu_only and not test.is_gpu_available():
+ tf_logging.info("Skipping Conv2D2x2FilterStride1x2 test.")
+ return
+ # expected_output = [0.0, 0.0, 8.0, 28.0, 48.0, 68.0]
+ self._VerifyValues(
+ tensor_in_sizes=[1, 3, 6, 1],
+ filter_in_sizes=[2, 2, 1, 1],
+ bias=[-90.0],
+ strides=[1, 2],
+ padding="VALID")
+
+ def testConv2DKernelSmallerThanStrideValid(self, gpu_only=True):
+ if gpu_only and not test.is_gpu_available():
+ tf_logging.info("Skipping Conv2DKernelSmallerThanStrideValid test.")
+ return
+ # expected_output = [0, 0, 175, 205]
+ self._VerifyValues(
+ tensor_in_sizes=[1, 7, 7, 1],
+ filter_in_sizes=[2, 2, 1, 1],
+ bias=[-100.0],
+ strides=[3, 3],
+ padding="VALID")
+
+ def testConv2DKernelSmallerThanStrideSame(self, gpu_only=True):
+ if gpu_only and not test.is_gpu_available():
+ tf_logging.info("Skipping Conv2DKernelSmallerThanStrideSame test.")
+ return
+ # expected = [0, 0, 2, 4]
+ self._VerifyValues(
+ tensor_in_sizes=[1, 3, 3, 1],
+ filter_in_sizes=[1, 1, 1, 1],
+ bias=[-5.0],
+ strides=[2, 2],
+ padding="SAME")
+
+ # expected = [0, 0, 4, 6]
+ self._VerifyValues(
+ tensor_in_sizes=[1, 4, 4, 1],
+ filter_in_sizes=[1, 1, 1, 1],
+ bias=[-5.0],
+ strides=[2, 2],
+ padding="SAME")
+
+ # expected = [4, 0, 1, 0]
+ self._VerifyValues(
+ tensor_in_sizes=[1, 4, 4, 1],
+ filter_in_sizes=[2, 2, 1, 1],
+ bias=[-40.0],
+ strides=[3, 3],
+ padding="SAME")
+
+ def testConv2DKernelSizeMatchesInputSize(self, gpu_only=True):
+ if gpu_only and not test.is_gpu_available():
+ tf_logging.info("Skipping Conv2DKernelSizeMatchesInputSize test.")
+ return
+ # expected = [0, 5]
+ self._VerifyValues(
+ tensor_in_sizes=[1, 2, 2, 1],
+ filter_in_sizes=[2, 2, 1, 2],
+ bias=[-50.0, -55.0],
+ strides=[1, 1],
+ padding="VALID")
+
+ # expected = [0, 2, 282, 322]
+ self._VerifyValues(
+ tensor_in_sizes=[1, 8, 8, 1],
+ filter_in_sizes=[2, 2, 1, 1],
+ bias=[-200.0],
+ strides=[4, 4],
+ padding="SAME")
+
+ def testShapeFunctionEdgeCases(self):
+ # All shapes unknown.
+ c1 = fused_conv2d_bias_activation_op.fused_conv2d_bias_activation(
+ array_ops.placeholder(dtypes.float32),
+ array_ops.placeholder(dtypes.float32),
+ array_ops.placeholder(dtypes.float32),
+ strides=[1, 1, 1, 1],
+ padding="SAME",
+ activation_mode="Relu")
+ self.assertEqual([None, None, None, None], c1.get_shape().as_list())
+
+ # Incorrect input shape.
+ with self.assertRaises(ValueError):
+ fused_conv2d_bias_activation_op.fused_conv2d_bias_activation(
+ array_ops.placeholder(dtypes.float32, shape=[1, 3]),
+ array_ops.placeholder(dtypes.float32),
+ array_ops.placeholder(dtypes.float32),
+ strides=[1, 1, 1, 1],
+ padding="SAME",
+ activation_mode="Relu")
+
+ # Incorrect filter shape.
+ with self.assertRaises(ValueError):
+ fused_conv2d_bias_activation_op.fused_conv2d_bias_activation(
+ array_ops.placeholder(dtypes.float32),
+ array_ops.placeholder(dtypes.float32, shape=[1, 3]),
+ array_ops.placeholder(dtypes.float32),
+ strides=[1, 1, 1, 1],
+ padding="SAME",
+ activation_mode="Relu")
+
+ # Depth mismatch.
+ with self.assertRaises(ValueError):
+ fused_conv2d_bias_activation_op.fused_conv2d_bias_activation(
+ array_ops.placeholder(dtypes.float32, shape=[32, 20, 20, 3]),
+ array_ops.placeholder(dtypes.float32, shape=[4, 4, 2, 2]),
+ array_ops.placeholder(dtypes.float32),
+ strides=[1, 1, 1, 1],
+ padding="SAME",
+ activation_mode="Relu")
+
+ def testOpEdgeCases(self, gpu_only=True):
+ if gpu_only and not test.is_gpu_available():
+ tf_logging.info("Skipping OpEdgeCases tests.")
+ return
+ with self.cached_session() as sess, self.test_scope():
+ # Illegal strides.
+ with self.assertRaisesRegexp(
+ errors_impl.UnimplementedError,
+ ".*strides.*in the batch and depth dimensions"):
+ sess.run(
+ fused_conv2d_bias_activation_op.fused_conv2d_bias_activation(
+ _IotaNdF32Constant([1, 1, 1, 1]),
+ _IotaNdF32Constant([1, 1, 1, 1]),
+ _IotaNdF32Constant([1]),
+ strides=[2, 1, 1, 1],
+ padding="SAME",
+ activation_mode="Relu"))
+ with self.assertRaisesRegexp(
+ errors_impl.UnimplementedError,
+ ".*strides.*in the batch and depth dimensions"):
+ sess.run(
+ fused_conv2d_bias_activation_op.fused_conv2d_bias_activation(
+ _IotaNdF32Constant([1, 1, 1, 1]),
+ _IotaNdF32Constant([1, 1, 1, 1]),
+ _IotaNdF32Constant([1]),
+ strides=[1, 1, 1, 2],
+ padding="SAME",
+ activation_mode="Relu"))
+
+ # Illegal activation mode.
+ with self.assertRaisesRegexp(ValueError,
+ "Op passed string 'Tanh' not in:"):
+ sess.run(
+ fused_conv2d_bias_activation_op.fused_conv2d_bias_activation(
+ _IotaNdF32Constant([1, 1, 1, 1]),
+ _IotaNdF32Constant([1, 1, 1, 1]),
+ _IotaNdF32Constant([1]),
+ strides=[1, 1, 1, 1],
+ padding="SAME",
+ activation_mode="Tanh"))
+
+ # Filter larger than input.
+ with self.assertRaisesRegexp(ValueError, "Negative dimension size"):
+ sess.run(
+ fused_conv2d_bias_activation_op.fused_conv2d_bias_activation(
+ _IotaNdF32Constant([32, 20, 20, 3]),
+ _IotaNdF32Constant([20, 21, 3, 2]),
+ _IotaNdF32Constant([2]),
+ strides=[1, 1, 1, 1],
+ padding="VALID",
+ activation_mode="Relu"))
+ with self.assertRaisesRegexp(ValueError, "Negative dimension size"):
+ sess.run(
+ fused_conv2d_bias_activation_op.fused_conv2d_bias_activation(
+ _IotaNdF32Constant([32, 20, 20, 3]),
+ _IotaNdF32Constant([21, 20, 3, 2]),
+ _IotaNdF32Constant([2]),
+ strides=[1, 1, 1, 1],
+ padding="VALID",
+ activation_mode="Relu"))
+
+
+# Add InceptionFwd tests to FusedConv2DBiasActivationTest.
+for index, (input_size_, filter_size_, output_size_, stride_,
+ padding_) in enumerate(_GetShrunkInceptionShapes()):
+ setattr(FusedConv2DBiasActivationTest, "testInceptionFwd_" + str(index),
+ _GetInceptionFwdTest(input_size_, filter_size_, stride_, padding_))
+
+# TODO(b/35359731)
+# Fwd, BckInput, and BackFilter to test that for certain input parameter
+# set, winograd nonfused algorithm will be excluded from conv autotune. If
+# in such case, winograd nonfused algorithm is added as one option of the
+# conv autotune, and cuDNN version is smaller than 7, the following tests
+# will fail.
+ishape = [1, 400, 400, 1]
+fshape = [1, 1, 1, 256]
+oshape = [1, 400, 400, 256]
+setattr(FusedConv2DBiasActivationTest, "testInceptionFwd_No_Winograd_Nonfused",
+ _GetInceptionFwdTest(ishape, fshape, 1, "SAME", gpu_only=True))
+
+
+def _CalculateConvolvedOutputDim(input_dim, filter_dim, stride, padding_type):
+ """Calculates the size of an output dimension of a strided convolution.
+
+ Given the sizes of the corresponding dimension of the input and filter shapes,
+ and the stride and padding_types, calculates the size of the output dimension.
+ This function can be called separately for each input dimension.
+
+ Args:
+ input_dim: An `int` specifying the size of the input dimension.
+ filter_dim: An `int` specifying the size of the filter dimension.
+ stride: An `int` specifying the step size of the convolution along the
+ input dimension.
+ padding_type: either 'VALID' or 'SAME'.
+
+ Returns:
+ The size of the output dimension.
+ """
+ if padding_type == "VALID":
+ return (input_dim - filter_dim + stride) // stride
+ else: # padding_type == 'SAME'
+ return (input_dim + stride - 1) // stride
+
+
+def _NchwVectCToNchw(in_tensor):
+ # [N, C / 4, H, W, 4] => [N, C / 4, 4, H, W] == [N, C, H, W]
+ t = array_ops.transpose(in_tensor, [0, 1, 4, 2, 3])
+ n = in_tensor.shape.dims[0].value
+ c = in_tensor.shape.dims[1].value * in_tensor.shape.dims[4].value
+ h = in_tensor.shape.dims[2].value
+ w = in_tensor.shape.dims[3].value
+ return array_ops.reshape(t, [n, c, h, w])
+
+
+def _OihwVectIToHwio(in_tensor):
+ # [O, I / 4, H, W, 4] => [O, I / 4, 4, H, W] == [O, I, H, W]
+ t = array_ops.transpose(in_tensor, [2, 3, 1, 4, 0])
+ o = in_tensor.shape.dims[0].value
+ i = in_tensor.shape.dims[1].value * in_tensor.shape.dims[4].value
+ h = in_tensor.shape.dims[2].value
+ w = in_tensor.shape.dims[3].value
+ return array_ops.reshape(t, [h, w, i, o])
+
+
+def _NchwToNchwVectC(in_tensor):
+ n, c, h, w = in_tensor.shape.as_list()
+ assert c % 4 == 0
+ t = array_ops.reshape(in_tensor, [n, c // 4, 4, h, w])
+ return array_ops.transpose(t, [0, 1, 3, 4, 2])
+
+
+def _HwioToOihw(in_tensor):
+ return array_ops.transpose(in_tensor, [3, 2, 0, 1])
+
+
+def _SimulateFusedConv2dBiasActivationInt8(conv_input_scale, conv_input, kernel,
+ padding, strides, side_input_scale,
+ side_input, biases, apply_relu):
+ """Simulates the int8 fused 2-D convolution op using separate float ops.
+
+ The arguments and return values have the same format, meanings and
+ restrictions as the actual op.
+ Args:
+ conv_input_scale: A scalar 'float'.
+ conv_input: A `Tensor` of type `qint8` in NCHW_VECT_C layout.
+ kernel: A `Tensor` of type `qint8` in OIHW_VECT_I layout.
+ padding: A `string` from: `"SAME", "VALID"`.
+ strides: A list of `ints`.
+ side_input_scale: A scalar 'float'.
+ side_input: A `Tensor` of type `qint8` in NCHW_VECT_C layout.
+ biases: A `Tensor` of type `float32` in NCHW layout.
+ apply_relu: A boolean to specify whether to apply "Relu" activation function
+ that clips outputs to the range [0, 127], or "None" activation that clips
+ to the range [-128, 127].
+ Returns:
+ A `Tensor` of type `qint8` in NCHW_VECT_C layout.
+ """
+ conv_result = nn_ops.conv2d(
+ _NchwVectCToNchw(gen_array_ops.dequantize(conv_input, -128, 127)),
+ _OihwVectIToHwio(gen_array_ops.dequantize(kernel, -128, 127)),
+ strides=strides,
+ padding=padding,
+ data_format="NCHW") * conv_input_scale
+
+ conv_and_side_inputs = conv_result + side_input_scale * _NchwVectCToNchw(
+ gen_array_ops.dequantize(side_input, -128, 127))
+
+ output = nn_ops.bias_add(conv_and_side_inputs, biases, data_format="NCHW")
+ if apply_relu:
+ output = nn_ops.relu(output)
+
+ result, _, _ = gen_array_ops.quantize_v2(
+ _NchwToNchwVectC(output), -128, 127, dtypes.qint8)
+ return result
+
+
+# TODO(b/114580749): XLA:CPU/GPU don't support int8 at the moment, so this test
+# doesn't currently use XLA.
+class FusedConvInt8Tests(object):
+ _test_params = [
+ {
+ "batch_size": 1,
+ "input_channels": 4,
+ "output_channels": 4,
+ "input_height": 8,
+ "input_width": 8,
+ "filter_height": 6,
+ "filter_width": 6,
+ "vertical_stride": 2,
+ "horizontal_stride": 2,
+ "conv_input_scale": 0.002,
+ "side_input_scale": 0.0,
+ "bias_scale": 1,
+ "padding_type": "SAME"
+ },
+ {
+ "batch_size": 1,
+ "input_channels": 4,
+ "output_channels": 4,
+ "input_height": 6,
+ "input_width": 6,
+ "filter_height": 6,
+ "filter_width": 6,
+ "vertical_stride": 2,
+ "horizontal_stride": 2,
+ "conv_input_scale": 0.002,
+ "side_input_scale": 0.0,
+ "bias_scale": 1,
+ "padding_type": "SAME"
+ },
+ {
+ "batch_size": 2,
+ "input_channels": 8,
+ "output_channels": 16,
+ "input_height": 8,
+ "input_width": 8,
+ "filter_height": 3,
+ "filter_width": 3,
+ "vertical_stride": 2,
+ "horizontal_stride": 2,
+ "conv_input_scale": 0.002,
+ "side_input_scale": 0.0,
+ "bias_scale": 1,
+ "padding_type": "VALID"
+ },
+ {
+ "batch_size": 2,
+ "input_channels": 8,
+ "output_channels": 16,
+ "input_height": 8,
+ "input_width": 8,
+ "filter_height": 3,
+ "filter_width": 3,
+ "vertical_stride": 2,
+ "horizontal_stride": 2,
+ "conv_input_scale": 0.002,
+ "side_input_scale": 0.0,
+ "bias_scale": 1,
+ "padding_type": "SAME"
+ },
+ {
+ "batch_size": 2,
+ "input_channels": 8,
+ "output_channels": 16,
+ "input_height": 8,
+ "input_width": 8,
+ "filter_height": 3,
+ "filter_width": 3,
+ "vertical_stride": 2,
+ "horizontal_stride": 2,
+ "conv_input_scale": 0.002,
+ "side_input_scale": 0.5,
+ "bias_scale": 1,
+ "padding_type": "VALID"
+ },
+ {
+ "batch_size": 2,
+ "input_channels": 16,
+ "output_channels": 16,
+ "input_height": 9,
+ "input_width": 9,
+ "filter_height": 3,
+ "filter_width": 3,
+ "vertical_stride": 1,
+ "horizontal_stride": 1,
+ "conv_input_scale": 0.001,
+ "side_input_scale": 0.5,
+ "bias_scale": 1,
+ "padding_type": "SAME"
+ },
+ {
+ "batch_size": 3,
+ "input_channels": 8,
+ "output_channels": 8,
+ "input_height": 9,
+ "input_width": 9,
+ "filter_height": 5,
+ "filter_width": 5,
+ "vertical_stride": 1,
+ "horizontal_stride": 1,
+ "conv_input_scale": 0.001,
+ "side_input_scale": 0.5,
+ "bias_scale": 1,
+ "padding_type": "SAME"
+ },
+ {
+ "batch_size": 3,
+ "input_channels": 8,
+ "output_channels": 8,
+ "input_height": 9,
+ "input_width": 9,
+ "filter_height": 7,
+ "filter_width": 1,
+ "vertical_stride": 2,
+ "horizontal_stride": 1,
+ "conv_input_scale": 0.002,
+ "side_input_scale": 0.5,
+ "bias_scale": 1,
+ "padding_type": "SAME"
+ },
+ {
+ "batch_size": 3,
+ "input_channels": 8,
+ "output_channels": 8,
+ "input_height": 9,
+ "input_width": 9,
+ "filter_height": 1,
+ "filter_width": 7,
+ "vertical_stride": 1,
+ "horizontal_stride": 1,
+ "conv_input_scale": 0.002,
+ "side_input_scale": 0.5,
+ "bias_scale": 1,
+ "padding_type": "SAME"
+ },
+ ]
+
+ @contextlib.contextmanager
+ def test_scope(self): # pylint: disable=invalid-name
+ """Can be overridden in base classes to provide a test scope."""
+ yield
+
+ def runTest(self, test_param, apply_relu):
+ batch_size = test_param["batch_size"]
+ input_channels = test_param["input_channels"]
+ output_channels = test_param["output_channels"]
+ input_height = test_param["input_height"]
+ input_width = test_param["input_width"]
+ filter_height = test_param["filter_height"]
+ filter_width = test_param["filter_width"]
+ vertical_stride = test_param["vertical_stride"]
+ horizontal_stride = test_param["horizontal_stride"]
+ conv_input_scale = test_param["conv_input_scale"]
+ side_input_scale = test_param["side_input_scale"]
+ bias_scale = test_param["bias_scale"]
+ padding_type = test_param["padding_type"]
+
+ with self.cached_session(use_gpu=True) as sess, self.test_scope():
+ conv_input, _, _ = gen_array_ops.quantize_v2(
+ random_ops.random_uniform(
+ [batch_size, input_channels // 4, input_height, input_width, 4],
+ minval=-0.0,
+ maxval=1.0,
+ dtype=dtypes.float32), -1.0, 1.0, dtypes.qint8)
+
+ kernel, _, _ = gen_array_ops.quantize_v2(
+ random_ops.random_uniform([
+ output_channels, input_channels // 4, filter_height, filter_width,
+ 4
+ ],
+ minval=-1.0,
+ maxval=1.0,
+ dtype=dtypes.float32), -1.0, 1.0,
+ dtypes.qint8)
+
+ output_height = _CalculateConvolvedOutputDim(
+ input_height, filter_height, vertical_stride, padding_type)
+ output_width = _CalculateConvolvedOutputDim(
+ input_width, filter_width, horizontal_stride, padding_type)
+ tf_logging.info("output_height=%s, output_width=%s", output_height,
+ output_width)
+
+ side_input, _, _ = gen_array_ops.quantize_v2(
+ random_ops.random_uniform([
+ batch_size, output_channels // 4, output_height, output_width, 4
+ ],
+ minval=0.0,
+ maxval=1.0,
+ dtype=dtypes.float32), -1.0, 1.0,
+ dtypes.qint8)
+
+ biases = random_ops.random_uniform([output_channels],
+ minval=-10 * bias_scale,
+ maxval=20 * bias_scale,
+ dtype=dtypes.float32)
+
+ strides = [1, 1, vertical_stride, horizontal_stride]
+
+ actual = fused_conv2d_bias_activation_op.fused_conv2d_bias_activation(
+ conv_input,
+ kernel,
+ biases,
+ strides=strides,
+ padding=padding_type,
+ conv_input_scale=conv_input_scale,
+ side_input_scale=side_input_scale,
+ side_input=side_input,
+ activation_mode="Relu" if apply_relu else "None",
+ data_format="NCHW_VECT_C",
+ filter_format="OIHW_VECT_I")
+
+ expected = _SimulateFusedConv2dBiasActivationInt8(
+ conv_input_scale, conv_input, kernel, padding_type, strides,
+ side_input_scale, side_input, biases, apply_relu)
+
+ actual_y, expected_y = sess.run([actual, expected])
+ self.assertAllClose(actual_y, expected_y, rtol=0, atol=1)
+
+ def testFusedConvInt8(self):
+ if not test.is_gpu_available(
+ cuda_only=True, min_cuda_compute_capability=(6, 1)):
+ tf_logging.info("int8 test skipped because not run with --config=cuda or "
+ "no GPUs with compute capability >= 6.1 are available.")
+ return
+ for apply_relu in [True, False]:
+ for test_param in self._test_params:
+ self.runTest(test_param, apply_relu)
+
+
+if __name__ == "__main__":
+ test.main()
diff --git a/tensorflow/contrib/graph_editor/tests/transform_test.py b/tensorflow/contrib/graph_editor/tests/transform_test.py
index 97f38c9..0ebcdc2 100644
--- a/tensorflow/contrib/graph_editor/tests/transform_test.py
+++ b/tensorflow/contrib/graph_editor/tests/transform_test.py
@@ -214,7 +214,7 @@
def test_graph_replace_gradients(self):
ops.reset_default_graph()
- w = variables.Variable(0.0, name="w")
+ w = variables.VariableV1(0.0, name="w")
y = math_ops.multiply(math_ops.multiply(w, w, name="mul1"), w, name="mul2")
g = gradients_impl.gradients(y, w, name="grad")[0]
diff --git a/tensorflow/contrib/ignite/BUILD b/tensorflow/contrib/ignite/BUILD
new file mode 100644
index 0000000..9393b70
--- /dev/null
+++ b/tensorflow/contrib/ignite/BUILD
@@ -0,0 +1,139 @@
+package(default_visibility = ["//tensorflow:internal"])
+
+licenses(["notice"]) # Apache 2.0
+
+exports_files(["LICENSE"])
+
+load(
+ "//tensorflow:tensorflow.bzl",
+ "if_not_windows",
+ "if_windows",
+ "tf_custom_op_library",
+ "tf_custom_op_py_library",
+ "tf_gen_op_libs",
+ "tf_gen_op_wrapper_py",
+ "tf_kernel_library",
+ "tf_py_test",
+)
+
+py_library(
+ name = "ignite",
+ srcs = ["__init__.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ ":dataset_ops",
+ ],
+)
+
+tf_custom_op_library(
+ name = "_dataset_ops.so",
+ srcs = ["ops/dataset_ops.cc"],
+ deps = [":dataset_kernels"],
+)
+
+tf_gen_op_libs(
+ op_lib_names = ["dataset_ops"],
+)
+
+cc_library(
+ name = "dataset_kernels",
+ srcs = [
+ "kernels/ignite_dataset_ops.cc",
+ "kernels/ignite_client.h",
+ "kernels/ignite_byte_swapper.h",
+ "kernels/ignite_plain_client.h",
+ "kernels/ignite_ssl_wrapper.h",
+ "kernels/ignite_ssl_wrapper.cc",
+ "kernels/ignite_binary_object_parser.h",
+ "kernels/ignite_binary_object_parser.cc",
+ "kernels/ignite_dataset.h",
+ "kernels/ignite_dataset.cc",
+ "kernels/ignite_dataset_iterator.h",
+ "kernels/ignite_dataset_iterator.cc",
+ ] + if_not_windows([
+ "kernels/ignite_plain_client_unix.cc",
+ ]) + if_windows([
+ "kernels/ignite_plain_client_windows.cc",
+ ]),
+ copts = if_windows([
+ "-DWIN32_LEAN_AND_MEAN",
+ ]),
+ deps = [
+ "//tensorflow/core:framework_headers_lib",
+ "//third_party/eigen3",
+ "@boringssl//:ssl",
+ "@protobuf_archive//:protobuf_headers",
+ ],
+ alwayslink = 1,
+)
+
+py_library(
+ name = "dataset_ops",
+ srcs = [
+ "python/ops/ignite_dataset_ops.py",
+ ],
+ srcs_version = "PY2AND3",
+ deps = [
+ ":ignite_op_loader",
+ "//tensorflow/python:dataset_ops_gen",
+ "//tensorflow/python:util",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/util:nest",
+ ],
+)
+
+tf_gen_op_wrapper_py(
+ name = "gen_dataset_ops",
+ out = "python/ops/gen_dataset_ops.py",
+ deps = ["//tensorflow/contrib/ignite:dataset_ops_op_lib"],
+)
+
+tf_kernel_library(
+ name = "dataset_ops_kernels",
+ deps = [
+ ":dataset_kernels",
+ "//tensorflow/core:framework",
+ ],
+ alwayslink = 1,
+)
+
+tf_custom_op_py_library(
+ name = "ignite_op_loader",
+ srcs = ["python/ops/ignite_op_loader.py"],
+ dso = ["//tensorflow/contrib/ignite:_dataset_ops.so"],
+ kernels = [
+ ":dataset_ops_kernels",
+ "//tensorflow/contrib/ignite:dataset_ops_op_lib",
+ ],
+ srcs_version = "PY2AND3",
+ deps = [
+ ":gen_dataset_ops",
+ "//tensorflow/contrib/util:util_py",
+ "//tensorflow/python:platform",
+ ],
+)
+
+# The Apache Ignite servers have to setup before the test and tear down
+# after the test manually. The docker engine has to be installed.
+#
+# To setup Apache Ignite servers:
+# $ bash ./python/tests/start_ignite.sh
+#
+# To tear down Apache Ignite servers:
+# $ bash ./python/tests/stop_ignite.sh
+tf_py_test(
+ name = "ignite_dataset_test",
+ srcs = ["python/tests/ignite_dataset_test.py"],
+ additional_deps = [
+ ":ignite",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:framework",
+ "//tensorflow/python:framework_test_lib",
+ "//tensorflow/python:platform_test",
+ ],
+ tags = [
+ "manual",
+ "no_windows",
+ "notap",
+ ],
+)
diff --git a/tensorflow/contrib/ignite/README.md b/tensorflow/contrib/ignite/README.md
new file mode 100644
index 0000000..55c89d2
--- /dev/null
+++ b/tensorflow/contrib/ignite/README.md
@@ -0,0 +1,167 @@
+# Ignite Dataset
+
+- [Overview](#overview)
+- [Features](#features)
+ * [Distributed In-Memory Datasource](#distributed-in-memory-datasource)
+ * [Structured Objects](#structured-objects)
+ * [Distributed Training](#distributed-training)
+ * [SSL Connection](#ssl-connection)
+ * [Windows Support](#windows-support)
+- [Try it out](#try-it-out)
+- [Limitations](#limitations)
+
+## Overview
+
+[Apache Ignite](https://ignite.apache.org/) is a memory-centric distributed database, 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](https://www.tensorflow.org/api_docs/python/tf/data) from TensorFlow side and [Binary Client Protocol](https://apacheignite.readme.io/v2.6/docs/binary-client-protocol) from Apache Ignite side. It allows to use Apache Ignite as a data source for neural network training, inference and all other computations supported by TensorFlow.
+
+## Features
+
+Ignite Dataset provides features that that you can use in a wide range of cases. The most important and interesting features are described below.
+
+### Distributed In-Memory Datasource
+[Apache Ignite](https://ignite.apache.org/) is a distributed in-memory database, caching, and processing platform that provides fast data access. It allows you to avoid limitations of hard drive and store and operate with as much data as you need in distributed cluster. You can utilize
+these benefits of Apache Ignite by using Ignite Dataset. Moreover, Ignite Dataset can be used for the following use-cases:
+- If you have a **gigabyte** of data you can keep it on a single machine on a hard drive, but you will face with hard drive speed limitations. At the same time, you can store your data in Apache Ignite on the same machine and use it as a datasource for TensorFlow and thus avoid these limitations.
+- If you have a **terabyte** of data you probably still can keep it on a single machine on a hard drive, but you will face with hard drive speed limitations again. At the same time, you can store your data in Apache Ignite distributed in-memory cluster and use it as a datasource for TensorFlow and thus avoid these limitations.
+- If you have a **petabyte** of data you can't keep it on a single machine. At the same time, you can store your data in Apache Ignite distributed in-memory cluster and use it as a datasource for TensorFlow.
+
+Note that Apache Ignite is not just a step of ETL pipeline between a database or a data warehouse and TensorFlow. Apache Ignite is a high-grade database itself. By choosing Apache Ignite and TensorFlow you are getting everything you need to work with operational or historical data and, at the same time, an ability to use this data for neural network training and inference.
+
+```bash
+$ apache-ignite-fabric/bin/ignite.sh
+$ apache-ignite-fabric/bin/sqlline.sh -u "jdbc:ignite:thin://localhost:10800/"
+
+jdbc:ignite:thin://localhost/> CREATE TABLE KITTEN_CACHE (ID LONG PRIMARY KEY, NAME VARCHAR);
+jdbc:ignite:thin://localhost/> INSERT INTO KITTEN_CACHE VALUES (1, 'WARM KITTY');
+jdbc:ignite:thin://localhost/> INSERT INTO KITTEN_CACHE VALUES (2, 'SOFT KITTY');
+jdbc:ignite:thin://localhost/> INSERT INTO KITTEN_CACHE VALUES (3, 'LITTLE BALL OF FUR');
+```
+
+```python
+>>> import tensorflow as tf
+>>> from tensorflow.contrib.ignite import IgniteDataset
+>>>
+>>> dataset = IgniteDataset(cache_name="SQL_PUBLIC_KITTEN_CACHE")
+>>> iterator = dataset.make_one_shot_iterator()
+>>> next_obj = iterator.get_next()
+>>>
+>>> with tf.Session() as sess:
+>>> for _ in range(3):
+>>> print(sess.run(next_obj))
+
+{'key': 1, 'val': {'NAME': b'WARM KITTY'}}
+{'key': 2, 'val': {'NAME': b'SOFT KITTY'}}
+{'key': 3, 'val': {'NAME': b'LITTLE BALL OF FUR'}}
+```
+
+### Structured Objects
+[Apache Ignite](https://ignite.apache.org/) allows to store any type of objects. These objects can have any hierarchy. Ignite Dataset provides an ability to work with such objects.
+
+```python
+>>> import tensorflow as tf
+>>> from tensorflow.contrib.ignite import IgniteDataset
+>>>
+>>> dataset = IgniteDataset(cache_name="IMAGES")
+>>> iterator = dataset.make_one_shot_iterator()
+>>> next_obj = iterator.get_next()
+>>>
+>>> with tf.Session() as sess:
+>>> print(sess.run(next_obj))
+
+{
+ 'key': 'kitten.png',
+ 'val': {
+ 'metadata': {
+ 'file_name': b'kitten.png',
+ 'label': b'little ball of fur',
+ width: 800,
+ height: 600
+ },
+ 'pixels': [0, 0, 0, 0, ..., 0]
+ }
+}
+```
+ Neural network training and other computations require transformations that can be done as part of [tf.data](https://www.tensorflow.org/api_docs/python/tf/data) pipeline if you use Ignite Dataset.
+
+```python
+>>> import tensorflow as tf
+>>> from tensorflow.contrib.ignite import IgniteDataset
+>>>
+>>> dataset = IgniteDataset(cache_name="IMAGES").map(lambda obj: obj['val']['pixels'])
+>>> iterator = dataset.make_one_shot_iterator()
+>>> next_obj = iterator.get_next()
+>>>
+>>> with tf.Session() as sess:
+>>> print(sess.run(next_obj))
+
+[0, 0, 0, 0, ..., 0]
+```
+
+### Distributed Training
+
+TensorFlow is a machine learning framework that [natively supports](https://www.tensorflow.org/deploy/distributed) distributed neural network training, inference and other computations. The main idea behind the distributed neural network training is the ability to calculate gradients of loss functions (squares of the errors) on every partition of data (in terms of horizontal partitioning) and then sum them to get loss function gradient of the whole dataset.
+
+<a href="https://www.codecogs.com/eqnedit.php?latex=\nabla[\sum_1^n(y&space;-&space;\hat{y})^2]&space;=&space;\nabla[\sum_1^{n_1}(y&space;-&space;\hat{y})^2]&space;+&space;\nabla[\sum_{n_1}^{n_2}(y&space;-&space;\hat{y})^2]&space;+&space;...&space;+&space;\nabla[\sum_{n_{k-1}}^n(y&space;-&space;\hat{y})^2]" target="_blank"><img src="https://latex.codecogs.com/gif.latex?\nabla[\sum_1^n(y&space;-&space;\hat{y})^2]&space;=&space;\nabla[\sum_1^{n_1}(y&space;-&space;\hat{y})^2]&space;+&space;\nabla[\sum_{n_1}^{n_2}(y&space;-&space;\hat{y})^2]&space;+&space;...&space;+&space;\nabla[\sum_{n_{k-1}}^n(y&space;-&space;\hat{y})^2]" title="\nabla[\sum_1^n(y - \hat{y})^2] = \nabla[\sum_1^{n_1}(y - \hat{y})^2] + \nabla[\sum_{n_1}^{n_2}(y - \hat{y})^2] + ... + \nabla[\sum_{n_{k-1}}^n(y - \hat{y})^2]" /></a>
+
+Using this ability we can calculate gradients on the nodes the data is stored on, reduce them and then finally update model parameters. It allows to avoid data transfers between nodes and thus to avoid network bottlenecks.
+
+Apache Ignite uses horizontal partitioning to store data in distributed cluster. When we create Apache Ignite cache (or table in terms of SQL), we can specify the number of partitions the data will be partitioned on. For example, if an Apache Ignite cluster consists of 10 machines and we create cache with 10 partitions, then every machine will maintain approximately one data partition.
+
+Ignite Dataset allows using these two aspects of distributed neural network training (using TensorFlow) and Apache Ignite partitioning. Ignite Dataset is a computation graph operation that can be performed on a remote worker. The remote worker can override Ignite Dataset parameters (such as `host`, `port` or `part`) by setting correstondent environment variables for worker process (such as `IGNITE_DATASET_HOST`, `IGNITE_DATASET_PORT` or `IGNITE_DATASET_PART`). Using this overriding approach, we can assign a specific partition to every worker so that one worker handles one partition and, at the same time, transparently work with single dataset.
+
+```python
+>>> import tensorflow as tf
+>>> from tensorflow.contrib.ignite import IgniteDataset
+>>>
+>>> dataset = IgniteDataset("IMAGES")
+>>>
+>>> # Compute gradients locally on every worker node.
+>>> gradients = []
+>>> for i in range(5):
+>>> with tf.device("/job:WORKER/task:%d" % i):
+>>> device_iterator = dataset.make_one_shot_iterator()
+>>> device_next_obj = device_iterator.get_next()
+>>> gradient = compute_gradient(device_next_obj)
+>>> gradients.append(gradient)
+>>>
+>>> # Aggregate them on master node.
+>>> result_gradient = tf.reduce_sum(gradients)
+>>>
+>>> with tf.Session("grpc://localhost:10000") as sess:
+>>> print(sess.run(result_gradient))
+```
+
+High-level TensorFlow API for [distributed training](https://www.tensorflow.org/api_docs/python/tf/contrib/distribute/DistributionStrategy) is supported as well.
+
+### SSL Connection
+
+Apache Ignite allows to protect data transfer channels by [SSL](https://en.wikipedia.org/wiki/Transport_Layer_Security) and authentification. Ignite Dataset supports both SSL connection with and without authntication. For more information, please refer to the [Apache Ignite SSL/TLS](https://apacheignite.readme.io/docs/ssltls) documentation.
+
+```python
+>>> import tensorflow as tf
+>>> from tensorflow.contrib.ignite import IgniteDataset
+>>>
+>>> dataset = IgniteDataset(cache_name="IMAGES", certfile="client.pem", cert_password="password", username="ignite", password="ignite")
+>>> ...
+```
+
+### Windows Support
+
+Ignite Dataset is fully compatible with Windows. You can use it as part of TensorFlow on your Windows workstation as well as on Linux/MacOS systems.
+
+## Try it out
+
+The simplest way to try Ignite Dataset is to run a [Docker](https://www.docker.com/) container with Apache Ignite and loaded [MNIST](http://yann.lecun.com/exdb/mnist/) data and after start interruct with it using Ignite Dataset. Such container is available on Docker Hub: [dmitrievanthony/ignite-with-mnist](https://hub.docker.com/r/dmitrievanthony/ignite-with-mnist/). You need to start this container on your machine:
+
+```
+docker run -it -p 10800:10800 dmitrievanthony/ignite-with-mnist
+```
+
+After that you will be able to work with it following way:
+
+
+
+## Limitations
+
+Presently, Ignite Dataset works with assumption that all objects in the cache have the same structure (homogeneous objects) and the cache contains at least one object. Another limitation concerns structured objects, Ignite Dataset does not support UUID, Maps and Object arrays that might be parts of an object structure.
diff --git a/tensorflow/contrib/ignite/__init__.py b/tensorflow/contrib/ignite/__init__.py
new file mode 100644
index 0000000..f429476
--- /dev/null
+++ b/tensorflow/contrib/ignite/__init__.py
@@ -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.
+# ==============================================================================
+"""IgniteDataset that allows to get data from Apache Ignite.
+
+Apache Ignite is a memory-centric distributed database, 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:
+https://apacheignite.readme.io/v2.6/docs/binary-client-protocol.
+
+@@IgniteDataset
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.contrib.ignite.python.ops.ignite_dataset_ops import IgniteDataset
+from tensorflow.python.util.all_util import remove_undocumented
+
+_allowed_symbols = [
+ "IgniteDataset",
+]
+
+remove_undocumented(__name__)
diff --git a/tensorflow/contrib/ignite/kernels/ignite_binary_object_parser.cc b/tensorflow/contrib/ignite/kernels/ignite_binary_object_parser.cc
new file mode 100644
index 0000000..2c8a7d4
--- /dev/null
+++ b/tensorflow/contrib/ignite/kernels/ignite_binary_object_parser.cc
@@ -0,0 +1,334 @@
+/* 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/contrib/ignite/kernels/ignite_binary_object_parser.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/lib/core/errors.h"
+
+namespace tensorflow {
+
+BinaryObjectParser::BinaryObjectParser() : byte_swapper_(ByteSwapper(false)) {}
+
+Status BinaryObjectParser::Parse(uint8_t** ptr,
+ std::vector<Tensor>* out_tensors,
+ std::vector<int32_t>* types) const {
+ uint8_t object_type_id = ParseByte(ptr);
+
+ // Skip non-leaf nodes.
+ if (object_type_id != WRAPPED_OBJ && object_type_id != COMPLEX_OBJ)
+ types->push_back(object_type_id);
+
+ switch (object_type_id) {
+ case BYTE: {
+ out_tensors->emplace_back(cpu_allocator(), DT_UINT8, TensorShape({}));
+ out_tensors->back().scalar<uint8>()() = ParseByte(ptr);
+ break;
+ }
+ case SHORT: {
+ out_tensors->emplace_back(cpu_allocator(), DT_INT16, TensorShape({}));
+ out_tensors->back().scalar<int16>()() = ParseShort(ptr);
+ break;
+ }
+ case USHORT: {
+ out_tensors->emplace_back(cpu_allocator(), DT_UINT16, TensorShape({}));
+ out_tensors->back().scalar<uint16>()() = ParseUnsignedShort(ptr);
+ break;
+ }
+ case INT: {
+ out_tensors->emplace_back(cpu_allocator(), DT_INT32, TensorShape({}));
+ out_tensors->back().scalar<int32>()() = ParseInt(ptr);
+ break;
+ }
+ case LONG: {
+ out_tensors->emplace_back(cpu_allocator(), DT_INT64, TensorShape({}));
+ out_tensors->back().scalar<int64>()() = ParseLong(ptr);
+ break;
+ }
+ case FLOAT: {
+ out_tensors->emplace_back(cpu_allocator(), DT_FLOAT, TensorShape({}));
+ out_tensors->back().scalar<float>()() = ParseFloat(ptr);
+ break;
+ }
+ case DOUBLE: {
+ out_tensors->emplace_back(cpu_allocator(), DT_DOUBLE, TensorShape({}));
+ out_tensors->back().scalar<double>()() = ParseDouble(ptr);
+ break;
+ }
+ case BOOL: {
+ out_tensors->emplace_back(cpu_allocator(), DT_BOOL, TensorShape({}));
+ out_tensors->back().scalar<bool>()() = ParseBool(ptr);
+ break;
+ }
+ case STRING: {
+ out_tensors->emplace_back(cpu_allocator(), DT_STRING, TensorShape({}));
+ out_tensors->back().scalar<string>()() = ParseString(ptr);
+ break;
+ }
+ case DATE: {
+ out_tensors->emplace_back(cpu_allocator(), DT_INT64, TensorShape({}));
+ out_tensors->back().scalar<int64>()() = ParseLong(ptr);
+ break;
+ }
+ case BYTE_ARR: {
+ int32_t length = ParseInt(ptr);
+ uint8_t* arr = ParseByteArr(ptr, length);
+ out_tensors->emplace_back(cpu_allocator(), DT_UINT8,
+ TensorShape({length}));
+ std::copy_n(arr, length, out_tensors->back().flat<uint8>().data());
+ break;
+ }
+ case SHORT_ARR: {
+ int32_t length = ParseInt(ptr);
+ int16_t* arr = ParseShortArr(ptr, length);
+ out_tensors->emplace_back(cpu_allocator(), DT_INT16,
+ TensorShape({length}));
+ std::copy_n(arr, length, out_tensors->back().flat<int16>().data());
+ break;
+ }
+ case USHORT_ARR: {
+ int32_t length = ParseInt(ptr);
+ uint16_t* arr = ParseUnsignedShortArr(ptr, length);
+ out_tensors->emplace_back(cpu_allocator(), DT_UINT16,
+ TensorShape({length}));
+ std::copy_n(arr, length, out_tensors->back().flat<uint16>().data());
+ break;
+ }
+ case INT_ARR: {
+ int32_t length = ParseInt(ptr);
+ int32_t* arr = ParseIntArr(ptr, length);
+ out_tensors->emplace_back(cpu_allocator(), DT_INT32,
+ TensorShape({length}));
+ std::copy_n(arr, length, out_tensors->back().flat<int32>().data());
+ break;
+ }
+ case LONG_ARR: {
+ int32_t length = ParseInt(ptr);
+ int64_t* arr = ParseLongArr(ptr, length);
+ out_tensors->emplace_back(cpu_allocator(), DT_INT64,
+ TensorShape({length}));
+ std::copy_n(arr, length, out_tensors->back().flat<int64>().data());
+ break;
+ }
+ case FLOAT_ARR: {
+ int32_t length = ParseInt(ptr);
+ float* arr = ParseFloatArr(ptr, length);
+ out_tensors->emplace_back(cpu_allocator(), DT_FLOAT,
+ TensorShape({length}));
+ std::copy_n(arr, length, out_tensors->back().flat<float>().data());
+ break;
+ }
+ case DOUBLE_ARR: {
+ int32_t length = ParseInt(ptr);
+ double* arr = ParseDoubleArr(ptr, length);
+ out_tensors->emplace_back(cpu_allocator(), DT_DOUBLE,
+ TensorShape({length}));
+ std::copy_n(arr, length, out_tensors->back().flat<double>().data());
+ break;
+ }
+ case BOOL_ARR: {
+ int32_t length = ParseInt(ptr);
+ bool* arr = ParseBoolArr(ptr, length);
+ out_tensors->emplace_back(cpu_allocator(), DT_BOOL,
+ TensorShape({length}));
+ std::copy_n(arr, length, out_tensors->back().flat<bool>().data());
+ break;
+ }
+ case STRING_ARR: {
+ int32_t length = ParseInt(ptr);
+ out_tensors->emplace_back(cpu_allocator(), DT_STRING,
+ TensorShape({length}));
+ for (int32_t i = 0; i < length; i++)
+ out_tensors->back().vec<string>()(i) = ParseString(ptr);
+ break;
+ }
+ case DATE_ARR: {
+ int32_t length = ParseInt(ptr);
+ int64_t* arr = ParseLongArr(ptr, length);
+ out_tensors->emplace_back(cpu_allocator(), DT_INT64,
+ TensorShape({length}));
+ std::copy_n(arr, length, out_tensors->back().flat<int64>().data());
+ break;
+ }
+ case WRAPPED_OBJ: {
+ int32_t byte_arr_size = ParseInt(ptr);
+ TF_RETURN_IF_ERROR(Parse(ptr, out_tensors, types));
+ int32_t offset = ParseInt(ptr);
+
+ break;
+ }
+ case COMPLEX_OBJ: {
+ uint8_t version = ParseByte(ptr);
+ int16_t flags = ParseShort(ptr);
+ int32_t type_id = ParseInt(ptr);
+ int32_t hash_code = ParseInt(ptr);
+ int32_t length = ParseInt(ptr);
+ int32_t schema_id = ParseInt(ptr);
+ int32_t schema_offset = ParseInt(ptr);
+
+ // 24 is size of header just read.
+ uint8_t* end = *ptr + schema_offset - 24;
+ int32_t i = 0;
+ while (*ptr < end) {
+ i++;
+ TF_RETURN_IF_ERROR(Parse(ptr, out_tensors, types));
+ }
+
+ *ptr += (length - schema_offset);
+
+ break;
+ }
+ default: {
+ return errors::Unknown("Unknowd binary type (type id ",
+ (int)object_type_id, ")");
+ }
+ }
+
+ return Status::OK();
+}
+
+uint8_t BinaryObjectParser::ParseByte(uint8_t** ptr) const {
+ uint8_t res = **ptr;
+ *ptr += 1;
+
+ return res;
+}
+
+int16_t BinaryObjectParser::ParseShort(uint8_t** ptr) const {
+ int16_t* res = *reinterpret_cast<int16_t**>(ptr);
+ byte_swapper_.SwapIfRequiredInt16(res);
+ *ptr += 2;
+
+ return *res;
+}
+
+uint16_t BinaryObjectParser::ParseUnsignedShort(uint8_t** ptr) const {
+ uint16_t* res = *reinterpret_cast<uint16_t**>(ptr);
+ byte_swapper_.SwapIfRequiredUnsignedInt16(res);
+ *ptr += 2;
+
+ return *res;
+}
+
+int32_t BinaryObjectParser::ParseInt(uint8_t** ptr) const {
+ int32_t* res = *reinterpret_cast<int32_t**>(ptr);
+ byte_swapper_.SwapIfRequiredInt32(res);
+ *ptr += 4;
+
+ return *res;
+}
+
+int64_t BinaryObjectParser::ParseLong(uint8_t** ptr) const {
+ int64_t* res = *reinterpret_cast<int64_t**>(ptr);
+ byte_swapper_.SwapIfRequiredInt64(res);
+ *ptr += 8;
+
+ return *res;
+}
+
+float BinaryObjectParser::ParseFloat(uint8_t** ptr) const {
+ float* res = *reinterpret_cast<float**>(ptr);
+ byte_swapper_.SwapIfRequiredFloat(res);
+ *ptr += 4;
+
+ return *res;
+}
+
+double BinaryObjectParser::ParseDouble(uint8_t** ptr) const {
+ double* res = *reinterpret_cast<double**>(ptr);
+ byte_swapper_.SwapIfRequiredDouble(res);
+ *ptr += 8;
+
+ return *res;
+}
+
+bool BinaryObjectParser::ParseBool(uint8_t** ptr) const {
+ bool res = **reinterpret_cast<bool**>(ptr);
+ *ptr += 1;
+
+ return res;
+}
+
+string BinaryObjectParser::ParseString(uint8_t** ptr) const {
+ int32_t length = ParseInt(ptr);
+ string res(*reinterpret_cast<char**>(ptr), length);
+ *ptr += length;
+
+ return res;
+}
+
+uint8_t* BinaryObjectParser::ParseByteArr(uint8_t** ptr, int length) const {
+ uint8_t* res = *reinterpret_cast<uint8_t**>(ptr);
+ *ptr += length;
+
+ return res;
+}
+
+int16_t* BinaryObjectParser::ParseShortArr(uint8_t** ptr, int length) const {
+ int16_t* res = *reinterpret_cast<int16_t**>(ptr);
+ byte_swapper_.SwapIfRequiredInt16Arr(res, length);
+ *ptr += length * 2;
+
+ return res;
+}
+
+uint16_t* BinaryObjectParser::ParseUnsignedShortArr(uint8_t** ptr,
+ int length) const {
+ uint16_t* res = *reinterpret_cast<uint16_t**>(ptr);
+ byte_swapper_.SwapIfRequiredUnsignedInt16Arr(res, length);
+ *ptr += length * 2;
+
+ return res;
+}
+
+int32_t* BinaryObjectParser::ParseIntArr(uint8_t** ptr, int length) const {
+ int32_t* res = *reinterpret_cast<int32_t**>(ptr);
+ byte_swapper_.SwapIfRequiredInt32Arr(res, length);
+ *ptr += length * 4;
+
+ return res;
+}
+
+int64_t* BinaryObjectParser::ParseLongArr(uint8_t** ptr, int length) const {
+ int64_t* res = *reinterpret_cast<int64_t**>(ptr);
+ byte_swapper_.SwapIfRequiredInt64Arr(res, length);
+ *ptr += length * 8;
+
+ return res;
+}
+
+float* BinaryObjectParser::ParseFloatArr(uint8_t** ptr, int length) const {
+ float* res = *reinterpret_cast<float**>(ptr);
+ byte_swapper_.SwapIfRequiredFloatArr(res, length);
+ *ptr += length * 4;
+
+ return res;
+}
+
+double* BinaryObjectParser::ParseDoubleArr(uint8_t** ptr, int length) const {
+ double* res = *reinterpret_cast<double**>(ptr);
+ byte_swapper_.SwapIfRequiredDoubleArr(res, length);
+ *ptr += length * 8;
+
+ return res;
+}
+
+bool* BinaryObjectParser::ParseBoolArr(uint8_t** ptr, int length) const {
+ bool* res = *reinterpret_cast<bool**>(ptr);
+ *ptr += length;
+
+ return res;
+}
+
+} // namespace tensorflow
diff --git a/tensorflow/contrib/ignite/kernels/ignite_binary_object_parser.h b/tensorflow/contrib/ignite/kernels/ignite_binary_object_parser.h
new file mode 100644
index 0000000..eb1f856
--- /dev/null
+++ b/tensorflow/contrib/ignite/kernels/ignite_binary_object_parser.h
@@ -0,0 +1,81 @@
+/* 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_CONTRIB_IGNITE_KERNELS_IGNITE_BINARY_OBJECT_PARSER_H_
+#define TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_BINARY_OBJECT_PARSER_H_
+
+#include <vector>
+#include "tensorflow/contrib/ignite/kernels/ignite_byte_swapper.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/lib/core/status.h"
+
+namespace tensorflow {
+
+class BinaryObjectParser {
+ public:
+ BinaryObjectParser();
+ Status Parse(uint8_t** ptr, std::vector<Tensor>* out_tensors,
+ std::vector<int32_t>* types) const;
+
+ private:
+ uint8_t ParseByte(uint8_t** ptr) const;
+ int16_t ParseShort(uint8_t** ptr) const;
+ uint16_t ParseUnsignedShort(uint8_t** ptr) const;
+ int32_t ParseInt(uint8_t** ptr) const;
+ int64_t ParseLong(uint8_t** ptr) const;
+ float ParseFloat(uint8_t** ptr) const;
+ double ParseDouble(uint8_t** ptr) const;
+ bool ParseBool(uint8_t** ptr) const;
+ string ParseString(uint8_t** ptr) const;
+ uint8_t* ParseByteArr(uint8_t** ptr, int length) const;
+ int16_t* ParseShortArr(uint8_t** ptr, int length) const;
+ uint16_t* ParseUnsignedShortArr(uint8_t** ptr, int length) const;
+ int32_t* ParseIntArr(uint8_t** ptr, int length) const;
+ int64_t* ParseLongArr(uint8_t** ptr, int length) const;
+ float* ParseFloatArr(uint8_t** ptr, int length) const;
+ double* ParseDoubleArr(uint8_t** ptr, int length) const;
+ bool* ParseBoolArr(uint8_t** ptr, int length) const;
+
+ const ByteSwapper byte_swapper_;
+};
+
+enum ObjectType {
+ BYTE = 1,
+ SHORT = 2,
+ INT = 3,
+ LONG = 4,
+ FLOAT = 5,
+ DOUBLE = 6,
+ USHORT = 7,
+ BOOL = 8,
+ STRING = 9,
+ DATE = 11,
+ BYTE_ARR = 12,
+ SHORT_ARR = 13,
+ INT_ARR = 14,
+ LONG_ARR = 15,
+ FLOAT_ARR = 16,
+ DOUBLE_ARR = 17,
+ USHORT_ARR = 18,
+ BOOL_ARR = 19,
+ STRING_ARR = 20,
+ DATE_ARR = 22,
+ WRAPPED_OBJ = 27,
+ COMPLEX_OBJ = 103
+};
+
+} // namespace tensorflow
+
+#endif // TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_BINARY_OBJECT_PARSER_H_
diff --git a/tensorflow/contrib/ignite/kernels/ignite_byte_swapper.h b/tensorflow/contrib/ignite/kernels/ignite_byte_swapper.h
new file mode 100644
index 0000000..46df3e3
--- /dev/null
+++ b/tensorflow/contrib/ignite/kernels/ignite_byte_swapper.h
@@ -0,0 +1,126 @@
+/* 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_CONTRIB_IGNITE_KERNELS_IGNITE_BYTE_SWAPPER_H_
+#define TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_BYTE_SWAPPER_H_
+
+#include <stdint.h>
+#include "tensorflow/core/platform/byte_order.h"
+
+namespace tensorflow {
+
+class ByteSwapper {
+ public:
+ ByteSwapper(bool big_endian) { swap_ = big_endian == port::kLittleEndian; }
+
+ inline void SwapIfRequiredInt16(int16_t *x) const {
+ if (swap_) {
+ Swap16(x);
+ }
+ }
+
+ inline void SwapIfRequiredUnsignedInt16(uint16_t *x) const {
+ if (swap_) {
+ Swap16(reinterpret_cast<int16_t *>(x));
+ }
+ }
+
+ inline void SwapIfRequiredInt32(int32_t *x) const {
+ if (swap_) {
+ Swap32(x);
+ }
+ }
+
+ inline void SwapIfRequiredFloat(float *x) const {
+ if (swap_) {
+ Swap32(reinterpret_cast<int32_t *>(x));
+ }
+ }
+
+ inline void SwapIfRequiredInt64(int64_t *x) const {
+ if (swap_) {
+ Swap64(x);
+ }
+ }
+
+ inline void SwapIfRequiredDouble(double *x) const {
+ if (swap_) {
+ Swap64(reinterpret_cast<int64_t *>(x));
+ }
+ }
+
+ inline void SwapIfRequiredInt16Arr(int16_t *x, int32_t length) const {
+ if (swap_) {
+ for (int32_t i = 0; i < length; i++) Swap16(&x[i]);
+ }
+ }
+
+ inline void SwapIfRequiredUnsignedInt16Arr(uint16_t *x,
+ int32_t length) const {
+ if (swap_) {
+ for (int32_t i = 0; i < length; i++)
+ Swap16(reinterpret_cast<int16_t *>(&x[i]));
+ }
+ }
+
+ inline void SwapIfRequiredInt32Arr(int32_t *x, int32_t length) const {
+ if (swap_) {
+ for (int32_t i = 0; i < length; i++) Swap32(&x[i]);
+ }
+ }
+
+ inline void SwapIfRequiredFloatArr(float *x, int32_t length) const {
+ if (swap_) {
+ for (int32_t i = 0; i < length; i++)
+ Swap32(reinterpret_cast<int32_t *>(&x[i]));
+ }
+ }
+
+ inline void SwapIfRequiredInt64Arr(int64_t *x, int32_t length) const {
+ if (swap_) {
+ for (int32_t i = 0; i < length; i++) Swap64(&x[i]);
+ }
+ }
+
+ inline void SwapIfRequiredDoubleArr(double *x, int32_t length) const {
+ if (swap_) {
+ for (int32_t i = 0; i < length; i++)
+ Swap64(reinterpret_cast<int64_t *>(&x[i]));
+ }
+ }
+
+ private:
+ inline void Swap16(int16_t *x) const {
+ *x = ((*x & 0xFF) << 8) | ((*x >> 8) & 0xFF);
+ }
+
+ inline void Swap32(int32_t *x) const {
+ *x = ((*x & 0xFF) << 24) | (((*x >> 8) & 0xFF) << 16) |
+ (((*x >> 16) & 0xFF) << 8) | ((*x >> 24) & 0xFF);
+ }
+
+ inline void Swap64(int64_t *x) const {
+ *x = ((*x & 0xFF) << 56) | (((*x >> 8) & 0xFF) << 48) |
+ (((*x >> 16) & 0xFF) << 40) | (((*x >> 24) & 0xFF) << 32) |
+ (((*x >> 32) & 0xFF) << 24) | (((*x >> 40) & 0xFF) << 16) |
+ (((*x >> 48) & 0xFF) << 8) | ((*x >> 56) & 0xFF);
+ }
+
+ bool swap_;
+};
+
+} // namespace tensorflow
+
+#endif // TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_BYTE_SWAPPER_H_
diff --git a/tensorflow/contrib/ignite/kernels/ignite_client.h b/tensorflow/contrib/ignite/kernels/ignite_client.h
new file mode 100644
index 0000000..459b50b
--- /dev/null
+++ b/tensorflow/contrib/ignite/kernels/ignite_client.h
@@ -0,0 +1,84 @@
+/* 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_CONTRIB_IGNITE_KERNELS_IGNITE_CLIENT_H_
+#define TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_CLIENT_H_
+
+#include "tensorflow/contrib/ignite/kernels/ignite_byte_swapper.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/status.h"
+
+namespace tensorflow {
+
+class Client {
+ public:
+ Client(bool big_endian) : byte_swapper_(ByteSwapper(big_endian)) {}
+ virtual Status Connect() = 0;
+ virtual Status Disconnect() = 0;
+ virtual bool IsConnected() = 0;
+ virtual int GetSocketDescriptor() = 0;
+ virtual Status ReadData(uint8_t *buf, const int32_t length) = 0;
+ virtual Status WriteData(const uint8_t *buf, const int32_t length) = 0;
+
+ inline Status ReadByte(uint8_t *data) { return ReadData(data, 1); }
+
+ inline Status ReadShort(int16_t *data) {
+ TF_RETURN_IF_ERROR(ReadData((uint8_t *)data, 2));
+ byte_swapper_.SwapIfRequiredInt16(data);
+
+ return Status::OK();
+ }
+
+ inline Status ReadInt(int32_t *data) {
+ TF_RETURN_IF_ERROR(ReadData((uint8_t *)data, 4));
+ byte_swapper_.SwapIfRequiredInt32(data);
+
+ return Status::OK();
+ }
+
+ inline Status ReadLong(int64_t *data) {
+ TF_RETURN_IF_ERROR(ReadData((uint8_t *)data, 8));
+ byte_swapper_.SwapIfRequiredInt64(data);
+
+ return Status::OK();
+ }
+
+ inline Status WriteByte(const uint8_t data) { return WriteData(&data, 1); }
+
+ inline Status WriteShort(const int16_t data) {
+ int16_t tmp = data;
+ byte_swapper_.SwapIfRequiredInt16(&tmp);
+ return WriteData((uint8_t *)&tmp, 2);
+ }
+
+ inline Status WriteInt(const int32_t data) {
+ int32_t tmp = data;
+ byte_swapper_.SwapIfRequiredInt32(&tmp);
+ return WriteData((uint8_t *)&tmp, 4);
+ }
+
+ inline Status WriteLong(const int64_t data) {
+ int64_t tmp = data;
+ byte_swapper_.SwapIfRequiredInt64(&tmp);
+ return WriteData((uint8_t *)&tmp, 8);
+ }
+
+ private:
+ const ByteSwapper byte_swapper_;
+};
+
+} // namespace tensorflow
+
+#endif // TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_CLIENT_H_
diff --git a/tensorflow/contrib/ignite/kernels/ignite_dataset.cc b/tensorflow/contrib/ignite/kernels/ignite_dataset.cc
new file mode 100644
index 0000000..c4a7d3c
--- /dev/null
+++ b/tensorflow/contrib/ignite/kernels/ignite_dataset.cc
@@ -0,0 +1,81 @@
+/* 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/contrib/ignite/kernels/ignite_dataset_iterator.h"
+#include "tensorflow/core/platform/logging.h"
+
+namespace tensorflow {
+
+IgniteDataset::IgniteDataset(OpKernelContext* ctx, string cache_name,
+ string host, int32 port, bool local, int32 part,
+ int32 page_size, string username, string password,
+ string certfile, string keyfile,
+ string cert_password, std::vector<int32> schema,
+ std::vector<int32> permutation,
+ DataTypeVector dtypes,
+ std::vector<PartialTensorShape> shapes)
+ : DatasetBase(DatasetContext(ctx)),
+ cache_name_(std::move(cache_name)),
+ host_(std::move(host)),
+ port_(port),
+ local_(local),
+ part_(part),
+ page_size_(page_size),
+ username_(std::move(username)),
+ password_(std::move(password)),
+ certfile_(std::move(certfile)),
+ keyfile_(std::move(keyfile)),
+ cert_password_(std::move(cert_password)),
+ schema_(std::move(schema)),
+ permutation_(std::move(permutation)),
+ dtypes_(dtypes),
+ shapes_(shapes) {
+ LOG(INFO) << "Ignite Dataset created [cache_name='" << cache_name_
+ << "', host='" << host_ << "', port=" << port_
+ << ", local=" << local_ << ", part=" << part_
+ << ", page_size=" << page_size_ << ", username='" << username_
+ << "', certfile='" << certfile_ << "', keyfile='"
+ << keyfile_ + "']";
+}
+
+IgniteDataset::~IgniteDataset() { LOG(INFO) << "Ignite Dataset destroyed"; }
+
+std::unique_ptr<IteratorBase> IgniteDataset::MakeIteratorInternal(
+ const string& prefix) const {
+ return std::unique_ptr<IteratorBase>(new IgniteDatasetIterator(
+ {this, strings::StrCat(prefix, "::Ignite")}, std::move(this->host_),
+ this->port_, std::move(this->cache_name_), this->local_, this->part_,
+ this->page_size_, std::move(this->username_), std::move(this->password_),
+ std::move(this->certfile_), std::move(this->keyfile_),
+ std::move(this->cert_password_), std::move(this->schema_),
+ std::move(this->permutation_)));
+}
+
+const DataTypeVector& IgniteDataset::output_dtypes() const { return dtypes_; }
+
+const std::vector<PartialTensorShape>& IgniteDataset::output_shapes() const {
+ return shapes_;
+}
+
+string IgniteDataset::DebugString() const { return "IgniteDatasetOp::Dataset"; }
+
+Status IgniteDataset::AsGraphDefInternal(SerializationContext* ctx,
+ DatasetGraphDefBuilder* b,
+ Node** output) const {
+ return errors::Unimplemented(
+ "IgniteDataset does not support 'AsGraphDefInternal'");
+}
+
+} // namespace tensorflow
diff --git a/tensorflow/contrib/ignite/kernels/ignite_dataset.h b/tensorflow/contrib/ignite/kernels/ignite_dataset.h
new file mode 100644
index 0000000..66bfdf2
--- /dev/null
+++ b/tensorflow/contrib/ignite/kernels/ignite_dataset.h
@@ -0,0 +1,63 @@
+/* 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_CONTRIB_IGNITE_KERNELS_IGNITE_DATASET_H_
+#define TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_DATASET_H_
+
+#include "tensorflow/core/framework/dataset.h"
+
+namespace tensorflow {
+
+class IgniteDataset : public DatasetBase {
+ public:
+ IgniteDataset(OpKernelContext* ctx, string cache_name, string host,
+ int32 port, bool local, int32 part, int32 page_size,
+ string username, string password, string certfile,
+ string keyfile, string cert_password, std::vector<int32> schema,
+ std::vector<int32> permutation, DataTypeVector dtypes,
+ std::vector<PartialTensorShape> shapes);
+ ~IgniteDataset();
+ std::unique_ptr<IteratorBase> MakeIteratorInternal(
+ const string& prefix) const override;
+ const DataTypeVector& output_dtypes() const override;
+ const std::vector<PartialTensorShape>& output_shapes() const override;
+ string DebugString() const override;
+
+ protected:
+ Status AsGraphDefInternal(SerializationContext* ctx,
+ DatasetGraphDefBuilder* b,
+ Node** output) const override;
+
+ private:
+ const string cache_name_;
+ const string host_;
+ const int32 port_;
+ const bool local_;
+ const int32 part_;
+ const int32 page_size_;
+ const string username_;
+ const string password_;
+ const string certfile_;
+ const string keyfile_;
+ const string cert_password_;
+ const std::vector<int32> schema_;
+ const std::vector<int32> permutation_;
+ const DataTypeVector dtypes_;
+ const std::vector<PartialTensorShape> shapes_;
+};
+
+} // namespace tensorflow
+
+#endif // TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_DATASET_H_
diff --git a/tensorflow/contrib/ignite/kernels/ignite_dataset_iterator.cc b/tensorflow/contrib/ignite/kernels/ignite_dataset_iterator.cc
new file mode 100644
index 0000000..5da9127
--- /dev/null
+++ b/tensorflow/contrib/ignite/kernels/ignite_dataset_iterator.cc
@@ -0,0 +1,422 @@
+/* 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/contrib/ignite/kernels/ignite_dataset_iterator.h"
+
+#include "tensorflow/contrib/ignite/kernels/ignite_plain_client.h"
+#include "tensorflow/contrib/ignite/kernels/ignite_ssl_wrapper.h"
+#include "tensorflow/core/lib/gtl/cleanup.h"
+#include "tensorflow/core/platform/logging.h"
+
+namespace tensorflow {
+
+IgniteDatasetIterator::IgniteDatasetIterator(
+ const Params& params, string host, int32 port, string cache_name,
+ bool local, int32 part, int32 page_size, string username, string password,
+ string certfile, string keyfile, string cert_password,
+ std::vector<int32> schema, std::vector<int32> permutation)
+ : DatasetIterator<IgniteDataset>(params),
+ cache_name_(std::move(cache_name)),
+ local_(local),
+ part_(part),
+ page_size_(page_size),
+ username_(std::move(username)),
+ password_(std::move(password)),
+ schema_(std::move(schema)),
+ permutation_(std::move(permutation)),
+ remainder_(-1),
+ cursor_id_(-1),
+ last_page_(false),
+ valid_state_(true) {
+ Client* p_client = new PlainClient(std::move(host), port, false);
+
+ if (certfile.empty())
+ client_ = std::unique_ptr<Client>(p_client);
+ else
+ client_ = std::unique_ptr<Client>(
+ new SslWrapper(std::unique_ptr<Client>(p_client), std::move(certfile),
+ std::move(keyfile), std::move(cert_password), false));
+
+ LOG(INFO) << "Ignite Dataset Iterator created";
+}
+
+IgniteDatasetIterator::~IgniteDatasetIterator() {
+ Status status = CloseConnection();
+ if (!status.ok()) LOG(ERROR) << status.ToString();
+
+ LOG(INFO) << "Ignite Dataset Iterator destroyed";
+}
+
+Status IgniteDatasetIterator::GetNextInternal(IteratorContext* ctx,
+ std::vector<Tensor>* out_tensors,
+ bool* end_of_sequence) {
+ mutex_lock l(mutex_);
+
+ if (valid_state_) {
+ Status status =
+ GetNextInternalWithValidState(ctx, out_tensors, end_of_sequence);
+
+ if (!status.ok()) valid_state_ = false;
+
+ return status;
+ }
+
+ return errors::Unknown("Iterator is invalid");
+}
+
+Status IgniteDatasetIterator::SaveInternal(IteratorStateWriter* writer) {
+ return errors::Unimplemented(
+ "Iterator for IgniteDataset does not support 'SaveInternal'");
+}
+
+Status IgniteDatasetIterator::RestoreInternal(IteratorContext* ctx,
+ IteratorStateReader* reader) {
+ return errors::Unimplemented(
+ "Iterator for IgniteDataset does not support 'RestoreInternal')");
+}
+
+Status IgniteDatasetIterator::GetNextInternalWithValidState(
+ IteratorContext* ctx, std::vector<Tensor>* out_tensors,
+ bool* end_of_sequence) {
+ if (remainder_ == 0 && last_page_) {
+ cursor_id_ = -1;
+ *end_of_sequence = true;
+
+ return Status::OK();
+ } else {
+ TF_RETURN_IF_ERROR(EstablishConnection());
+
+ if (remainder_ == -1) {
+ TF_RETURN_IF_ERROR(ScanQuery());
+ } else if (remainder_ == 0) {
+ TF_RETURN_IF_ERROR(LoadNextPage());
+ }
+
+ uint8_t* initial_ptr = ptr_;
+ std::vector<Tensor> tensors;
+ std::vector<int32_t> types;
+
+ TF_RETURN_IF_ERROR(parser_.Parse(&ptr_, &tensors, &types)); // Parse key
+ TF_RETURN_IF_ERROR(parser_.Parse(&ptr_, &tensors, &types)); // Parse val
+
+ remainder_ -= (ptr_ - initial_ptr);
+
+ TF_RETURN_IF_ERROR(CheckTypes(types));
+
+ for (size_t i = 0; i < tensors.size(); i++)
+ out_tensors->push_back(tensors[permutation_[i]]);
+
+ *end_of_sequence = false;
+
+ return Status::OK();
+ }
+
+ *end_of_sequence = true;
+
+ return Status::OK();
+}
+
+Status IgniteDatasetIterator::EstablishConnection() {
+ if (!client_->IsConnected()) {
+ TF_RETURN_IF_ERROR(client_->Connect());
+
+ Status status = Handshake();
+ if (!status.ok()) {
+ Status disconnect_status = client_->Disconnect();
+ if (!disconnect_status.ok()) LOG(ERROR) << disconnect_status.ToString();
+
+ return status;
+ }
+ }
+
+ return Status::OK();
+}
+
+Status IgniteDatasetIterator::CloseConnection() {
+ if (cursor_id_ != -1 && !last_page_) {
+ TF_RETURN_IF_ERROR(EstablishConnection());
+
+ TF_RETURN_IF_ERROR(client_->WriteInt(kCloseConnectionReqLength));
+ TF_RETURN_IF_ERROR(client_->WriteShort(kCloseConnectionOpcode));
+ TF_RETURN_IF_ERROR(client_->WriteLong(0)); // Request ID
+ TF_RETURN_IF_ERROR(client_->WriteLong(cursor_id_)); // Resource ID
+
+ int32_t res_len;
+ TF_RETURN_IF_ERROR(client_->ReadInt(&res_len));
+ if (res_len < kMinResLength)
+ return errors::Unknown("Close Resource Response is corrupted");
+
+ int64_t req_id;
+ TF_RETURN_IF_ERROR(client_->ReadLong(&req_id));
+ int32_t status;
+ TF_RETURN_IF_ERROR(client_->ReadInt(&status));
+ if (status != 0) {
+ uint8_t err_msg_header;
+ TF_RETURN_IF_ERROR(client_->ReadByte(&err_msg_header));
+ if (err_msg_header == kStringVal) {
+ int32_t err_msg_length;
+ TF_RETURN_IF_ERROR(client_->ReadInt(&err_msg_length));
+
+ uint8_t* err_msg_c = new uint8_t[err_msg_length];
+ auto clean = gtl::MakeCleanup([err_msg_c] { delete[] err_msg_c; });
+ TF_RETURN_IF_ERROR(client_->ReadData(err_msg_c, err_msg_length));
+ string err_msg(reinterpret_cast<char*>(err_msg_c), err_msg_length);
+
+ return errors::Unknown("Close Resource Error [status=", status,
+ ", message=", err_msg, "]");
+ }
+ return errors::Unknown("Close Resource Error [status=", status, "]");
+ }
+
+ cursor_id_ = -1;
+
+ return client_->Disconnect();
+ } else {
+ LOG(INFO) << "Query Cursor " << cursor_id_ << " is already closed";
+ }
+
+ return client_->IsConnected() ? client_->Disconnect() : Status::OK();
+}
+
+Status IgniteDatasetIterator::Handshake() {
+ int32_t msg_len = kHandshakeReqDefaultLength;
+
+ if (username_.empty())
+ msg_len += 1;
+ else
+ msg_len += 5 + username_.length(); // 1 byte header, 4 bytes length.
+
+ if (password_.empty())
+ msg_len += 1;
+ else
+ msg_len += 5 + password_.length(); // 1 byte header, 4 bytes length.
+
+ TF_RETURN_IF_ERROR(client_->WriteInt(msg_len));
+ TF_RETURN_IF_ERROR(client_->WriteByte(1));
+ TF_RETURN_IF_ERROR(client_->WriteShort(kProtocolMajorVersion));
+ TF_RETURN_IF_ERROR(client_->WriteShort(kProtocolMinorVersion));
+ TF_RETURN_IF_ERROR(client_->WriteShort(kProtocolPatchVersion));
+ TF_RETURN_IF_ERROR(client_->WriteByte(2));
+ if (username_.empty()) {
+ TF_RETURN_IF_ERROR(client_->WriteByte(kNullVal));
+ } else {
+ TF_RETURN_IF_ERROR(client_->WriteByte(kStringVal));
+ TF_RETURN_IF_ERROR(client_->WriteInt(username_.length()));
+ TF_RETURN_IF_ERROR(
+ client_->WriteData(reinterpret_cast<const uint8_t*>(username_.c_str()),
+ username_.length()));
+ }
+
+ if (password_.empty()) {
+ TF_RETURN_IF_ERROR(client_->WriteByte(kNullVal));
+ } else {
+ TF_RETURN_IF_ERROR(client_->WriteByte(kStringVal));
+ TF_RETURN_IF_ERROR(client_->WriteInt(password_.length()));
+ TF_RETURN_IF_ERROR(
+ client_->WriteData(reinterpret_cast<const uint8_t*>(password_.c_str()),
+ password_.length()));
+ }
+
+ int32_t handshake_res_len;
+ TF_RETURN_IF_ERROR(client_->ReadInt(&handshake_res_len));
+ uint8_t handshake_res;
+ TF_RETURN_IF_ERROR(client_->ReadByte(&handshake_res));
+
+ if (handshake_res != 1) {
+ int16_t serv_ver_major;
+ TF_RETURN_IF_ERROR(client_->ReadShort(&serv_ver_major));
+ int16_t serv_ver_minor;
+ TF_RETURN_IF_ERROR(client_->ReadShort(&serv_ver_minor));
+ int16_t serv_ver_patch;
+ TF_RETURN_IF_ERROR(client_->ReadShort(&serv_ver_patch));
+ uint8_t header;
+ TF_RETURN_IF_ERROR(client_->ReadByte(&header));
+
+ if (header == kStringVal) {
+ int32_t length;
+ TF_RETURN_IF_ERROR(client_->ReadInt(&length));
+
+ uint8_t* err_msg_c = new uint8_t[length];
+ auto clean = gtl::MakeCleanup([err_msg_c] { delete[] err_msg_c; });
+ TF_RETURN_IF_ERROR(client_->ReadData(err_msg_c, length));
+ string err_msg(reinterpret_cast<char*>(err_msg_c), length);
+
+ return errors::Unknown("Handshake Error [result=", handshake_res,
+ ", version=", serv_ver_major, ".", serv_ver_minor,
+ ".", serv_ver_patch, ", message='", err_msg, "']");
+ } else if (header == kNullVal) {
+ return errors::Unknown("Handshake Error [result=", handshake_res,
+ ", version=", serv_ver_major, ".", serv_ver_minor,
+ ".", serv_ver_patch, "]");
+ } else {
+ return errors::Unknown("Handshake Error [result=", handshake_res,
+ ", version=", serv_ver_major, ".", serv_ver_minor,
+ ".", serv_ver_patch, "]");
+ }
+ }
+
+ return Status::OK();
+}
+
+Status IgniteDatasetIterator::ScanQuery() {
+ TF_RETURN_IF_ERROR(client_->WriteInt(kScanQueryReqLength));
+ TF_RETURN_IF_ERROR(client_->WriteShort(kScanQueryOpcode));
+ TF_RETURN_IF_ERROR(client_->WriteLong(0)); // Request ID
+ TF_RETURN_IF_ERROR(
+ client_->WriteInt(JavaHashCode(cache_name_))); // Cache name
+ TF_RETURN_IF_ERROR(client_->WriteByte(0)); // Flags
+ TF_RETURN_IF_ERROR(client_->WriteByte(kNullVal)); // Filter object
+ TF_RETURN_IF_ERROR(client_->WriteInt(page_size_)); // Cursor page size
+ TF_RETURN_IF_ERROR(client_->WriteInt(part_)); // part_ition to query
+ TF_RETURN_IF_ERROR(client_->WriteByte(local_)); // local_ flag
+
+ uint64 wait_start = Env::Default()->NowMicros();
+ int32_t res_len;
+ TF_RETURN_IF_ERROR(client_->ReadInt(&res_len));
+ int64_t wait_stop = Env::Default()->NowMicros();
+
+ LOG(INFO) << "Scan Query waited " << (wait_stop - wait_start) / 1000 << " ms";
+
+ if (res_len < kMinResLength)
+ return errors::Unknown("Scan Query Response is corrupted");
+
+ int64_t req_id;
+ TF_RETURN_IF_ERROR(client_->ReadLong(&req_id));
+
+ int32_t status;
+ TF_RETURN_IF_ERROR(client_->ReadInt(&status));
+
+ if (status != 0) {
+ uint8_t err_msg_header;
+ TF_RETURN_IF_ERROR(client_->ReadByte(&err_msg_header));
+
+ if (err_msg_header == kStringVal) {
+ int32_t err_msg_length;
+ TF_RETURN_IF_ERROR(client_->ReadInt(&err_msg_length));
+
+ uint8_t* err_msg_c = new uint8_t[err_msg_length];
+ auto clean = gtl::MakeCleanup([err_msg_c] { delete[] err_msg_c; });
+ TF_RETURN_IF_ERROR(client_->ReadData(err_msg_c, err_msg_length));
+ string err_msg(reinterpret_cast<char*>(err_msg_c), err_msg_length);
+
+ return errors::Unknown("Scan Query Error [status=", status,
+ ", message=", err_msg, "]");
+ }
+ return errors::Unknown("Scan Query Error [status=", status, "]");
+ }
+
+ TF_RETURN_IF_ERROR(client_->ReadLong(&cursor_id_));
+
+ int32_t row_cnt;
+ TF_RETURN_IF_ERROR(client_->ReadInt(&row_cnt));
+
+ int32_t page_size = res_len - kScanQueryResHeaderLength;
+
+ return ReceivePage(page_size);
+}
+
+Status IgniteDatasetIterator::LoadNextPage() {
+ TF_RETURN_IF_ERROR(client_->WriteInt(kLoadNextPageReqLength));
+ TF_RETURN_IF_ERROR(client_->WriteShort(kLoadNextPageOpcode));
+ TF_RETURN_IF_ERROR(client_->WriteLong(0)); // Request ID
+ TF_RETURN_IF_ERROR(client_->WriteLong(cursor_id_)); // Cursor ID
+
+ uint64 wait_start = Env::Default()->NowMicros();
+ int32_t res_len;
+ TF_RETURN_IF_ERROR(client_->ReadInt(&res_len));
+ uint64 wait_stop = Env::Default()->NowMicros();
+
+ LOG(INFO) << "Load Next Page waited " << (wait_stop - wait_start) / 1000
+ << " ms";
+
+ if (res_len < kMinResLength)
+ return errors::Unknown("Load Next Page Response is corrupted");
+
+ int64_t req_id;
+ TF_RETURN_IF_ERROR(client_->ReadLong(&req_id));
+
+ int32_t status;
+ TF_RETURN_IF_ERROR(client_->ReadInt(&status));
+
+ if (status != 0) {
+ uint8_t err_msg_header;
+ TF_RETURN_IF_ERROR(client_->ReadByte(&err_msg_header));
+
+ if (err_msg_header == kStringVal) {
+ int32_t err_msg_length;
+ TF_RETURN_IF_ERROR(client_->ReadInt(&err_msg_length));
+
+ uint8_t* err_msg_c = new uint8_t[err_msg_length];
+ auto clean = gtl::MakeCleanup([err_msg_c] { delete[] err_msg_c; });
+ TF_RETURN_IF_ERROR(client_->ReadData(err_msg_c, err_msg_length));
+ string err_msg(reinterpret_cast<char*>(err_msg_c), err_msg_length);
+
+ return errors::Unknown("Load Next Page Error [status=", status,
+ ", message=", err_msg, "]");
+ }
+ return errors::Unknown("Load Next Page Error [status=", status, "]");
+ }
+
+ int32_t row_cnt;
+ TF_RETURN_IF_ERROR(client_->ReadInt(&row_cnt));
+
+ int32_t page_size = res_len - kLoadNextPageResHeaderLength;
+
+ return ReceivePage(page_size);
+}
+
+Status IgniteDatasetIterator::ReceivePage(int32_t page_size) {
+ remainder_ = page_size;
+ page_ = std::unique_ptr<uint8_t>(new uint8_t[remainder_]);
+ ptr_ = page_.get();
+
+ uint64 start = Env::Default()->NowMicros();
+ TF_RETURN_IF_ERROR(client_->ReadData(ptr_, remainder_));
+ uint64 stop = Env::Default()->NowMicros();
+
+ double size_in_mb = 1.0 * remainder_ / 1024 / 1024;
+ double time_in_s = 1.0 * (stop - start) / 1000 / 1000;
+ LOG(INFO) << "Page size " << size_in_mb << " Mb, time " << time_in_s * 1000
+ << " ms download speed " << size_in_mb / time_in_s << " Mb/sec";
+
+ uint8_t last_page_b;
+ TF_RETURN_IF_ERROR(client_->ReadByte(&last_page_b));
+
+ last_page_ = !last_page_b;
+
+ return Status::OK();
+}
+
+Status IgniteDatasetIterator::CheckTypes(const std::vector<int32_t>& types) {
+ if (schema_.size() != types.size())
+ return errors::Unknown("Object has unexpected schema");
+
+ for (size_t i = 0; i < schema_.size(); i++) {
+ if (schema_[i] != types[permutation_[i]])
+ return errors::Unknown("Object has unexpected schema");
+ }
+
+ return Status::OK();
+}
+
+int32_t IgniteDatasetIterator::JavaHashCode(string str) const {
+ int32_t h = 0;
+ for (char& c : str) {
+ h = 31 * h + c;
+ }
+ return h;
+}
+
+} // namespace tensorflow
diff --git a/tensorflow/contrib/ignite/kernels/ignite_dataset_iterator.h b/tensorflow/contrib/ignite/kernels/ignite_dataset_iterator.h
new file mode 100644
index 0000000..c499e2c
--- /dev/null
+++ b/tensorflow/contrib/ignite/kernels/ignite_dataset_iterator.h
@@ -0,0 +1,99 @@
+/* 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_CONTRIB_IGNITE_KERNELS_IGNITE_DATASET_ITERATOR_H_
+#define TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_DATASET_ITERATOR_H_
+
+#include "tensorflow/contrib/ignite/kernels/ignite_binary_object_parser.h"
+#include "tensorflow/contrib/ignite/kernels/ignite_client.h"
+#include "tensorflow/contrib/ignite/kernels/ignite_dataset.h"
+#include "tensorflow/core/platform/mutex.h"
+
+namespace tensorflow {
+
+class IgniteDatasetIterator : public DatasetIterator<IgniteDataset> {
+ public:
+ IgniteDatasetIterator(const Params& params, string host, int32 port,
+ string cache_name, bool local, int32 part,
+ int32 page_size, string username, string password,
+ string certfile, string keyfile, string cert_password,
+ std::vector<int32> schema,
+ std::vector<int32> permutation);
+ ~IgniteDatasetIterator();
+ Status GetNextInternal(IteratorContext* ctx, std::vector<Tensor>* out_tensors,
+ bool* end_of_sequence) override;
+
+ protected:
+ Status SaveInternal(IteratorStateWriter* writer) override;
+ Status RestoreInternal(IteratorContext* ctx,
+ IteratorStateReader* reader) override;
+
+ private:
+ Status GetNextInternalWithValidState(IteratorContext* ctx,
+ std::vector<Tensor>* out_tensors,
+ bool* end_of_sequence);
+
+ Status EstablishConnection();
+ Status CloseConnection();
+ Status Handshake();
+ Status ScanQuery();
+ Status LoadNextPage();
+ Status ReceivePage(int32_t page_size);
+ Status CheckTypes(const std::vector<int32_t>& types);
+ int32_t JavaHashCode(string str) const;
+
+ std::unique_ptr<Client> client_;
+ BinaryObjectParser parser_;
+
+ const string cache_name_;
+ const bool local_;
+ const int32 part_;
+ const int32 page_size_;
+ const string username_;
+ const string password_;
+ const std::vector<int32> schema_;
+ const std::vector<int32> permutation_;
+
+ int32_t remainder_;
+ int64_t cursor_id_;
+ bool last_page_;
+
+ bool valid_state_;
+
+ mutex mutex_;
+
+ std::unique_ptr<uint8_t> page_;
+ uint8_t* ptr_;
+};
+
+constexpr uint8_t kNullVal = 101;
+constexpr uint8_t kStringVal = 9;
+constexpr uint8_t kProtocolMajorVersion = 1;
+constexpr uint8_t kProtocolMinorVersion = 1;
+constexpr uint8_t kProtocolPatchVersion = 0;
+constexpr int16_t kScanQueryOpcode = 2000;
+constexpr int16_t kLoadNextPageOpcode = 2001;
+constexpr int16_t kCloseConnectionOpcode = 0;
+constexpr int32_t kScanQueryReqLength = 25;
+constexpr int32_t kScanQueryResHeaderLength = 25;
+constexpr int32_t kLoadNextPageReqLength = 18;
+constexpr int32_t kLoadNextPageResHeaderLength = 17;
+constexpr int32_t kCloseConnectionReqLength = 18;
+constexpr int32_t kHandshakeReqDefaultLength = 8;
+constexpr int32_t kMinResLength = 12;
+
+} // namespace tensorflow
+
+#endif // TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_DATASET_ITERATOR_H_
diff --git a/tensorflow/contrib/ignite/kernels/ignite_dataset_ops.cc b/tensorflow/contrib/ignite/kernels/ignite_dataset_ops.cc
new file mode 100644
index 0000000..f75b1c5
--- /dev/null
+++ b/tensorflow/contrib/ignite/kernels/ignite_dataset_ops.cc
@@ -0,0 +1,198 @@
+/* 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 <stdlib.h>
+
+#include "tensorflow/contrib/ignite/kernels/ignite_binary_object_parser.h"
+#include "tensorflow/contrib/ignite/kernels/ignite_dataset.h"
+#include "tensorflow/core/framework/dataset.h"
+#include "tensorflow/core/lib/strings/numbers.h"
+
+namespace tensorflow {
+namespace {
+
+Status SchemaToTypes(const std::vector<int32>& schema, DataTypeVector* dtypes) {
+ for (auto e : schema) {
+ if (e == BYTE || e == BYTE_ARR) {
+ dtypes->push_back(DT_UINT8);
+ } else if (e == SHORT || e == SHORT_ARR) {
+ dtypes->push_back(DT_INT16);
+ } else if (e == INT || e == INT_ARR) {
+ dtypes->push_back(DT_INT32);
+ } else if (e == LONG || e == LONG_ARR) {
+ dtypes->push_back(DT_INT64);
+ } else if (e == FLOAT || e == FLOAT_ARR) {
+ dtypes->push_back(DT_FLOAT);
+ } else if (e == DOUBLE || e == DOUBLE_ARR) {
+ dtypes->push_back(DT_DOUBLE);
+ } else if (e == USHORT || e == USHORT_ARR) {
+ dtypes->push_back(DT_UINT8);
+ } else if (e == BOOL || e == BOOL_ARR) {
+ dtypes->push_back(DT_BOOL);
+ } else if (e == STRING || e == STRING_ARR) {
+ dtypes->push_back(DT_STRING);
+ } else {
+ return errors::Unknown("Unexpected type in schema [type_id=", e, "]");
+ }
+ }
+
+ return Status::OK();
+}
+
+Status SchemaToShapes(const std::vector<int32>& schema,
+ std::vector<PartialTensorShape>* shapes) {
+ for (auto e : schema) {
+ if (e >= 1 && e < 10) {
+ shapes->push_back(PartialTensorShape({}));
+ } else if (e >= 12 && e < 21) {
+ shapes->push_back(PartialTensorShape({-1}));
+ } else {
+ return errors::Unknown("Unexpected type in schema [type_id=", e, "]");
+ }
+ }
+
+ return Status::OK();
+}
+
+class IgniteDatasetOp : public DatasetOpKernel {
+ public:
+ using DatasetOpKernel::DatasetOpKernel;
+
+ void MakeDataset(OpKernelContext* ctx, DatasetBase** output) override {
+ string cache_name = "";
+ string host = "";
+ int32 port = -1;
+ bool local = false;
+ int32 part = -1;
+ int32 page_size = -1;
+ string username = "";
+ string password = "";
+ string certfile = "";
+ string keyfile = "";
+ string cert_password = "";
+
+ const char* env_cache_name = std::getenv("IGNITE_DATASET_CACHE_NAME");
+ const char* env_host = std::getenv("IGNITE_DATASET_HOST");
+ const char* env_port = std::getenv("IGNITE_DATASET_PORT");
+ const char* env_local = std::getenv("IGNITE_DATASET_LOCAL");
+ const char* env_part = std::getenv("IGNITE_DATASET_PART");
+ const char* env_page_size = std::getenv("IGNITE_DATASET_PAGE_SIZE");
+ const char* env_username = std::getenv("IGNITE_DATASET_USERNAME");
+ const char* env_password = std::getenv("IGNITE_DATASET_PASSWORD");
+ const char* env_certfile = std::getenv("IGNITE_DATASET_CERTFILE");
+ const char* env_keyfile = std::getenv("IGNITE_DATASET_KEYFILE");
+ const char* env_cert_password = std::getenv("IGNITE_DATASET_CERT_PASSWORD");
+
+ if (env_cache_name) {
+ cache_name = string(env_cache_name);
+ } else {
+ OP_REQUIRES_OK(
+ ctx, ParseScalarArgument<string>(ctx, "cache_name", &cache_name));
+ }
+
+ if (env_host) {
+ host = string(env_host);
+ } else {
+ OP_REQUIRES_OK(ctx, ParseScalarArgument<string>(ctx, "host", &host));
+ }
+
+ if (env_port) {
+ OP_REQUIRES(ctx, strings::safe_strto32(env_port, &port),
+ errors::InvalidArgument("IGNITE_DATASET_PORT environment "
+ "variable is not a valid integer: ",
+ env_port));
+ } else {
+ OP_REQUIRES_OK(ctx, ParseScalarArgument<int32>(ctx, "port", &port));
+ }
+
+ if (env_local) {
+ local = true;
+ } else {
+ OP_REQUIRES_OK(ctx, ParseScalarArgument<bool>(ctx, "local", &local));
+ }
+
+ if (env_part) {
+ OP_REQUIRES(ctx, strings::safe_strto32(env_part, &part),
+ errors::InvalidArgument("IGNITE_DATASET_PART environment "
+ "variable is not a valid integer: ",
+ env_part));
+ } else {
+ OP_REQUIRES_OK(ctx, ParseScalarArgument<int32>(ctx, "part", &part));
+ }
+
+ if (env_page_size) {
+ OP_REQUIRES(ctx, strings::safe_strto32(env_page_size, &page_size),
+ errors::InvalidArgument("IGNITE_DATASET_PAGE_SIZE "
+ "environment variable is not a valid "
+ "integer: ",
+ env_page_size));
+ } else {
+ OP_REQUIRES_OK(ctx,
+ ParseScalarArgument<int32>(ctx, "page_size", &page_size));
+ }
+
+ if (env_username) username = string(env_username);
+
+ if (env_password) password = string(env_password);
+
+ if (env_certfile) certfile = string(env_certfile);
+
+ if (env_keyfile) keyfile = string(env_keyfile);
+
+ if (env_cert_password) cert_password = string(env_cert_password);
+
+ const Tensor* schema_tensor;
+ OP_REQUIRES_OK(ctx, ctx->input("schema", &schema_tensor));
+ OP_REQUIRES(ctx, schema_tensor->dims() == 1,
+ errors::InvalidArgument("`schema` must be a vector."));
+
+ std::vector<int32> schema;
+ schema.reserve(schema_tensor->NumElements());
+ for (int i = 0; i < schema_tensor->NumElements(); i++) {
+ schema.push_back(schema_tensor->flat<int32>()(i));
+ }
+
+ const Tensor* permutation_tensor;
+ OP_REQUIRES_OK(ctx, ctx->input("permutation", &permutation_tensor));
+ OP_REQUIRES(ctx, permutation_tensor->dims() == 1,
+ errors::InvalidArgument("`permutation` must be a vector."));
+
+ std::vector<int32> permutation;
+ permutation.resize(permutation_tensor->NumElements());
+ for (int i = 0; i < permutation_tensor->NumElements(); i++) {
+ // Inversed permutation.
+ permutation[permutation_tensor->flat<int32>()(i)] = i;
+ }
+
+ DataTypeVector dtypes;
+ std::vector<PartialTensorShape> shapes;
+
+ OP_REQUIRES_OK(ctx, SchemaToTypes(schema, &dtypes));
+ OP_REQUIRES_OK(ctx, SchemaToShapes(schema, &shapes));
+
+ *output = new IgniteDataset(
+ ctx, std::move(cache_name), std::move(host), port, local, part,
+ page_size, std::move(username), std::move(password),
+ std::move(certfile), std::move(keyfile), std::move(cert_password),
+ std::move(schema), std::move(permutation), std::move(dtypes),
+ std::move(shapes));
+ }
+};
+
+REGISTER_KERNEL_BUILDER(Name("IgniteDataset").Device(DEVICE_CPU),
+ IgniteDatasetOp);
+
+} // namespace
+} // namespace tensorflow
diff --git a/tensorflow/contrib/ignite/kernels/ignite_plain_client.h b/tensorflow/contrib/ignite/kernels/ignite_plain_client.h
new file mode 100644
index 0000000..75424c1
--- /dev/null
+++ b/tensorflow/contrib/ignite/kernels/ignite_plain_client.h
@@ -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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_PLAIN_CLIENT_H_
+#define TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_PLAIN_CLIENT_H_
+
+#include "tensorflow/contrib/ignite/kernels/ignite_client.h"
+
+namespace tensorflow {
+
+class PlainClient : public Client {
+ public:
+ PlainClient(string host, int port, bool big_endian);
+ ~PlainClient();
+
+ Status Connect() override;
+ Status Disconnect() override;
+ bool IsConnected() override;
+ int GetSocketDescriptor() override;
+ Status ReadData(uint8_t* buf, const int32_t length) override;
+ Status WriteData(const uint8_t* buf, const int32_t length) override;
+
+ private:
+ const string host_;
+ const int port_;
+ int sock_;
+};
+
+} // namespace tensorflow
+
+#endif // TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_PLAIN_CLIENT_H_
diff --git a/tensorflow/contrib/ignite/kernels/ignite_plain_client_unix.cc b/tensorflow/contrib/ignite/kernels/ignite_plain_client_unix.cc
new file mode 100644
index 0000000..cf67294
--- /dev/null
+++ b/tensorflow/contrib/ignite/kernels/ignite_plain_client_unix.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/contrib/ignite/kernels/ignite_plain_client.h"
+
+#include <arpa/inet.h>
+#include <netdb.h>
+#include <sys/socket.h>
+#include <unistd.h>
+
+#include <iostream>
+#include <map>
+
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/platform/logging.h"
+
+namespace tensorflow {
+
+PlainClient::PlainClient(string host, int port, bool big_endian)
+ : Client(big_endian), host_(std::move(host)), port_(port), sock_(-1) {}
+
+PlainClient::~PlainClient() {
+ if (IsConnected()) {
+ Status status = Disconnect();
+ if (!status.ok()) LOG(WARNING) << status.ToString();
+ }
+}
+
+Status PlainClient::Connect() {
+ if (sock_ == -1) {
+ sock_ = socket(AF_INET, SOCK_STREAM, 0);
+ if (sock_ == -1) return errors::Internal("Failed to create socket");
+ }
+
+ sockaddr_in server;
+
+ server.sin_addr.s_addr = inet_addr(host_.c_str());
+ if (server.sin_addr.s_addr == -1) {
+ hostent* he;
+ in_addr** addr_list;
+
+ if ((he = gethostbyname(host_.c_str())) == NULL)
+ return errors::Internal("Failed to resolve hostname \"", host_, "\"");
+
+ addr_list = (in_addr**)he->h_addr_list;
+ if (addr_list[0] != NULL) server.sin_addr = *addr_list[0];
+ }
+
+ server.sin_family = AF_INET;
+ server.sin_port = htons(port_);
+
+ if (connect(sock_, (sockaddr*)&server, sizeof(server)) < 0)
+ return errors::Internal("Failed to connect to \"", host_, ":", port_, "\"");
+
+ LOG(INFO) << "Connection to \"" << host_ << ":" << port_ << "\" established";
+
+ return Status::OK();
+}
+
+Status PlainClient::Disconnect() {
+ int close_res = close(sock_);
+ sock_ = -1;
+
+ LOG(INFO) << "Connection to \"" << host_ << ":" << port_ << "\" is closed";
+
+ return close_res == 0
+ ? Status::OK()
+ : errors::Internal("Failed to correctly close connection");
+}
+
+bool PlainClient::IsConnected() { return sock_ != -1; }
+
+int PlainClient::GetSocketDescriptor() { return sock_; }
+
+Status PlainClient::ReadData(uint8_t* buf, const int32_t length) {
+ int received = 0;
+
+ while (received < length) {
+ int res = recv(sock_, buf, length - received, 0);
+
+ if (res < 0)
+ return errors::Internal("Error occurred while reading from socket: ", res,
+ ", ", string(strerror(errno)));
+
+ if (res == 0) return errors::Internal("Server closed connection");
+
+ received += res;
+ buf += res;
+ }
+
+ return Status::OK();
+}
+
+Status PlainClient::WriteData(const uint8_t* buf, const int32_t length) {
+ int sent = 0;
+
+ while (sent < length) {
+ int res = send(sock_, buf, length - sent, 0);
+
+ if (res < 0)
+ return errors::Internal("Error occurred while writing into socket: ", res,
+ ", ", string(strerror(errno)));
+
+ sent += res;
+ buf += res;
+ }
+
+ return Status::OK();
+}
+
+} // namespace tensorflow
diff --git a/tensorflow/contrib/ignite/kernels/ignite_plain_client_windows.cc b/tensorflow/contrib/ignite/kernels/ignite_plain_client_windows.cc
new file mode 100644
index 0000000..dad5aac
--- /dev/null
+++ b/tensorflow/contrib/ignite/kernels/ignite_plain_client_windows.cc
@@ -0,0 +1,142 @@
+/* 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/contrib/ignite/kernels/ignite_plain_client.h"
+
+#define WIN32_LEAN_AND_MEAN
+#include <windows.h>
+#include <winsock2.h>
+#include <ws2tcpip.h>
+
+#pragma comment(lib, "Ws2_32.lib")
+#pragma comment(lib, "Mswsock.lib")
+#pragma comment(lib, "AdvApi32.lib")
+
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/gtl/cleanup.h"
+#include "tensorflow/core/platform/logging.h"
+
+namespace tensorflow {
+
+PlainClient::PlainClient(string host, int port, bool big_endian)
+ : Client(big_endian),
+ host_(std::move(host)),
+ port_(port),
+ sock_(INVALID_SOCKET) {}
+
+PlainClient::~PlainClient() {
+ if (IsConnected()) {
+ Status status = Disconnect();
+ if (!status.ok()) LOG(WARNING) << status.ToString();
+ }
+}
+
+Status PlainClient::Connect() {
+ WSADATA wsaData;
+ addrinfo *result = NULL, *ptr = NULL, hints;
+
+ int res = WSAStartup(MAKEWORD(2, 2), &wsaData);
+ if (res != 0) return errors::Internal("WSAStartup failed with error: ", res);
+
+ ZeroMemory(&hints, sizeof(hints));
+ hints.ai_family = AF_UNSPEC;
+ hints.ai_socktype = SOCK_STREAM;
+ hints.ai_protocol = IPPROTO_TCP;
+
+ res = getaddrinfo(host_.c_str(), std::to_string(port_).c_str(), &hints,
+ &result);
+ if (res != 0) return errors::Internal("Getaddrinfo failed with error: ", res);
+
+ auto clean = gtl::MakeCleanup([result] { freeaddrinfo(result); });
+
+ for (ptr = result; ptr != NULL; ptr = ptr->ai_next) {
+ sock_ = socket(ptr->ai_family, ptr->ai_socktype, ptr->ai_protocol);
+ if (sock_ == INVALID_SOCKET) {
+ WSACleanup();
+ return errors::Internal("Socket failed with error: ", WSAGetLastError());
+ }
+
+ res = connect(sock_, ptr->ai_addr, (int)ptr->ai_addrlen);
+ if (res == SOCKET_ERROR) {
+ closesocket(sock_);
+ sock_ = INVALID_SOCKET;
+ continue;
+ }
+
+ break;
+ }
+
+ if (sock_ == INVALID_SOCKET) {
+ WSACleanup();
+ return errors::Internal("Unable to connect to server");
+ }
+
+ LOG(INFO) << "Connection to \"" << host_ << ":" << port_ << "\" established";
+
+ return Status::OK();
+}
+
+Status PlainClient::Disconnect() {
+ int res = shutdown(sock_, SD_SEND);
+ closesocket(sock_);
+ WSACleanup();
+
+ if (res == SOCKET_ERROR)
+ return errors::Internal("Shutdown failed with error: ", WSAGetLastError());
+ else
+ return Status::OK();
+}
+
+bool PlainClient::IsConnected() { return sock_ != INVALID_SOCKET; }
+
+int PlainClient::GetSocketDescriptor() { return sock_; }
+
+Status PlainClient::ReadData(uint8_t *buf, const int32_t length) {
+ int received = 0;
+
+ while (received < length) {
+ int res = recv(sock_, (char *)buf, length - received, 0);
+
+ if (res < 0)
+ return errors::Internal("Error occurred while reading from socket: ",
+ res);
+
+ if (res == 0) return errors::Internal("Server closed connection");
+
+ received += res;
+ buf += res;
+ }
+
+ return Status::OK();
+}
+
+Status PlainClient::WriteData(const uint8_t *buf, const int32_t length) {
+ int sent = 0;
+
+ while (sent < length) {
+ int res = send(sock_, (char *)buf, length - sent, 0);
+
+ if (res < 0)
+ return errors::Internal("Error occurred while writing into socket: ",
+ res);
+
+ sent += res;
+ buf += res;
+ }
+
+ return Status::OK();
+}
+
+} // namespace tensorflow
diff --git a/tensorflow/contrib/ignite/kernels/ignite_ssl_wrapper.cc b/tensorflow/contrib/ignite/kernels/ignite_ssl_wrapper.cc
new file mode 100644
index 0000000..ceb479b
--- /dev/null
+++ b/tensorflow/contrib/ignite/kernels/ignite_ssl_wrapper.cc
@@ -0,0 +1,151 @@
+/* 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/contrib/ignite/kernels/ignite_ssl_wrapper.h"
+
+#include <openssl/err.h>
+#include <openssl/ssl.h>
+
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/platform/logging.h"
+
+namespace tensorflow {
+
+static int PasswordCb(char *buf, int size, int rwflag, void *password) {
+ strncpy(buf, (char *)(password), size);
+ buf[size - 1] = '\0';
+ return (strlen(buf));
+}
+
+SslWrapper::SslWrapper(std::shared_ptr<Client> client, string certfile,
+ string keyfile, string cert_password, bool big_endian)
+ : Client(big_endian),
+ client_(client),
+ certfile_(std::move(certfile)),
+ keyfile_(std::move(keyfile)),
+ cert_password_(std::move(cert_password)),
+ ctx_(nullptr),
+ ssl_(nullptr) {}
+
+SslWrapper::~SslWrapper() {
+ if (IsConnected()) {
+ Status status = Disconnect();
+ if (!status.ok()) LOG(WARNING) << status.ToString();
+ }
+
+ if (ctx_ != nullptr) {
+ SSL_CTX_free(ctx_);
+ ctx_ = nullptr;
+ }
+
+ if (ssl_ != nullptr) {
+ SSL_free(ssl_);
+ ssl_ = nullptr;
+ }
+}
+
+Status SslWrapper::InitSslContext() {
+ OpenSSL_add_all_algorithms();
+ SSL_load_error_strings();
+
+ ctx_ = SSL_CTX_new(SSLv23_method());
+ if (ctx_ == NULL) return errors::Internal("Couldn't create SSL context");
+
+ SSL_CTX_set_default_passwd_cb(ctx_, PasswordCb);
+ SSL_CTX_set_default_passwd_cb_userdata(ctx_, (void *)cert_password_.c_str());
+
+ if (SSL_CTX_use_certificate_chain_file(ctx_, certfile_.c_str()) != 1)
+ return errors::Internal("Couldn't load cetificate chain (file '", certfile_,
+ "')");
+
+ string private_key_file = keyfile_.empty() ? certfile_ : keyfile_;
+ if (SSL_CTX_use_PrivateKey_file(ctx_, private_key_file.c_str(),
+ SSL_FILETYPE_PEM) != 1)
+ return errors::Internal("Couldn't load private key (file '",
+ private_key_file, "')");
+
+ return Status::OK();
+}
+
+Status SslWrapper::Connect() {
+ if (ctx_ == NULL) {
+ TF_RETURN_IF_ERROR(InitSslContext());
+ }
+
+ ssl_ = SSL_new(ctx_);
+ if (ssl_ == NULL)
+ return errors::Internal("Failed to establish SSL connection");
+
+ TF_RETURN_IF_ERROR(client_->Connect());
+
+ SSL_set_fd(ssl_, client_->GetSocketDescriptor());
+ if (SSL_connect(ssl_) != 1)
+ return errors::Internal("Failed to establish SSL connection");
+
+ LOG(INFO) << "SSL connection established";
+
+ return Status::OK();
+}
+
+Status SslWrapper::Disconnect() {
+ SSL_free(ssl_);
+ ssl_ = nullptr;
+
+ LOG(INFO) << "SSL connection closed";
+
+ return client_->Disconnect();
+}
+
+bool SslWrapper::IsConnected() { return client_->IsConnected(); }
+
+int SslWrapper::GetSocketDescriptor() { return client_->GetSocketDescriptor(); }
+
+Status SslWrapper::ReadData(uint8_t *buf, const int32_t length) {
+ int received = 0;
+
+ while (received < length) {
+ int res = SSL_read(ssl_, buf, length - received);
+
+ if (res < 0)
+ return errors::Internal("Error occurred while reading from SSL socket: ",
+ res);
+
+ if (res == 0) return errors::Internal("Server closed SSL connection");
+
+ received += res;
+ buf += res;
+ }
+
+ return Status::OK();
+}
+
+Status SslWrapper::WriteData(const uint8_t *buf, const int32_t length) {
+ int sent = 0;
+
+ while (sent < length) {
+ int res = SSL_write(ssl_, buf, length - sent);
+
+ if (res < 0)
+ return errors::Internal("Error occurred while writing into socket: ",
+ res);
+
+ sent += res;
+ buf += res;
+ }
+
+ return Status::OK();
+}
+
+} // namespace tensorflow
diff --git a/tensorflow/contrib/ignite/kernels/ignite_ssl_wrapper.h b/tensorflow/contrib/ignite/kernels/ignite_ssl_wrapper.h
new file mode 100644
index 0000000..0406644
--- /dev/null
+++ b/tensorflow/contrib/ignite/kernels/ignite_ssl_wrapper.h
@@ -0,0 +1,51 @@
+/* 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_CONTRIB_IGNITE_KERNELS_IGNITE_SSL_WRAPPER_H_
+#define TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_SSL_WRAPPER_H_
+
+#include "tensorflow/contrib/ignite/kernels/ignite_client.h"
+
+#include <openssl/ssl.h>
+
+namespace tensorflow {
+
+class SslWrapper : public Client {
+ public:
+ SslWrapper(std::shared_ptr<Client> client, string certfile, string keyfile,
+ string cert_password, bool big_endian);
+ ~SslWrapper();
+
+ Status Connect() override;
+ Status Disconnect() override;
+ bool IsConnected() override;
+ int GetSocketDescriptor() override;
+ Status ReadData(uint8_t* buf, const int32_t length) override;
+ Status WriteData(const uint8_t* buf, const int32_t length) override;
+
+ private:
+ Status InitSslContext();
+
+ std::shared_ptr<Client> client_;
+ string certfile_;
+ string keyfile_;
+ string cert_password_;
+ SSL_CTX* ctx_;
+ SSL* ssl_;
+};
+
+} // namespace tensorflow
+
+#endif // TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_SSL_WRAPPER_H_
diff --git a/tensorflow/contrib/ignite/ops/dataset_ops.cc b/tensorflow/contrib/ignite/ops/dataset_ops.cc
new file mode 100644
index 0000000..3d6fbe0
--- /dev/null
+++ b/tensorflow/contrib/ignite/ops/dataset_ops.cc
@@ -0,0 +1,56 @@
+/* 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/framework/common_shape_fns.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/shape_inference.h"
+
+namespace tensorflow {
+
+REGISTER_OP("IgniteDataset")
+ .Input("cache_name: string")
+ .Input("host: string")
+ .Input("port: int32")
+ .Input("local: bool")
+ .Input("part: int32")
+ .Input("page_size: int32")
+ .Input("schema: int32")
+ .Input("permutation: int32")
+ .Output("handle: variant")
+ .SetIsStateful()
+ .SetShapeFn(shape_inference::ScalarShape)
+ .Doc(R"doc(
+IgniteDataset that allows to get data from Apache Ignite.
+
+Apache Ignite is a memory-centric distributed database, 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.
+
+cache_name: Ignite Cache Name.
+host: Ignite Thin Client Host.
+port: Ignite Thin Client Port.
+local: Local flag that defines that data should be fetched from local host only.
+part: Partition data should be fetched from.
+page_size: Page size for Ignite Thin Client.
+schema: Internal structure that defines schema of cache objects.
+permutation: Internal structure that defines permutation of cache objects.
+)doc");
+
+} // namespace tensorflow
diff --git a/tensorflow/contrib/ignite/python/ops/ignite_dataset_ops.py b/tensorflow/contrib/ignite/python/ops/ignite_dataset_ops.py
new file mode 100644
index 0000000..cfe59b6
--- /dev/null
+++ b/tensorflow/contrib/ignite/python/ops/ignite_dataset_ops.py
@@ -0,0 +1,772 @@
+# 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.
+# ==============================================================================
+"""Ignite Dataset."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import abc
+import socket
+import ssl
+import struct
+
+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.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+
+
+class Readable(object):
+ """Readable abstract class that exposes methods to do reading-related
+
+ operations.
+ """
+
+ @abc.abstractmethod
+ def __init__(self):
+ pass
+
+ def read_byte(self):
+ """Reads and returnes byte."""
+ return self._read("b", 1)
+
+ def read_short(self):
+ """Reads and returns short (2 bytes, little-endian)."""
+ return self._read("h", 2)
+
+ def read_int(self):
+ """Reads and returns int (4 bytes, little-endian)."""
+ return self._read("i", 4)
+
+ def read_long(self):
+ """Reads and returns long (8 bytes, little-endian)."""
+ return self._read("q", 8)
+
+ def skip(self, length):
+ """Skips the specified number of bytes."""
+ self.read_data(length)
+
+ @abc.abstractmethod
+ def read_data(self, length):
+ """Reads the specified number of bytes and returns them as a buffer."""
+ return None
+
+ def _read(self, data_type, length):
+ """Reads, unpacks and returns specified type (little-endian)."""
+ data_buffer = self.read_data(length)
+ return struct.unpack("<" + data_type, data_buffer)[0]
+
+
+class DataBuffer(Readable):
+ """DataBuffer class that exposes methods to read data from a byte buffer."""
+
+ def __init__(self, data_buffer):
+ """Constructs a new instance based on the specified byte buffer.
+
+ Args:
+ data_buffer: Buffer to be read.
+ """
+ Readable.__init__(self)
+ self.buffer = data_buffer
+ self.ptr = 0
+
+ def read_data(self, length):
+ """Reads the specified number of bytes and returns them as a buffer."""
+ data_buffer = self.buffer[self.ptr:][:length]
+ self.ptr += length
+ return data_buffer
+
+
+class TcpClient(Readable):
+ """TcpClient class that exposes methods to read data from a socket."""
+
+ def __init__(self, host, port, certfile=None, keyfile=None, password=None):
+ """Constructs a new instance based on the specified host and port.
+
+ Args:
+ host: Host to be connected.
+ port: Port to be connected.
+ certfile: File in PEM format containing the certificate as well as any
+ number of CA certificates needed to establish the certificate's
+ authenticity.
+ keyfile: File containing the private key (otherwise the private key will
+ be taken from certfile as well).
+ password: Password to be used if the private key is encrypted and a
+ password is necessary.
+
+ Raises:
+ ValueError: If the wrong combination of arguments is provided.
+ """
+ Readable.__init__(self)
+ self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+
+ if certfile is not None:
+ context = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
+ context.load_cert_chain(certfile, keyfile, password)
+ self.sock = context.wrap_socket(self.sock)
+ else:
+ if keyfile is not None:
+ raise ValueError("SSL is disabled, keyfile must not be specified "
+ "(to enable SSL specify certfile)")
+ if password is not None:
+ raise ValueError("SSL is disabled, password must not be specified "
+ "(to enable SSL specify certfile)")
+
+ self.host = host
+ self.port = port
+
+ def __enter__(self):
+ """Connects to host and port specified in the constructor."""
+ self.sock.connect((self.host, self.port))
+ return self
+
+ def __exit__(self, t, v, traceback):
+ """Disconnects the socket."""
+ self.sock.close()
+
+ def write_byte(self, v):
+ """Writes the specified byte."""
+ self._write(v, "b")
+
+ def write_short(self, v):
+ """Writes the specified short (2 bytes, little-endian)."""
+ self._write(v, "h")
+
+ def write_int(self, v):
+ """Writes the specified short (4 bytes, little-endian)."""
+ self._write(v, "i")
+
+ def write_long(self, v):
+ """Writes the specified int (8 bytes, little-endian)."""
+ self._write(v, "q")
+
+ def write_string(self, v):
+ """Writes the specified string."""
+ self.sock.sendall(v.encode("UTF-8"))
+
+ def read_data(self, length):
+ """Reads the specified number of bytes and returns them as a buffer."""
+ data_buffer = None
+ rem = length
+ while rem > 0:
+ buf = self.sock.recv(rem)
+ rem = rem - len(buf)
+ if data_buffer is None:
+ data_buffer = buf
+ else:
+ data_buffer += buf
+ return data_buffer
+
+ def _write(self, value, data_type):
+ """Packs and writes data using the specified type (little-endian)."""
+ data_buffer = struct.pack("<" + data_type, value)
+ self.sock.sendall(data_buffer)
+
+
+class BinaryType(object):
+ """BinaryType class that encapsulated type id, type name and fields."""
+
+ def __init__(self, type_id, type_name, fields):
+ """Constructs a new instance of BinaryType."""
+ self.type_id = type_id
+ self.type_name = type_name
+ self.fields = fields
+
+
+class BinaryField(object):
+ """BinaryField class that encapsulated field name, type id and field id."""
+
+ def __init__(self, field_name, type_id, field_id):
+ """Constructs a new instance of BinaryField."""
+ self.field_name = field_name
+ self.type_id = type_id
+ self.field_id = field_id
+
+
+# Binary types defined in Apache Ignite Thin client and supported by
+# TensorFlow on Apache Ignite, see
+# https://apacheignite.readme.io/v2.6/docs/binary-client-protocol.
+# True means that type is a vector, False means type is scalar.
+types = {
+ 1: (dtypes.uint8, False),
+ 2: (dtypes.int16, False),
+ 3: (dtypes.int32, False),
+ 4: (dtypes.int64, False),
+ 5: (dtypes.float32, False),
+ 6: (dtypes.float64, False),
+ 7: (dtypes.uint16, False),
+ 8: (dtypes.bool, False),
+ 9: (dtypes.string, False),
+ 12: (dtypes.uint8, True),
+ 13: (dtypes.int16, True),
+ 14: (dtypes.int32, True),
+ 15: (dtypes.int64, True),
+ 16: (dtypes.float32, True),
+ 17: (dtypes.float64, True),
+ 18: (dtypes.uint16, True),
+ 19: (dtypes.bool, True),
+ 20: (dtypes.string, True)
+}
+
+
+class TypeTreeNode(object):
+ """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.
+
+ Args:
+ name: Name of the object tree node.
+ type_id: Type id of the object tree node.
+ fields: List of fields (children of the object tree node).
+ permutation: Permutation that should be applied to order object children.
+ """
+ self.name = name
+ self.type_id = type_id
+ self.fields = fields
+ self.permutation = permutation
+
+ def to_output_classes(self):
+ """Formats the tree object as required by `Dataset.output_classes`."""
+ if self.fields is None:
+ return ops.Tensor
+ output_classes = {}
+ for field in self.fields:
+ output_classes[field.name] = field.to_output_classes()
+ return output_classes
+
+ def to_output_shapes(self):
+ """Formats the tree object as required by `Dataset.output_shapes`."""
+ if self.fields is None:
+ if self.type_id in types:
+ object_type = types[self.type_id]
+ is_array = object_type[1]
+ if is_array:
+ return tensor_shape.TensorShape([None])
+ return tensor_shape.TensorShape([])
+ raise ValueError("Unsupported type [type_id=%d]" % self.type_id)
+ output_shapes = {}
+ for field in self.fields:
+ output_shapes[field.name] = field.to_output_shapes()
+ return output_shapes
+
+ def to_output_types(self):
+ """Formats the tree object as required by `Dataset.output_types`."""
+ if self.fields is None:
+ if self.type_id in types:
+ object_type = types[self.type_id]
+ return object_type[0]
+ raise ValueError("Unsupported type [type_id=%d]" % self.type_id)
+ else:
+ output_types = {}
+ for field in self.fields:
+ output_types[field.name] = field.to_output_types()
+ return output_types
+
+ def to_flat(self):
+ """Returns a list of node types."""
+ return self.to_flat_rec([])
+
+ def to_permutation(self):
+ """Returns a permutation that should be applied to order object leaves."""
+ correct_order_dict = {}
+ self.traversal_rec(correct_order_dict, 0)
+ object_order = []
+ self.traversal_permutation_rec(object_order)
+ return [correct_order_dict[o] for o in object_order]
+
+ def to_flat_rec(self, flat):
+ """Formats a list of leaf node types in pre-order."""
+ if self.fields is None:
+ flat.append(self.type_id)
+ else:
+ for field in self.fields:
+ field.to_flat_rec(flat)
+ return flat
+
+ def traversal_permutation_rec(self, permutation):
+ """Collects nodes in accordance with permutation."""
+ if self.fields is None:
+ permutation.append(self)
+ else:
+ for idx in self.permutation:
+ field = self.fields[idx]
+ field.traversal_permutation_rec(permutation)
+
+ def traversal_rec(self, d, i):
+ """Collects nodes in pre-order traversal."""
+ if self.fields is None:
+ d[self] = i
+ i += 1
+ else:
+ for field in self.fields:
+ i = field.traversal_rec(d, i)
+ return i
+
+
+class IgniteClient(TcpClient):
+ """IgniteClient enables working with Apache Ignite using a thin client.
+
+ This client works with assumption that all object in the cache
+ have the same structure (homogeneous objects) and the cache contains at
+ least one object.
+ """
+
+ def __init__(self,
+ host,
+ port,
+ username=None,
+ password=None,
+ certfile=None,
+ keyfile=None,
+ cert_password=None):
+ """Constructs a new instance of IgniteClient.
+
+ Args:
+ host: Apache Ignite Thin client host to be connected.
+ port: Apache Ignite Thin client port to be connected.
+ username: Apache Ignite Thin Client authentication username.
+ password: Apache Ignite Thin Client authentication password.
+ certfile: File in PEM format containing the certificate as well as any
+ number of CA certificates needed to establish the certificate's
+ authenticity.
+ keyfile: File containing the private key (otherwise the private key will
+ be taken from certfile as well).
+ cert_password: Password to be used if the private key is encrypted and a
+ password is necessary.
+ """
+ TcpClient.__init__(self, host, port, certfile, keyfile, cert_password)
+ self.username = username
+ self.password = password
+
+ def handshake(self):
+ """Makes a handshake after connect and before any other calls."""
+ msg_len = 8
+
+ if self.username is None:
+ msg_len += 1
+ else:
+ msg_len += 5 + len(self.username)
+
+ if self.password is None:
+ msg_len += 1
+ else:
+ msg_len += 5 + len(self.password)
+
+ self.write_int(msg_len) # Message length
+ self.write_byte(1) # Handshake operation
+ self.write_short(1) # Version (1.1.0)
+ self.write_short(1)
+ self.write_short(0)
+ self.write_byte(2) # Thin client
+
+ if self.username is None: # Username
+ self.write_byte(101)
+ else:
+ self.write_byte(9)
+ self.write_int(len(self.username))
+ self.write_string(self.username)
+
+ if self.password is None: # Password
+ self.write_byte(101)
+ else:
+ self.write_byte(9)
+ self.write_int(len(self.password))
+ self.write_string(self.password)
+
+ self.read_int() # Result length
+ res = self.read_byte()
+
+ if res != 1:
+ serv_ver_major = self.read_short()
+ serv_ver_minor = self.read_short()
+ serv_ver_patch = self.read_short()
+ err_msg = self._parse_string()
+ if err_msg is None:
+ raise RuntimeError(
+ "Handshake Error [result=%d, version=%d.%d.%d]" %
+ (res, serv_ver_major, serv_ver_minor, serv_ver_patch))
+ else:
+ raise RuntimeError(
+ "Handshake Error [result=%d, version=%d.%d.%d, message='%s']" %
+ (res, serv_ver_major, serv_ver_minor, serv_ver_patch, err_msg))
+
+ def get_cache_type(self, cache_name):
+ """Collects type information about objects stored in the specified cache."""
+ cache_name_hash = self._java_hash_code(cache_name)
+ self.write_int(25) # Message length
+ self.write_short(2000) # Operation code
+ self.write_long(0) # Request ID
+ self.write_int(cache_name_hash) # Cache name
+ self.write_byte(0) # Flags
+ self.write_byte(101) # Filter (NULL)
+ self.write_int(1) # Cursor page size
+ self.write_int(-1) # Partition to query
+ self.write_byte(0) # Local flag
+
+ result_length = self.read_int()
+ self.read_long() # Request id
+ status = self.read_int()
+
+ if status != 0:
+ err_msg = self._parse_string()
+ if err_msg is None:
+ raise RuntimeError("Scan Query Error [status=%s]" % status)
+ else:
+ raise RuntimeError(
+ "Scan Query Error [status=%s, message='%s']" % (status, err_msg))
+
+ self.read_long() # Cursor id
+ row_count = self.read_int()
+
+ if row_count == 0:
+ raise RuntimeError("Scan Query returned empty result, so it's "
+ "impossible to derive the cache type")
+
+ payload = DataBuffer(self.read_data(result_length - 25))
+
+ self.read_byte() # Next page
+
+ res = TypeTreeNode("root", 0, [
+ self._collect_types("key", payload),
+ self._collect_types("val", payload)
+ ], [0, 1])
+
+ return res
+
+ def _java_hash_code(self, s):
+ """Computes hash code of the specified string using Java code."""
+ h = 0
+ for c in s:
+ h = (31 * h + ord(c)) & 0xFFFFFFFF
+ return ((h + 0x80000000) & 0xFFFFFFFF) - 0x80000000
+
+ def _collect_types(self, field_name, data):
+ """Extracts type information from the specified object."""
+ type_id = data.read_byte()
+
+ # Byte scalar.
+ if type_id == 1:
+ data.skip(1)
+ return TypeTreeNode(field_name, type_id)
+
+ # Short scalar.
+ if type_id == 2:
+ data.skip(2)
+ return TypeTreeNode(field_name, type_id)
+
+ # Integer scalar.
+ if type_id == 3:
+ data.skip(4)
+ return TypeTreeNode(field_name, type_id)
+
+ # Long scalar.
+ if type_id == 4:
+ data.skip(8)
+ return TypeTreeNode(field_name, type_id)
+
+ # Float scalar.
+ if type_id == 5:
+ data.skip(4)
+ return TypeTreeNode(field_name, type_id)
+
+ # Double scalar.
+ if type_id == 6:
+ data.skip(8)
+ return TypeTreeNode(field_name, type_id)
+
+ # Char scalar.
+ if type_id == 7:
+ data.skip(2)
+ return TypeTreeNode(field_name, type_id)
+
+ # Bool scalar.
+ if type_id == 8:
+ data.skip(1)
+ return TypeTreeNode(field_name, type_id)
+
+ # String scalar.
+ if type_id == 9:
+ length = data.read_int()
+ data.skip(length)
+ return TypeTreeNode(field_name, type_id)
+
+ # UUID scalar.
+ if type_id == 10:
+ data.skip(16)
+ return TypeTreeNode(field_name, type_id)
+
+ # Date scalar.
+ if type_id == 11:
+ data.skip(8)
+ return TypeTreeNode(field_name, type_id)
+
+ # Byte array.
+ if type_id == 12:
+ length = data.read_int()
+ data.skip(length)
+ return TypeTreeNode(field_name, type_id)
+
+ # Short array.
+ if type_id == 13:
+ length = data.read_int()
+ data.skip(length * 2)
+ return TypeTreeNode(field_name, type_id)
+
+ # Integer array.
+ if type_id == 14:
+ length = data.read_int()
+ data.skip(length * 4)
+ return TypeTreeNode(field_name, type_id)
+
+ # Long array.
+ if type_id == 15:
+ length = data.read_int()
+ data.skip(length * 8)
+ return TypeTreeNode(field_name, type_id)
+
+ # Float array.
+ if type_id == 16:
+ length = data.read_int()
+ data.skip(length * 4)
+ return TypeTreeNode(field_name, type_id)
+
+ # Double array.
+ if type_id == 17:
+ length = data.read_int()
+ data.skip(length * 8)
+ return TypeTreeNode(field_name, type_id)
+
+ # Char array.
+ if type_id == 18:
+ length = data.read_int()
+ data.skip(length * 2)
+ return TypeTreeNode(field_name, type_id)
+
+ # Bool array.
+ if type_id == 19:
+ length = data.read_int()
+ data.skip(length)
+ return TypeTreeNode(field_name, type_id)
+
+ # String array.
+ if type_id == 20:
+ length = data.read_int()
+ for _ in range(length):
+ header = data.read_byte()
+ if header == 9:
+ str_length = data.read_int()
+ data.skip(str_length)
+ elif header == 101:
+ pass
+ else:
+ raise RuntimeError(
+ "Unknown binary type when expected string [type_id=%d]" % header)
+ return TypeTreeNode(field_name, type_id)
+
+ # UUID array.
+ if type_id == 21:
+ length = data.read_int()
+ data.skip(length * 16) # TODO(dmitrievanthony): support NULL values.
+ return TypeTreeNode(field_name, type_id)
+
+ # Date array.
+ if type_id == 22:
+ length = data.read_int()
+ data.skip(length * 8)
+ return TypeTreeNode(field_name, type_id)
+
+ # Wrapped Binary Object.
+ if type_id == 27:
+ length = data.read_int()
+ inner_data = data.read_data(length)
+ data.read_int() # Offset
+ return self._collect_types(field_name, DataBuffer(inner_data))
+
+ # Complex Object.
+ if type_id == 103:
+ data.read_byte() # Object version
+ data.read_short() # Object flags
+ obj_type_id = data.read_int()
+ data.read_int() # Object hash code
+ obj_length = data.read_int()
+ data.read_int() # Object schema id
+ obj_schema_offset = data.read_int()
+
+ obj_type = self._get_type(obj_type_id)
+ children = []
+
+ for obj_field in obj_type.fields:
+ child = self._collect_types(obj_field.field_name, data)
+ children.append(child)
+
+ children_sorted = sorted(children, key=lambda child: child.name)
+ permutation = [children_sorted.index(child) for child in children]
+ children = children_sorted
+
+ data.skip(obj_length - obj_schema_offset)
+
+ return TypeTreeNode(field_name, type_id, children, permutation)
+
+ raise RuntimeError("Unknown binary type [type_id=%d]" % type_id)
+
+ def _get_type(self, type_id):
+ """Queries Apache Ignite information about type by type id."""
+ self.write_int(14) # Message length
+ self.write_short(3002) # Operation code
+ self.write_long(0) # Request ID
+ self.write_int(type_id) # Type ID
+
+ self.read_int() # Result length
+ self.read_long() # Request id
+ status = self.read_int()
+
+ if status != 0:
+ err_msg = self._parse_string()
+ if err_msg is None:
+ raise RuntimeError("Get Binary Type Error [status=%d, message='%s']" %
+ (status, err_msg))
+ else:
+ raise RuntimeError("Get Binary Type Error [status=%d]" % status)
+
+ binary_type_exists = self.read_byte()
+
+ if binary_type_exists == 0:
+ raise RuntimeError("Binary type not found [type_id=%d] " % type_id)
+
+ binary_type_id = self.read_int()
+ binary_type_name = self._parse_string()
+ self._parse_string() # Affinity field name
+
+ fields = []
+ for _ in range(self.read_int()):
+ field_name = self._parse_string()
+ field_type_id = self.read_int()
+ field_id = self.read_int()
+
+ field = BinaryField(field_name, field_type_id, field_id)
+ fields.append(field)
+
+ is_enum = self.read_byte()
+ if is_enum == 1:
+ raise RuntimeError("Enum fields are not supported yet")
+
+ schema_cnt = self.read_int()
+ for _ in range(schema_cnt):
+ self.read_int() # Schema id
+ field_cnt = self.read_int()
+ self.skip(field_cnt * 4)
+
+ return BinaryType(binary_type_id, binary_type_name, fields)
+
+ def _parse_string(self):
+ """Parses string."""
+ header = self.read_byte()
+ if header == 9:
+ length = self.read_int()
+ return self.read_data(length).decode("utf-8")
+ if header == 101:
+ return None
+ raise RuntimeError(
+ "Unknown binary type when expected string [type_id=%d]" % header)
+
+
+class IgniteDataset(dataset_ops.Dataset):
+ """Apache Ignite is a memory-centric distributed database, 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.
+ """
+
+ def __init__(self,
+ cache_name,
+ host="localhost",
+ port=10800,
+ local=False,
+ part=-1,
+ page_size=100,
+ username=None,
+ password=None,
+ certfile=None,
+ keyfile=None,
+ cert_password=None):
+ """Create a IgniteDataset.
+
+ Args:
+ cache_name: Cache name to be used as datasource.
+ host: Apache Ignite Thin Client host to be connected.
+ port: Apache Ignite Thin Client port to be connected.
+ local: Local flag that defines to query only local data.
+ part: Number of partitions to be queried.
+ page_size: Apache Ignite Thin Client page size.
+ username: Apache Ignite Thin Client authentication username.
+ password: Apache Ignite Thin Client authentication password.
+ certfile: File in PEM format containing the certificate as well as any
+ number of CA certificates needed to establish the certificate's
+ authenticity.
+ keyfile: File containing the private key (otherwise the private key will
+ be taken from certfile as well).
+ cert_password: Password to be used if the private key is encrypted and a
+ password is necessary.
+ """
+ super(IgniteDataset, self).__init__()
+
+ with IgniteClient(host, port, username, password, certfile, keyfile,
+ cert_password) as client:
+ client.handshake()
+ self.cache_type = client.get_cache_type(cache_name)
+
+ self.cache_name = ops.convert_to_tensor(
+ cache_name, dtype=dtypes.string, name="cache_name")
+ self.host = ops.convert_to_tensor(host, dtype=dtypes.string, name="host")
+ self.port = ops.convert_to_tensor(port, dtype=dtypes.int32, name="port")
+ self.local = ops.convert_to_tensor(local, dtype=dtypes.bool, name="local")
+ self.part = ops.convert_to_tensor(part, dtype=dtypes.int32, name="part")
+ self.page_size = ops.convert_to_tensor(
+ page_size, dtype=dtypes.int32, name="page_size")
+ self.schema = ops.convert_to_tensor(
+ self.cache_type.to_flat(), dtype=dtypes.int32, name="schema")
+ self.permutation = ops.convert_to_tensor(
+ self.cache_type.to_permutation(),
+ dtype=dtypes.int32,
+ name="permutation")
+
+ def _as_variant_tensor(self):
+ return gen_dataset_ops.ignite_dataset(self.cache_name, self.host, self.port,
+ self.local, self.part, self.page_size,
+ 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()
diff --git a/tensorflow/contrib/data/python/ops/contrib_op_loader.py b/tensorflow/contrib/ignite/python/ops/ignite_op_loader.py
similarity index 94%
rename from tensorflow/contrib/data/python/ops/contrib_op_loader.py
rename to tensorflow/contrib/ignite/python/ops/ignite_op_loader.py
index 8f495a9..c9af738 100644
--- a/tensorflow/contrib/data/python/ops/contrib_op_loader.py
+++ b/tensorflow/contrib/ignite/python/ops/ignite_op_loader.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Python helper for loading contrib ops and kernels."""
+"""Python helper for loading Ignite ops and kernels."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/tensorflow/contrib/data/python/ops/contrib_op_loader.py b/tensorflow/contrib/ignite/python/tests/bin/start-plain.sh
old mode 100644
new mode 100755
similarity index 64%
copy from tensorflow/contrib/data/python/ops/contrib_op_loader.py
copy to tensorflow/contrib/ignite/python/tests/bin/start-plain.sh
index 8f495a9..f4607ce
--- a/tensorflow/contrib/data/python/ops/contrib_op_loader.py
+++ b/tensorflow/contrib/ignite/python/tests/bin/start-plain.sh
@@ -1,3 +1,4 @@
+#!/usr/bin/env bash
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -12,13 +13,12 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Python helper for loading contrib ops and kernels."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-from tensorflow.contrib.util import loader
-from tensorflow.python.platform import resource_loader
+nohup apache-ignite-fabric/bin/ignite.sh /data/config/ignite-config-plain.xml &
+sleep 5 # Wait Apache Ignite to be started
-_dataset_ops = loader.load_op_library(
- resource_loader.get_path_to_datafile("../../_dataset_ops.so"))
+./apache-ignite-fabric/bin/sqlline.sh \
+-u "jdbc:ignite:thin://127.0.0.1/" \
+--run=/data/sql/init.sql
+
+tail -f nohup.out
diff --git a/tensorflow/contrib/ignite/python/tests/config/ignite-config-plain.xml b/tensorflow/contrib/ignite/python/tests/config/ignite-config-plain.xml
new file mode 100644
index 0000000..d900174
--- /dev/null
+++ b/tensorflow/contrib/ignite/python/tests/config/ignite-config-plain.xml
@@ -0,0 +1,39 @@
+<?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.
+-->
+
+<beans xmlns="http://www.springframework.org/schema/beans"
+ xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+ xmlns:util="http://www.springframework.org/schema/util"
+ xsi:schemaLocation="http://www.springframework.org/schema/beans
+ http://www.springframework.org/schema/beans/spring-beans.xsd
+ http://www.springframework.org/schema/util
+ http://www.springframework.org/schema/util/spring-util.xsd">
+
+ <bean class="org.apache.ignite.configuration.IgniteConfiguration">
+ <property name="discoverySpi">
+ <bean class="org.apache.ignite.spi.discovery.tcp.TcpDiscoverySpi">
+ <property name="ipFinder">
+ <bean class="org.apache.ignite.spi.discovery.tcp.ipfinder.vm.TcpDiscoveryVmIpFinder">
+ <property name="addresses">
+ <list>
+ <value>127.0.0.1</value>
+ </list>
+ </property>
+ </bean>
+ </property>
+ </bean>
+ </property>
+ </bean>
+
+</beans>
diff --git a/tensorflow/contrib/ignite/python/tests/ignite_dataset_test.py b/tensorflow/contrib/ignite/python/tests/ignite_dataset_test.py
new file mode 100644
index 0000000..1856a4f
--- /dev/null
+++ b/tensorflow/contrib/ignite/python/tests/ignite_dataset_test.py
@@ -0,0 +1,118 @@
+# 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 IgniteDataset."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+
+from tensorflow.contrib.ignite import IgniteDataset
+from tensorflow.python.client import session
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.platform import test
+
+
+class IgniteDatasetTest(test.TestCase):
+ """The Apache Ignite servers have to setup before the test and tear down
+
+ after the test manually. The docker engine has to be installed.
+
+ To setup Apache Ignite servers:
+ $ bash start_ignite.sh
+
+ To tear down Apache Ignite servers:
+ $ bash stop_ignite.sh
+ """
+
+ def test_ignite_dataset_with_plain_client(self):
+ """Test Ignite Dataset with plain client.
+
+ """
+ self._clear_env()
+ ds = IgniteDataset(cache_name="SQL_PUBLIC_TEST_CACHE", port=42300)
+ self._check_dataset(ds)
+
+ def test_ignite_dataset_with_ssl_client(self):
+ """Test Ignite Dataset with ssl client.
+
+ """
+ self._clear_env()
+ os.environ["IGNITE_DATASET_CERTFILE"] = os.path.dirname(
+ os.path.realpath(__file__)) + "/keystore/client.pem"
+ os.environ["IGNITE_DATASET_CERT_PASSWORD"] = "123456"
+
+ ds = IgniteDataset(
+ cache_name="SQL_PUBLIC_TEST_CACHE",
+ port=42301,
+ certfile=os.environ["IGNITE_DATASET_CERTFILE"],
+ cert_password=os.environ["IGNITE_DATASET_CERT_PASSWORD"])
+ self._check_dataset(ds)
+
+ def test_ignite_dataset_with_ssl_client_and_auth(self):
+ """Test Ignite Dataset with ssl client and authentication.
+
+ """
+ self._clear_env()
+ os.environ["IGNITE_DATASET_USERNAME"] = "ignite"
+ os.environ["IGNITE_DATASET_PASSWORD"] = "ignite"
+ os.environ["IGNITE_DATASET_CERTFILE"] = os.path.dirname(
+ os.path.realpath(__file__)) + "/keystore/client.pem"
+ os.environ["IGNITE_DATASET_CERT_PASSWORD"] = "123456"
+
+ ds = IgniteDataset(
+ cache_name="SQL_PUBLIC_TEST_CACHE",
+ port=42302,
+ certfile=os.environ["IGNITE_DATASET_CERTFILE"],
+ cert_password=os.environ["IGNITE_DATASET_CERT_PASSWORD"],
+ username=os.environ["IGNITE_DATASET_USERNAME"],
+ password=os.environ["IGNITE_DATASET_PASSWORD"])
+ self._check_dataset(ds)
+
+ def _clear_env(self):
+ """Clears environment variables used by Ignite Dataset.
+
+ """
+ if "IGNITE_DATASET_USERNAME" in os.environ:
+ del os.environ["IGNITE_DATASET_USERNAME"]
+ if "IGNITE_DATASET_PASSWORD" in os.environ:
+ del os.environ["IGNITE_DATASET_PASSWORD"]
+ if "IGNITE_DATASET_CERTFILE" in os.environ:
+ del os.environ["IGNITE_DATASET_CERTFILE"]
+ if "IGNITE_DATASET_CERT_PASSWORD" in os.environ:
+ del os.environ["IGNITE_DATASET_CERT_PASSWORD"]
+
+ def _check_dataset(self, dataset):
+ """Checks that dataset provides correct data."""
+ self.assertEqual(dtypes.int64, dataset.output_types["key"])
+ self.assertEqual(dtypes.string, dataset.output_types["val"]["NAME"])
+ self.assertEqual(dtypes.int64, dataset.output_types["val"]["VAL"])
+
+ it = dataset.make_one_shot_iterator()
+ ne = it.get_next()
+
+ with session.Session() as sess:
+ rows = [sess.run(ne), sess.run(ne), sess.run(ne)]
+ with self.assertRaises(errors.OutOfRangeError):
+ sess.run(ne)
+
+ self.assertEqual({"key": 1, "val": {"NAME": b"TEST1", "VAL": 42}}, rows[0])
+ self.assertEqual({"key": 2, "val": {"NAME": b"TEST2", "VAL": 43}}, rows[1])
+ self.assertEqual({"key": 3, "val": {"NAME": b"TEST3", "VAL": 44}}, rows[2])
+
+
+if __name__ == "__main__":
+ test.main()
diff --git a/tensorflow/contrib/ignite/python/tests/sql/init.sql b/tensorflow/contrib/ignite/python/tests/sql/init.sql
new file mode 100644
index 0000000..5a192ae
--- /dev/null
+++ b/tensorflow/contrib/ignite/python/tests/sql/init.sql
@@ -0,0 +1,20 @@
+-- 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.
+-- ==============================================================================
+
+CREATE TABLE TEST_CACHE (ID LONG PRIMARY KEY, NAME VARCHAR, VAL LONG);
+
+INSERT INTO TEST_CACHE VALUES (1, 'TEST1', 42);
+INSERT INTO TEST_CACHE VALUES (2, 'TEST2', 43);
+INSERT INTO TEST_CACHE VALUES (3, 'TEST3', 44);
diff --git a/tensorflow/contrib/data/python/ops/contrib_op_loader.py b/tensorflow/contrib/ignite/python/tests/start_ignite.sh
old mode 100644
new mode 100755
similarity index 64%
copy from tensorflow/contrib/data/python/ops/contrib_op_loader.py
copy to tensorflow/contrib/ignite/python/tests/start_ignite.sh
index 8f495a9..a67bd44
--- a/tensorflow/contrib/data/python/ops/contrib_op_loader.py
+++ b/tensorflow/contrib/ignite/python/tests/start_ignite.sh
@@ -1,3 +1,4 @@
+#!/usr/bin/env bash
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -12,13 +13,10 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Python helper for loading contrib ops and kernels."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-from tensorflow.contrib.util import loader
-from tensorflow.python.platform import resource_loader
+IGNITE_VERSION=2.6.0
+SCRIPT_PATH="$( cd "$(dirname "$0")" ; pwd -P )"
-_dataset_ops = loader.load_op_library(
- resource_loader.get_path_to_datafile("../../_dataset_ops.so"))
+# Start Apache Ignite with plain client listener.
+docker run -itd --name ignite-plain -p 42300:10800 \
+-v ${SCRIPT_PATH}:/data apacheignite/ignite:${IGNITE_VERSION} /data/bin/start-plain.sh
diff --git a/tensorflow/contrib/data/python/ops/contrib_op_loader.py b/tensorflow/contrib/ignite/python/tests/stop_ignite.sh
old mode 100644
new mode 100755
similarity index 64%
copy from tensorflow/contrib/data/python/ops/contrib_op_loader.py
copy to tensorflow/contrib/ignite/python/tests/stop_ignite.sh
index 8f495a9..8f03dbd
--- a/tensorflow/contrib/data/python/ops/contrib_op_loader.py
+++ b/tensorflow/contrib/ignite/python/tests/stop_ignite.sh
@@ -1,3 +1,4 @@
+#!/usr/bin/env bash
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -12,13 +13,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Python helper for loading contrib ops and kernels."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-from tensorflow.contrib.util import loader
-from tensorflow.python.platform import resource_loader
-
-_dataset_ops = loader.load_op_library(
- resource_loader.get_path_to_datafile("../../_dataset_ops.so"))
+docker rm -f ignite-plain
+docker rm -f ignite-ssl
+docker rm -f ignite-ssl-auth
diff --git a/tensorflow/contrib/image/kernels/image_ops.cc b/tensorflow/contrib/image/kernels/image_ops.cc
index 370a8ca..788bf04 100644
--- a/tensorflow/contrib/image/kernels/image_ops.cc
+++ b/tensorflow/contrib/image/kernels/image_ops.cc
@@ -156,6 +156,7 @@
TF_CALL_uint8(DECLARE_FUNCTOR);
TF_CALL_int32(DECLARE_FUNCTOR);
TF_CALL_int64(DECLARE_FUNCTOR);
+TF_CALL_half(DECLARE_FUNCTOR);
TF_CALL_float(DECLARE_FUNCTOR);
TF_CALL_double(DECLARE_FUNCTOR);
@@ -175,6 +176,7 @@
TF_CALL_uint8(REGISTER);
TF_CALL_int32(REGISTER);
TF_CALL_int64(REGISTER);
+TF_CALL_half(REGISTER);
TF_CALL_float(REGISTER);
TF_CALL_double(REGISTER);
diff --git a/tensorflow/contrib/image/kernels/image_ops.h b/tensorflow/contrib/image/kernels/image_ops.h
index 6b63eed..7fac774 100644
--- a/tensorflow/contrib/image/kernels/image_ops.h
+++ b/tensorflow/contrib/image/kernels/image_ops.h
@@ -71,14 +71,7 @@
(transform[3] * output_x + transform[4] * output_y + transform[5]) /
projection;
- // TODO(ringwalt): Add a fill value input.
-#if (defined __CUDA_ARCH__) && (CUDART_VERSION < 8000)
- // On CUDA versions previous to 8.0, only __shared__ variables
- // could be declared as static in the device code.
const T fill_value = T(0);
-#else
- static const T fill_value = T(0);
-#endif
switch (interpolation_) {
case INTERPOLATION_NEAREST:
// Switch the order of x and y again for indexing into the image.
diff --git a/tensorflow/contrib/image/kernels/image_ops_gpu.cu.cc b/tensorflow/contrib/image/kernels/image_ops_gpu.cu.cc
index 8743a5f..36b9a23 100644
--- a/tensorflow/contrib/image/kernels/image_ops_gpu.cu.cc
+++ b/tensorflow/contrib/image/kernels/image_ops_gpu.cu.cc
@@ -32,6 +32,7 @@
template class FillProjectiveTransform<GPUDevice, uint8>;
template class FillProjectiveTransform<GPUDevice, int32>;
template class FillProjectiveTransform<GPUDevice, int64>;
+template class FillProjectiveTransform<GPUDevice, Eigen::half>;
template class FillProjectiveTransform<GPUDevice, float>;
template class FillProjectiveTransform<GPUDevice, double>;
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 376c075..4997c31 100644
--- a/tensorflow/contrib/image/python/kernel_tests/image_ops_test.py
+++ b/tensorflow/contrib/image/python/kernel_tests/image_ops_test.py
@@ -272,6 +272,15 @@
with self.cached_session():
self.assertAllEqual([[[[1], [0]], [[0], [1]]]], result.eval())
+ def test_transform_data_types(self):
+ for dtype in _DTYPES:
+ image = constant_op.constant([[1, 2], [3, 4]], dtype=dtype)
+ value = image_ops.transform(image, [1] * 8)
+ with self.test_session(use_gpu=True):
+ self.assertAllEqual(
+ value.eval(),
+ np.array([[4, 4], [4, 4]]).astype(dtype.as_numpy_dtype()))
+
class BipartiteMatchTest(test_util.TensorFlowTestCase):
diff --git a/tensorflow/contrib/layers/python/layers/embedding_ops.py b/tensorflow/contrib/layers/python/layers/embedding_ops.py
index 17ee8c0..60e1d85 100644
--- a/tensorflow/contrib/layers/python/layers/embedding_ops.py
+++ b/tensorflow/contrib/layers/python/layers/embedding_ops.py
@@ -112,11 +112,9 @@
dtype = sparse_weights.dtype if sparse_weights is not None else None
if isinstance(embedding_weights, variables.PartitionedVariable):
embedding_weights = list(embedding_weights)
- if not isinstance(embedding_weights[0],
- resource_variable_ops.ResourceVariable):
- embedding_weights = [
- ops.convert_to_tensor(w, dtype=dtype) for w in embedding_weights
- ]
+ embedding_weights = [
+ ops.convert_to_tensor(w, dtype=dtype) for w in embedding_weights
+ ]
contrib_tensor_util.assert_same_float_dtype(embedding_weights +
[sparse_weights])
diff --git a/tensorflow/contrib/layers/python/layers/layers_test.py b/tensorflow/contrib/layers/python/layers/layers_test.py
index 85af9de..3b7ae72 100644
--- a/tensorflow/contrib/layers/python/layers/layers_test.py
+++ b/tensorflow/contrib/layers/python/layers/layers_test.py
@@ -2360,7 +2360,7 @@
batch_size * height * width, expected_var)
images = constant_op.constant(
image_values, shape=image_shape, dtype=dtypes.float32)
- is_training = variables_lib.Variable(True)
+ is_training = variables_lib.VariableV1(True)
output = _layers.batch_norm(
images,
decay=0.1,
@@ -2507,7 +2507,7 @@
batch_size * height * width, expected_var)
images = constant_op.constant(
image_values, shape=image_shape, dtype=dtypes.float32)
- is_training = variables_lib.Variable(True)
+ is_training = variables_lib.VariableV1(True)
output = _layers.batch_norm(
images,
decay=0.1,
diff --git a/tensorflow/contrib/learn/python/learn/graph_actions_test.py b/tensorflow/contrib/learn/python/learn/graph_actions_test.py
index 33180b7..a160cb5 100644
--- a/tensorflow/contrib/learn/python/learn/graph_actions_test.py
+++ b/tensorflow/contrib/learn/python/learn/graph_actions_test.py
@@ -162,9 +162,9 @@
Tuple of 3 `Tensor` objects, 2 input and 1 output.
"""
variables_lib.create_global_step()
- in0 = variables.Variable(1.0)
+ in0 = variables.VariableV1(1.0)
in1 = variables_lib.local_variable(2.0)
- fake_table = variables.Variable(
+ fake_table = variables.VariableV1(
3.0,
trainable=False,
collections=['fake_tables'],
@@ -312,8 +312,8 @@
def test_evaluate_ready_for_local_init(self):
with ops.Graph().as_default() as g, self.session(g):
variables_lib.create_global_step()
- v = variables.Variable(1.0)
- variables.Variable(
+ v = variables.VariableV1(1.0)
+ variables.VariableV1(
v + 1, collections=[ops.GraphKeys.LOCAL_VARIABLES], trainable=False)
ready_for_local_init_op = variables.report_uninitialized_variables(
variables.global_variables())
@@ -456,9 +456,9 @@
Tuple of 3 `Tensor` objects, 2 input and 1 output.
"""
variables_lib.create_global_step()
- in0 = variables.Variable(1.0)
+ in0 = variables.VariableV1(1.0)
in1 = variables_lib.local_variable(2.0)
- fake_table = variables.Variable(
+ fake_table = variables.VariableV1(
3.0,
trainable=False,
collections=['fake_tables'],
diff --git a/tensorflow/contrib/learn/python/learn/monitors_test.py b/tensorflow/contrib/learn/python/learn/monitors_test.py
index 83e48a3..d4a7169 100644
--- a/tensorflow/contrib/learn/python/learn/monitors_test.py
+++ b/tensorflow/contrib/learn/python/learn/monitors_test.py
@@ -247,7 +247,7 @@
def test_logging_trainable(self):
with ops.Graph().as_default() as g, self.session(g):
- var = variables.Variable(constant_op.constant(42.0), name='foo')
+ var = variables.VariableV1(constant_op.constant(42.0), name='foo')
var.initializer.run()
cof = constant_op.constant(1.0)
loss = math_ops.subtract(
@@ -261,7 +261,7 @@
with ops.Graph().as_default() as g, self.session(g):
log_dir = 'log/dir'
summary_writer = testing.FakeSummaryWriter(log_dir, g)
- var = variables.Variable(0.0)
+ var = variables.VariableV1(0.0)
var.initializer.run()
tensor = state_ops.assign_add(var, 1.0)
summary_op = summary.scalar('my_summary', tensor)
@@ -526,8 +526,8 @@
monitor0 = learn.monitors.GraphDump()
monitor1 = learn.monitors.GraphDump()
with ops.Graph().as_default() as g, self.session(g):
- const_var = variables.Variable(42.0, name='my_const')
- counter_var = variables.Variable(0.0, name='my_counter')
+ const_var = variables.VariableV1(42.0, name='my_const')
+ counter_var = variables.VariableV1(0.0, name='my_counter')
assign_add = state_ops.assign_add(counter_var, 1.0, name='my_assign_add')
variables.global_variables_initializer().run()
@@ -569,7 +569,7 @@
monitor = learn.monitors.CaptureVariable(
var_name='my_assign_add:0', every_n=8, first_n=2)
with ops.Graph().as_default() as g, self.session(g):
- var = variables.Variable(0.0, name='my_var')
+ var = variables.VariableV1(0.0, name='my_var')
var.initializer.run()
state_ops.assign_add(var, 1.0, name='my_assign_add')
self._run_monitor(monitor, num_epochs=3, num_steps_per_epoch=10)
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 9ecf023..8466dc3 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
@@ -125,7 +125,7 @@
],
example_ids=[str(i) for i in range(num_examples)])
- weights = variables_lib.Variable(
+ weights = variables_lib.VariableV1(
array_ops.zeros([dim], dtype=dtypes.float32))
variables_dict = dict(
sparse_features_weights=[weights],
@@ -184,7 +184,7 @@
dense_tensors.append(dense_tensor)
# Add variables of shape [feature_column_dimension].
dense_weights.append(
- variables_lib.Variable(
+ variables_lib.VariableV1(
array_ops.zeros(
[dense_tensor.get_shape().as_list()[1]], dtype=dtypes.float32)))
@@ -341,7 +341,7 @@
examples = make_example_dict(example_protos, example_weights)
# Explicitly make age a [1]-shaped Variable (which cannot be
# partitioned), while making gender a PartitionedVariable.
- age_weights = variables_lib.Variable(
+ age_weights = variables_lib.VariableV1(
array_ops.zeros([1], dtype=dtypes.float32))
with variable_scope.variable_scope(
name_or_scope=('variables/shard_{}'.format(num_shards)
@@ -801,7 +801,7 @@
labels=[1.0, 0.0])
# Replace with a variable of size 1 instead of 2.
variables['dense_features_weights'] = [
- variables_lib.Variable(array_ops.zeros(
+ variables_lib.VariableV1(array_ops.zeros(
[1], dtype=dtypes.float32))
]
options = dict(
diff --git a/tensorflow/contrib/lite/BUILD b/tensorflow/contrib/lite/BUILD
index f320b53..f3ebe3b 100644
--- a/tensorflow/contrib/lite/BUILD
+++ b/tensorflow/contrib/lite/BUILD
@@ -26,6 +26,14 @@
},
)
+# Enables inclusion of TensorFlow kernels via the TF Lite Flex delegate.
+# WARNING: This build flag is experimental and subject to change.
+config_setting(
+ name = "with_tflite_flex",
+ define_values = {"with_tflite_flex": "true"},
+ visibility = ["//visibility:public"],
+)
+
cc_library(
name = "schema_fbs_version",
hdrs = ["version.h"],
@@ -157,6 +165,10 @@
"stderr_reporter.h",
],
copts = tflite_copts(),
+ defines = select({
+ ":with_tflite_flex": ["TFLITE_FLEX"],
+ "//conditions:default": [],
+ }),
linkopts = [
] + select({
"//tensorflow:android": [
@@ -180,7 +192,12 @@
"//tensorflow/contrib/lite/nnapi:nnapi_lib",
"//tensorflow/contrib/lite/profiling:profiler",
"//tensorflow/contrib/lite/schema:schema_fbs",
- ],
+ ] + select({
+ ":with_tflite_flex": [
+ "//tensorflow/contrib/lite/delegates/flex:delegate",
+ ],
+ "//conditions:default": [],
+ }),
)
cc_library(
diff --git a/tensorflow/contrib/lite/build_def.bzl b/tensorflow/contrib/lite/build_def.bzl
index fc4d9b4..7ef26de 100644
--- a/tensorflow/contrib/lite/build_def.bzl
+++ b/tensorflow/contrib/lite/build_def.bzl
@@ -301,7 +301,7 @@
"""Returns a list of conversion modes."""
# TODO(nupurgarg): Add "pb2lite" when it's in open source. b/113614050.
- return ["toco-extended", ""]
+ return ["toco-flex", ""]
def generated_test_models_all():
"""Generates a list of all tests with the different converters.
@@ -335,7 +335,7 @@
# TODO(nupurgarg): Comment in when pb2lite is in open source. b/113614050.
# if conversion_mode == "pb2lite":
# toco = "//tensorflow/contrib/lite/experimental/pb2lite:pb2lite"
- flags = "--ignore_toco_errors --run_with_extended"
+ flags = "--ignore_toco_errors --run_with_flex"
kwargs["tags"].append("skip_already_failing")
kwargs["tags"].append("no_oss")
kwargs["tags"].append("notap")
@@ -391,3 +391,41 @@
(tool, model, out, tflite_path[2:]),
tools = [tool],
)
+
+def gen_full_model_test(conversion_modes, models, data, test_suite_tag):
+ """Generates Python test targets for testing TFLite models.
+
+ Args:
+ conversion_modes: List of conversion modes to test the models on.
+ models: List of models to test.
+ data: List of BUILD targets linking the data.
+ test_suite_tag: Tag identifying the model test suite.
+ """
+ options = [
+ (conversion_mode, model)
+ for model in models
+ for conversion_mode in conversion_modes
+ ]
+
+ for conversion_mode, model_name in options:
+ native.py_test(
+ name = "model_coverage_test_%s_%s" % (model_name, conversion_mode.lower()),
+ srcs = ["model_coverage_test.py"],
+ main = "model_coverage_test.py",
+ args = [
+ "--model_name=%s" % model_name,
+ "--converter_mode=%s" % conversion_mode,
+ ],
+ data = data,
+ srcs_version = "PY2AND3",
+ tags = [
+ "no_oss",
+ "no_windows",
+ "notap",
+ ] + [test_suite_tag],
+ deps = [
+ "//tensorflow/contrib/lite/testing:model_coverage_lib",
+ "//tensorflow/contrib/lite/python:lite",
+ "//tensorflow/python:client_testlib",
+ ],
+ )
diff --git a/tensorflow/contrib/lite/delegates/eager/BUILD b/tensorflow/contrib/lite/delegates/flex/BUILD
similarity index 95%
rename from tensorflow/contrib/lite/delegates/eager/BUILD
rename to tensorflow/contrib/lite/delegates/flex/BUILD
index bf5d918..9dd3895 100644
--- a/tensorflow/contrib/lite/delegates/eager/BUILD
+++ b/tensorflow/contrib/lite/delegates/flex/BUILD
@@ -20,7 +20,7 @@
"//tensorflow/contrib/lite:kernel_api",
] + select({
"//tensorflow:android": [
- "//tensorflow/core:android_tensorflow_lib_lite_no_runtime",
+ "//tensorflow/core:android_tensorflow_lib_lite",
],
"//conditions:default": [
"//tensorflow/core:framework",
@@ -60,7 +60,7 @@
"//tensorflow/contrib/lite:util",
] + select({
"//tensorflow:android": [
- "//tensorflow/core:android_tensorflow_lib_lite_no_runtime",
+ "//tensorflow/core:android_tensorflow_lib_lite",
],
"//conditions:default": [
"//tensorflow/core:lib",
@@ -178,7 +178,7 @@
"//tensorflow/contrib/lite:kernel_api",
] + select({
"//tensorflow:android": [
- "//tensorflow/core:android_tensorflow_lib_lite_no_runtime",
+ "//tensorflow/core:android_tensorflow_lib",
],
"//conditions:default": [
"//tensorflow/core:lib",
diff --git a/tensorflow/contrib/lite/delegates/eager/buffer_map.cc b/tensorflow/contrib/lite/delegates/flex/buffer_map.cc
similarity index 95%
rename from tensorflow/contrib/lite/delegates/eager/buffer_map.cc
rename to tensorflow/contrib/lite/delegates/flex/buffer_map.cc
index e5a19c3..63e3919 100644
--- a/tensorflow/contrib/lite/delegates/eager/buffer_map.cc
+++ b/tensorflow/contrib/lite/delegates/flex/buffer_map.cc
@@ -12,15 +12,15 @@
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
-#include "tensorflow/contrib/lite/delegates/eager/buffer_map.h"
+#include "tensorflow/contrib/lite/delegates/flex/buffer_map.h"
#include "tensorflow/c/c_api_internal.h"
-#include "tensorflow/contrib/lite/delegates/eager/util.h"
+#include "tensorflow/contrib/lite/delegates/flex/util.h"
#include "tensorflow/core/framework/allocation_description.pb.h"
#include "tensorflow/core/framework/log_memory.h"
namespace tflite {
-namespace eager {
+namespace flex {
namespace {
// A tensor buffer that is allocated, deallocated and populated by TF Lite.
class TfLiteTensorBuffer : public tensorflow::TensorBuffer {
@@ -107,5 +107,5 @@
id_to_tensor_[tensor_index] = std::move(tensor);
}
-} // namespace eager
+} // namespace flex
} // namespace tflite
diff --git a/tensorflow/contrib/lite/delegates/eager/buffer_map.h b/tensorflow/contrib/lite/delegates/flex/buffer_map.h
similarity index 86%
rename from tensorflow/contrib/lite/delegates/eager/buffer_map.h
rename to tensorflow/contrib/lite/delegates/flex/buffer_map.h
index aaaa045..4ce8865 100644
--- a/tensorflow/contrib/lite/delegates/eager/buffer_map.h
+++ b/tensorflow/contrib/lite/delegates/flex/buffer_map.h
@@ -12,8 +12,8 @@
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
-#ifndef TENSORFLOW_CONTRIB_LITE_DELEGATES_EAGER_BUFFER_MAP_H_
-#define TENSORFLOW_CONTRIB_LITE_DELEGATES_EAGER_BUFFER_MAP_H_
+#ifndef TENSORFLOW_CONTRIB_LITE_DELEGATES_FLEX_BUFFER_MAP_H_
+#define TENSORFLOW_CONTRIB_LITE_DELEGATES_FLEX_BUFFER_MAP_H_
#include <map>
@@ -21,12 +21,12 @@
#include "tensorflow/core/framework/tensor.h"
namespace tflite {
-namespace eager {
+namespace flex {
// Maps a TF Lite tensor index into a TensorFlow tensor.
//
// The TF Lite interpreter assigns integer indices to each of its tensors, but
-// the Eager delegate deals in terms of TensorFlow tensors. This class maps
+// the Flex delegate deals in terms of TensorFlow tensors. This class maps
// from indices to tensors and allows the creation of new tensors to be
// associated with a given index.
class BufferMap {
@@ -55,7 +55,7 @@
std::map<int, tensorflow::Tensor> id_to_tensor_;
};
-} // namespace eager
+} // namespace flex
} // namespace tflite
-#endif // TENSORFLOW_CONTRIB_LITE_DELEGATES_EAGER_BUFFER_MAP_H_
+#endif // TENSORFLOW_CONTRIB_LITE_DELEGATES_FLEX_BUFFER_MAP_H_
diff --git a/tensorflow/contrib/lite/delegates/eager/buffer_map_test.cc b/tensorflow/contrib/lite/delegates/flex/buffer_map_test.cc
similarity index 97%
rename from tensorflow/contrib/lite/delegates/eager/buffer_map_test.cc
rename to tensorflow/contrib/lite/delegates/flex/buffer_map_test.cc
index a046943..bb80e25 100644
--- a/tensorflow/contrib/lite/delegates/eager/buffer_map_test.cc
+++ b/tensorflow/contrib/lite/delegates/flex/buffer_map_test.cc
@@ -12,7 +12,7 @@
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
-#include "tensorflow/contrib/lite/delegates/eager/buffer_map.h"
+#include "tensorflow/contrib/lite/delegates/flex/buffer_map.h"
#include <gmock/gmock.h>
#include <gtest/gtest.h>
@@ -21,7 +21,7 @@
#include "tensorflow/contrib/lite/util.h"
namespace tflite {
-namespace eager {
+namespace flex {
namespace {
using ::testing::ElementsAre;
@@ -164,7 +164,7 @@
}
} // namespace
-} // namespace eager
+} // namespace flex
} // namespace tflite
int main(int argc, char** argv) {
diff --git a/tensorflow/contrib/lite/delegates/eager/delegate.cc b/tensorflow/contrib/lite/delegates/flex/delegate.cc
similarity index 76%
rename from tensorflow/contrib/lite/delegates/eager/delegate.cc
rename to tensorflow/contrib/lite/delegates/flex/delegate.cc
index 45fc158..ba065a8 100644
--- a/tensorflow/contrib/lite/delegates/eager/delegate.cc
+++ b/tensorflow/contrib/lite/delegates/flex/delegate.cc
@@ -12,19 +12,19 @@
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
-#include "tensorflow/contrib/lite/delegates/eager/delegate.h"
+#include "tensorflow/contrib/lite/delegates/flex/delegate.h"
#include <vector>
#include "tensorflow/contrib/lite/context_util.h"
-#include "tensorflow/contrib/lite/delegates/eager/buffer_map.h"
-#include "tensorflow/contrib/lite/delegates/eager/kernel.h"
-#include "tensorflow/contrib/lite/delegates/eager/util.h"
+#include "tensorflow/contrib/lite/delegates/flex/buffer_map.h"
+#include "tensorflow/contrib/lite/delegates/flex/kernel.h"
+#include "tensorflow/contrib/lite/delegates/flex/util.h"
#include "tensorflow/contrib/lite/util.h"
#include "tensorflow/core/lib/core/status.h"
namespace tflite {
-namespace eager {
+namespace flex {
namespace delegate {
TfLiteStatus Prepare(TfLiteContext* context, TfLiteDelegate* delegate) {
@@ -32,7 +32,7 @@
TfLiteIntArray* plan;
TF_LITE_ENSURE_STATUS(context->GetExecutionPlan(context, &plan));
- // Add all custom ops starting with "Eager" to list of supported nodes.
+ // Add all custom ops starting with "Flex" to list of supported nodes.
std::vector<int> supported_nodes;
for (int node_index : TfLiteIntArrayView(plan)) {
TfLiteNode* node;
@@ -40,7 +40,7 @@
TF_LITE_ENSURE_STATUS(context->GetNodeAndRegistration(
context, node_index, &node, ®istration));
- if (IsEagerOp(registration->custom_name)) {
+ if (IsFlexOp(registration->custom_name)) {
supported_nodes.push_back(node_index);
}
}
@@ -81,28 +81,28 @@
}
} // namespace delegate
-} // namespace eager
+} // namespace flex
-std::unique_ptr<EagerDelegate> EagerDelegate::Create() {
- std::unique_ptr<eager::DelegateData> delegate_data;
- if (!eager::DelegateData::Create(&delegate_data).ok()) {
+std::unique_ptr<FlexDelegate> FlexDelegate::Create() {
+ std::unique_ptr<flex::DelegateData> delegate_data;
+ if (!flex::DelegateData::Create(&delegate_data).ok()) {
fprintf(stderr, "Unable to initialize TensorFlow context.\n");
return nullptr;
}
- return std::unique_ptr<EagerDelegate>(
- new EagerDelegate(std::move(delegate_data)));
+ return std::unique_ptr<FlexDelegate>(
+ new FlexDelegate(std::move(delegate_data)));
}
-EagerDelegate::EagerDelegate(std::unique_ptr<eager::DelegateData> delegate_data)
+FlexDelegate::FlexDelegate(std::unique_ptr<flex::DelegateData> delegate_data)
: TfLiteDelegate{
/*data_=*/delegate_data.get(),
- /*nullptr,*/ &eager::delegate::Prepare,
- /*CopyFromBufferHandle=*/&eager::delegate::CopyFromBufferHandle,
+ /*nullptr,*/ &flex::delegate::Prepare,
+ /*CopyFromBufferHandle=*/&flex::delegate::CopyFromBufferHandle,
/*CopyToBufferHandle=*/nullptr,
/*FreeBufferHandle=*/nullptr},
delegate_data_(std::move(delegate_data)) {}
-EagerDelegate::~EagerDelegate() {}
+FlexDelegate::~FlexDelegate() {}
} // namespace tflite
diff --git a/tensorflow/contrib/lite/delegates/eager/delegate.h b/tensorflow/contrib/lite/delegates/flex/delegate.h
similarity index 64%
rename from tensorflow/contrib/lite/delegates/eager/delegate.h
rename to tensorflow/contrib/lite/delegates/flex/delegate.h
index 70f3c15..1017780 100644
--- a/tensorflow/contrib/lite/delegates/eager/delegate.h
+++ b/tensorflow/contrib/lite/delegates/flex/delegate.h
@@ -12,11 +12,11 @@
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
-#ifndef TENSORFLOW_CONTRIB_LITE_DELEGATES_EAGER_DELEGATE_H_
-#define TENSORFLOW_CONTRIB_LITE_DELEGATES_EAGER_DELEGATE_H_
+#ifndef TENSORFLOW_CONTRIB_LITE_DELEGATES_FLEX_DELEGATE_H_
+#define TENSORFLOW_CONTRIB_LITE_DELEGATES_FLEX_DELEGATE_H_
#include "tensorflow/contrib/lite/c/c_api_internal.h"
-#include "tensorflow/contrib/lite/delegates/eager/delegate_data.h"
+#include "tensorflow/contrib/lite/delegates/flex/delegate_data.h"
namespace tflite {
@@ -24,12 +24,12 @@
// Delegate that can be used to extract parts of a graph that are designed to be
// executed by TensorFlow's runtime via Eager.
//
-// The interpreter must be constructed after the EagerDelegate and destructed
-// before the EagerDelegate. This delegate may be used with multiple
+// The interpreter must be constructed after the FlexDelegate and destructed
+// before the FlexDelegate. This delegate may be used with multiple
// interpreters, but it is *not* thread-safe.
//
// Usage:
-// auto delegate = EagerDelegate::Create();
+// auto delegate = FlexDelegate::Create();
// ... build interpreter ...
//
// if (delegate) {
@@ -39,21 +39,21 @@
// ... run inference ...
// ... destroy interpreter ...
// ... destroy delegate ...
-class EagerDelegate : public TfLiteDelegate {
+class FlexDelegate : public TfLiteDelegate {
public:
// Creates a delegate that supports TF ops.
//
- // If the underyling TF Eager context creation fails, returns null.
- static std::unique_ptr<EagerDelegate> Create();
+ // If the underyling TF Flex context creation fails, returns null.
+ static std::unique_ptr<FlexDelegate> Create();
- ~EagerDelegate();
+ ~FlexDelegate();
private:
- explicit EagerDelegate(std::unique_ptr<eager::DelegateData> delegate_data);
+ explicit FlexDelegate(std::unique_ptr<flex::DelegateData> delegate_data);
- std::unique_ptr<eager::DelegateData> delegate_data_;
+ std::unique_ptr<flex::DelegateData> delegate_data_;
};
} // namespace tflite
-#endif // TENSORFLOW_CONTRIB_LITE_DELEGATES_EAGER_DELEGATE_H_
+#endif // TENSORFLOW_CONTRIB_LITE_DELEGATES_FLEX_DELEGATE_H_
diff --git a/tensorflow/contrib/lite/delegates/eager/delegate_data.cc b/tensorflow/contrib/lite/delegates/flex/delegate_data.cc
similarity index 94%
rename from tensorflow/contrib/lite/delegates/eager/delegate_data.cc
rename to tensorflow/contrib/lite/delegates/flex/delegate_data.cc
index 0fd5c97..8f985f7 100644
--- a/tensorflow/contrib/lite/delegates/eager/delegate_data.cc
+++ b/tensorflow/contrib/lite/delegates/flex/delegate_data.cc
@@ -12,13 +12,13 @@
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
-#include "tensorflow/contrib/lite/delegates/eager/delegate_data.h"
+#include "tensorflow/contrib/lite/delegates/flex/delegate_data.h"
#include "tensorflow/core/common_runtime/device_factory.h"
#include "tensorflow/core/lib/core/status.h"
namespace tflite {
-namespace eager {
+namespace flex {
tensorflow::Status DelegateData::Create(std::unique_ptr<DelegateData>* data) {
std::vector<tensorflow::Device*> devices;
@@ -43,5 +43,5 @@
DelegateData::~DelegateData() {}
-} // namespace eager
+} // namespace flex
} // namespace tflite
diff --git a/tensorflow/contrib/lite/delegates/eager/delegate_data.h b/tensorflow/contrib/lite/delegates/flex/delegate_data.h
similarity index 78%
rename from tensorflow/contrib/lite/delegates/eager/delegate_data.h
rename to tensorflow/contrib/lite/delegates/flex/delegate_data.h
index 772d26f..8d75f0b 100644
--- a/tensorflow/contrib/lite/delegates/eager/delegate_data.h
+++ b/tensorflow/contrib/lite/delegates/flex/delegate_data.h
@@ -12,16 +12,16 @@
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
-#ifndef TENSORFLOW_CONTRIB_LITE_DELEGATES_EAGER_DELEGATE_DATA_H_
-#define TENSORFLOW_CONTRIB_LITE_DELEGATES_EAGER_DELEGATE_DATA_H_
+#ifndef TENSORFLOW_CONTRIB_LITE_DELEGATES_FLEX_DELEGATE_DATA_H_
+#define TENSORFLOW_CONTRIB_LITE_DELEGATES_FLEX_DELEGATE_DATA_H_
-#include "tensorflow/contrib/lite/delegates/eager/buffer_map.h"
+#include "tensorflow/contrib/lite/delegates/flex/buffer_map.h"
#include "tensorflow/core/common_runtime/eager/context.h"
namespace tflite {
-namespace eager {
+namespace flex {
-// Data kept by the Eager delegate for the lifetime of an Interpreter.
+// Data kept by the Flex delegate for the lifetime of an Interpreter.
class DelegateData {
public:
// Create a new DelegateData, initialized with a newly-created EagerContext.
@@ -29,7 +29,7 @@
~DelegateData();
- // The EagerContext that is required for execution of Eager Ops.
+ // The EagerContext that is required for execution of Flex Ops.
tensorflow::EagerContext* GetEagerContext() { return eager_context_.get(); }
// Map from TF Lite tensor index to TensorFlow tensor for a given context.
@@ -46,7 +46,7 @@
std::unordered_map<const TfLiteContext*, BufferMap> buffer_map_;
};
-} // namespace eager
+} // namespace flex
} // namespace tflite
-#endif // TENSORFLOW_CONTRIB_LITE_DELEGATES_EAGER_DELEGATE_DATA_H_
+#endif // TENSORFLOW_CONTRIB_LITE_DELEGATES_FLEX_DELEGATE_DATA_H_
diff --git a/tensorflow/contrib/lite/delegates/eager/delegate_data_test.cc b/tensorflow/contrib/lite/delegates/flex/delegate_data_test.cc
similarity index 93%
rename from tensorflow/contrib/lite/delegates/eager/delegate_data_test.cc
rename to tensorflow/contrib/lite/delegates/flex/delegate_data_test.cc
index def0633..30b10f4 100644
--- a/tensorflow/contrib/lite/delegates/eager/delegate_data_test.cc
+++ b/tensorflow/contrib/lite/delegates/flex/delegate_data_test.cc
@@ -12,7 +12,7 @@
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
-#include "tensorflow/contrib/lite/delegates/eager/delegate_data.h"
+#include "tensorflow/contrib/lite/delegates/flex/delegate_data.h"
#include <gmock/gmock.h>
#include <gtest/gtest.h>
@@ -20,7 +20,7 @@
#include "tensorflow/contrib/lite/testing/util.h"
namespace tflite {
-namespace eager {
+namespace flex {
namespace {
TEST(DelegateDataTest, Basic) {
@@ -39,7 +39,7 @@
}
} // namespace
-} // namespace eager
+} // namespace flex
} // namespace tflite
int main(int argc, char** argv) {
diff --git a/tensorflow/contrib/lite/delegates/eager/delegate_test.cc b/tensorflow/contrib/lite/delegates/flex/delegate_test.cc
similarity index 95%
rename from tensorflow/contrib/lite/delegates/eager/delegate_test.cc
rename to tensorflow/contrib/lite/delegates/flex/delegate_test.cc
index 43ec5d5..1813952 100644
--- a/tensorflow/contrib/lite/delegates/eager/delegate_test.cc
+++ b/tensorflow/contrib/lite/delegates/flex/delegate_test.cc
@@ -12,23 +12,23 @@
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
-#include "tensorflow/contrib/lite/delegates/eager/delegate.h"
+#include "tensorflow/contrib/lite/delegates/flex/delegate.h"
#include <gmock/gmock.h>
#include <gtest/gtest.h>
-#include "tensorflow/contrib/lite/delegates/eager/test_util.h"
+#include "tensorflow/contrib/lite/delegates/flex/test_util.h"
namespace tflite {
-namespace eager {
+namespace flex {
namespace {
using ::testing::ContainsRegex;
using ::testing::ElementsAre;
-class DelegateTest : public testing::EagerModelTest {
+class DelegateTest : public testing::FlexModelTest {
public:
DelegateTest() {
- delegate_ = EagerDelegate::Create();
+ delegate_ = FlexDelegate::Create();
interpreter_.reset(new Interpreter(&error_reporter_));
}
@@ -46,7 +46,7 @@
}
private:
- std::unique_ptr<EagerDelegate> delegate_;
+ std::unique_ptr<FlexDelegate> delegate_;
};
TEST_F(DelegateTest, FullGraph) {
@@ -236,7 +236,7 @@
}
} // namespace
-} // namespace eager
+} // namespace flex
} // namespace tflite
int main(int argc, char** argv) {
diff --git a/tensorflow/contrib/lite/delegates/eager/kernel.cc b/tensorflow/contrib/lite/delegates/flex/kernel.cc
similarity index 91%
rename from tensorflow/contrib/lite/delegates/eager/kernel.cc
rename to tensorflow/contrib/lite/delegates/flex/kernel.cc
index 48a2f56..e4f1aea 100644
--- a/tensorflow/contrib/lite/delegates/eager/kernel.cc
+++ b/tensorflow/contrib/lite/delegates/flex/kernel.cc
@@ -12,14 +12,14 @@
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
-#include "tensorflow/contrib/lite/delegates/eager/kernel.h"
+#include "tensorflow/contrib/lite/delegates/flex/kernel.h"
#include "flatbuffers/flexbuffers.h" // TF:flatbuffers
#include "tensorflow/contrib/lite/builtin_ops.h"
#include "tensorflow/contrib/lite/c/c_api_internal.h"
#include "tensorflow/contrib/lite/context_util.h"
-#include "tensorflow/contrib/lite/delegates/eager/delegate_data.h"
-#include "tensorflow/contrib/lite/delegates/eager/util.h"
+#include "tensorflow/contrib/lite/delegates/flex/delegate_data.h"
+#include "tensorflow/contrib/lite/delegates/flex/util.h"
#include "tensorflow/contrib/lite/kernels/kernel_util.h"
#include "tensorflow/contrib/lite/string.h"
#include "tensorflow/core/common_runtime/eager/context.h"
@@ -28,10 +28,10 @@
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/node_def_util.h"
-// Note: this is part of TF Lite's Eager delegation code which is to be
+// Note: this is part of TF Lite's Flex delegation code which is to be
// completed soon.
-// This is the TF Lite op that is created by the eager delegate to handle
+// This is the TF Lite op that is created by the flex delegate to handle
// execution of a supported subgraph. The usual flow is that the delegate
// informs the interpreter of supported nodes in a graph, and each supported
// subgraph is replaced with one instance of this kernel.
@@ -46,7 +46,7 @@
// corresponding TensorFlow/Eager Op.
namespace tflite {
-namespace eager {
+namespace flex {
namespace kernel {
// Controls the lifetime of tensor handles in a vector.
@@ -72,11 +72,11 @@
// Executes the TensorFlow op given by 'op_name', with the attributes specified
// in 'nodedef'. Inputs and outputs are given as indices into the 'buffer_map'.
-tensorflow::Status ExecuteEagerOp(tensorflow::EagerContext* eager_context,
- BufferMap* buffer_map, const string& op_name,
- const tensorflow::NodeDef& nodedef,
- const std::vector<int>& inputs,
- const std::vector<int>& outputs) {
+tensorflow::Status ExecuteFlexOp(tensorflow::EagerContext* eager_context,
+ BufferMap* buffer_map, const string& op_name,
+ const tensorflow::NodeDef& nodedef,
+ const std::vector<int>& inputs,
+ const std::vector<int>& outputs) {
const tensorflow::AttrTypeMap* attr_types;
TF_RETURN_WITH_CONTEXT_IF_ERROR(
tensorflow::AttrTypeMapForOp(op_name.c_str(), &attr_types),
@@ -258,13 +258,13 @@
// Execute the TensorFlow Ops sequentially.
for (const auto& node_data : op_data->nodes) {
if (node_data.nodedef.op().empty()) {
- context->ReportError(context, "Invalid NodeDef in Eager op '%s'",
+ context->ReportError(context, "Invalid NodeDef in Flex op '%s'",
node_data.name.c_str());
return kTfLiteError;
}
auto status =
- ExecuteEagerOp(eager_context, buffer_map, node_data.name,
- node_data.nodedef, node_data.inputs, node_data.outputs);
+ ExecuteFlexOp(eager_context, buffer_map, node_data.name,
+ node_data.nodedef, node_data.inputs, node_data.outputs);
TF_LITE_ENSURE_OK(context, ConvertStatus(context, status));
}
@@ -295,5 +295,5 @@
return registration;
}
-} // namespace eager
+} // namespace flex
} // namespace tflite
diff --git a/tensorflow/contrib/lite/delegates/eager/kernel.h b/tensorflow/contrib/lite/delegates/flex/kernel.h
similarity index 79%
rename from tensorflow/contrib/lite/delegates/eager/kernel.h
rename to tensorflow/contrib/lite/delegates/flex/kernel.h
index 2478abc..ac9313a 100644
--- a/tensorflow/contrib/lite/delegates/eager/kernel.h
+++ b/tensorflow/contrib/lite/delegates/flex/kernel.h
@@ -12,23 +12,23 @@
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
-#ifndef TENSORFLOW_CONTRIB_LITE_DELEGATES_EAGER_KERNEL_H_
-#define TENSORFLOW_CONTRIB_LITE_DELEGATES_EAGER_KERNEL_H_
+#ifndef TENSORFLOW_CONTRIB_LITE_DELEGATES_FLEX_KERNEL_H_
+#define TENSORFLOW_CONTRIB_LITE_DELEGATES_FLEX_KERNEL_H_
#include "tensorflow/contrib/lite/c/c_api_internal.h"
namespace tflite {
-namespace eager {
+namespace flex {
// Return the registration object used to initialize and execute ops that will
// be delegated to TensorFlow's Eager runtime. This TF Lite op is created by
-// the eager delegate to handle execution of a supported subgraph. The usual
+// the flex delegate to handle execution of a supported subgraph. The usual
// flow is that the delegate informs the interpreter of supported nodes in a
// graph, and each supported subgraph is replaced with one instance of this
// kernel.
TfLiteRegistration GetKernel();
-} // namespace eager
+} // namespace flex
} // namespace tflite
-#endif // TENSORFLOW_CONTRIB_LITE_DELEGATES_EAGER_KERNEL_H_
+#endif // TENSORFLOW_CONTRIB_LITE_DELEGATES_FLEX_KERNEL_H_
diff --git a/tensorflow/contrib/lite/delegates/eager/kernel_test.cc b/tensorflow/contrib/lite/delegates/flex/kernel_test.cc
similarity index 94%
rename from tensorflow/contrib/lite/delegates/eager/kernel_test.cc
rename to tensorflow/contrib/lite/delegates/flex/kernel_test.cc
index 66f2226..94a6f8b 100644
--- a/tensorflow/contrib/lite/delegates/eager/kernel_test.cc
+++ b/tensorflow/contrib/lite/delegates/flex/kernel_test.cc
@@ -12,15 +12,15 @@
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
-#include "tensorflow/contrib/lite/delegates/eager/kernel.h"
+#include "tensorflow/contrib/lite/delegates/flex/kernel.h"
#include <gmock/gmock.h>
#include <gtest/gtest.h>
-#include "tensorflow/contrib/lite/delegates/eager/delegate_data.h"
-#include "tensorflow/contrib/lite/delegates/eager/test_util.h"
+#include "tensorflow/contrib/lite/delegates/flex/delegate_data.h"
+#include "tensorflow/contrib/lite/delegates/flex/test_util.h"
namespace tflite {
-namespace eager {
+namespace flex {
namespace {
using ::testing::ContainsRegex;
@@ -31,12 +31,12 @@
TfLiteIntArray* size_and_nodes =
ConvertVectorToTfLiteIntArray(supported_nodes);
TF_LITE_ENSURE_STATUS(context->ReplaceSubgraphsWithDelegateKernels(
- context, eager::GetKernel(), size_and_nodes, delegate));
+ context, flex::GetKernel(), size_and_nodes, delegate));
TfLiteIntArrayFree(size_and_nodes);
return kTfLiteOk;
}
-class KernelTest : public testing::EagerModelTest {
+class KernelTest : public testing::FlexModelTest {
public:
KernelTest() {
CHECK(DelegateData::Create(&delegate_data_).ok());
@@ -167,7 +167,7 @@
ASSERT_FALSE(Invoke());
ASSERT_THAT(error_reporter().error_messages(),
- ContainsRegex("Invalid NodeDef in Eager op"));
+ ContainsRegex("Invalid NodeDef in Flex op"));
}
TEST_F(KernelTest, MixedGraph) {
@@ -220,7 +220,7 @@
}
} // namespace
-} // namespace eager
+} // namespace flex
} // namespace tflite
int main(int argc, char** argv) {
diff --git a/tensorflow/contrib/lite/delegates/eager/test_util.cc b/tensorflow/contrib/lite/delegates/flex/test_util.cc
similarity index 75%
rename from tensorflow/contrib/lite/delegates/eager/test_util.cc
rename to tensorflow/contrib/lite/delegates/flex/test_util.cc
index d47be76..69c336a 100644
--- a/tensorflow/contrib/lite/delegates/eager/test_util.cc
+++ b/tensorflow/contrib/lite/delegates/flex/test_util.cc
@@ -13,25 +13,24 @@
limitations under the License.
==============================================================================*/
-#include "tensorflow/contrib/lite/delegates/eager/test_util.h"
+#include "tensorflow/contrib/lite/delegates/flex/test_util.h"
#include "absl/memory/memory.h"
#include "flatbuffers/flexbuffers.h" // TF:flatbuffers
#include "tensorflow/contrib/lite/string.h"
namespace tflite {
-namespace eager {
+namespace flex {
namespace testing {
-bool EagerModelTest::Invoke() { return interpreter_->Invoke() == kTfLiteOk; }
+bool FlexModelTest::Invoke() { return interpreter_->Invoke() == kTfLiteOk; }
-void EagerModelTest::SetShape(int tensor_index,
- const std::vector<int>& values) {
+void FlexModelTest::SetShape(int tensor_index, const std::vector<int>& values) {
ASSERT_EQ(interpreter_->ResizeInputTensor(tensor_index, values), kTfLiteOk);
ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
}
-std::vector<int> EagerModelTest::GetShape(int tensor_index) {
+std::vector<int> FlexModelTest::GetShape(int tensor_index) {
std::vector<int> result;
auto* dims = interpreter_->tensor(tensor_index)->dims;
result.reserve(dims->size);
@@ -41,13 +40,13 @@
return result;
}
-TfLiteType EagerModelTest::GetType(int tensor_index) {
+TfLiteType FlexModelTest::GetType(int tensor_index) {
return interpreter_->tensor(tensor_index)->type;
}
-void EagerModelTest::AddTensors(int num_tensors, const std::vector<int>& inputs,
- const std::vector<int>& outputs,
- TfLiteType type, const std::vector<int>& dims) {
+void FlexModelTest::AddTensors(int num_tensors, const std::vector<int>& inputs,
+ const std::vector<int>& outputs, TfLiteType type,
+ const std::vector<int>& dims) {
interpreter_->AddTensors(num_tensors);
for (int i = 0; i < num_tensors; ++i) {
TfLiteQuantizationParams quant;
@@ -66,8 +65,8 @@
CHECK_EQ(interpreter_->SetOutputs(outputs), kTfLiteOk);
}
-void EagerModelTest::AddTfLiteMulOp(const std::vector<int>& inputs,
- const std::vector<int>& outputs) {
+void FlexModelTest::AddTfLiteMulOp(const std::vector<int>& inputs,
+ const std::vector<int>& outputs) {
static TfLiteRegistration reg = {nullptr, nullptr, nullptr, nullptr};
reg.builtin_code = BuiltinOperator_MUL;
reg.prepare = [](TfLiteContext* context, TfLiteNode* node) {
@@ -90,8 +89,8 @@
kTfLiteOk);
}
-void EagerModelTest::AddTfOp(TfOpType op, const std::vector<int>& inputs,
- const std::vector<int>& outputs) {
+void FlexModelTest::AddTfOp(TfOpType op, const std::vector<int>& inputs,
+ const std::vector<int>& outputs) {
auto attr = [](const string& key, const string& value) {
return " attr{ key: '" + key + "' value {" + value + "}}";
};
@@ -107,28 +106,28 @@
if (op == kUnpack) {
string attributes =
type_attribute + attr("num", "i: 2") + attr("axis", "i: 0");
- AddTfOp("EagerUnpack", "Unpack", attributes, inputs, outputs);
+ AddTfOp("FlexUnpack", "Unpack", attributes, inputs, outputs);
} else if (op == kIdentity) {
string attributes = type_attribute;
- AddTfOp("EagerIdentity", "Identity", attributes, inputs, outputs);
+ AddTfOp("FlexIdentity", "Identity", attributes, inputs, outputs);
} else if (op == kAdd) {
string attributes = type_attribute;
- AddTfOp("EagerAdd", "Add", attributes, inputs, outputs);
+ AddTfOp("FlexAdd", "Add", attributes, inputs, outputs);
} else if (op == kMul) {
string attributes = type_attribute;
- AddTfOp("EagerMul", "Mul", attributes, inputs, outputs);
+ AddTfOp("FlexMul", "Mul", attributes, inputs, outputs);
} else if (op == kNonExistent) {
AddTfOp("NonExistentOp", "NonExistentOp", "", inputs, outputs);
} else if (op == kIncompatibleNodeDef) {
// "Cast" op is created without attributes - making it incompatible.
- AddTfOp("EagerCast", "Cast", "", inputs, outputs);
+ AddTfOp("FlexCast", "Cast", "", inputs, outputs);
}
}
-void EagerModelTest::AddTfOp(const char* tflite_name, const string& tf_name,
- const string& nodedef_str,
- const std::vector<int>& inputs,
- const std::vector<int>& outputs) {
+void FlexModelTest::AddTfOp(const char* tflite_name, const string& tf_name,
+ const string& nodedef_str,
+ const std::vector<int>& inputs,
+ const std::vector<int>& outputs) {
static TfLiteRegistration reg = {nullptr, nullptr, nullptr, nullptr};
reg.builtin_code = BuiltinOperator_CUSTOM;
reg.custom_name = tflite_name;
@@ -154,5 +153,5 @@
}
} // namespace testing
-} // namespace eager
+} // namespace flex
} // namespace tflite
diff --git a/tensorflow/contrib/lite/delegates/eager/test_util.h b/tensorflow/contrib/lite/delegates/flex/test_util.h
similarity index 89%
rename from tensorflow/contrib/lite/delegates/eager/test_util.h
rename to tensorflow/contrib/lite/delegates/flex/test_util.h
index 816db41..a8c81b9 100644
--- a/tensorflow/contrib/lite/delegates/eager/test_util.h
+++ b/tensorflow/contrib/lite/delegates/flex/test_util.h
@@ -13,14 +13,14 @@
limitations under the License.
==============================================================================*/
-#ifndef TENSORFLOW_CONTRIB_LITE_DELEGATES_EAGER_TEST_UTIL_H_
-#define TENSORFLOW_CONTRIB_LITE_DELEGATES_EAGER_TEST_UTIL_H_
+#ifndef TENSORFLOW_CONTRIB_LITE_DELEGATES_FLEX_TEST_UTIL_H_
+#define TENSORFLOW_CONTRIB_LITE_DELEGATES_FLEX_TEST_UTIL_H_
#include "tensorflow/c/c_api_internal.h"
#include "tensorflow/contrib/lite/kernels/test_util.h"
namespace tflite {
-namespace eager {
+namespace flex {
namespace testing {
enum TfOpType {
@@ -35,12 +35,12 @@
};
// This class creates models with TF and TFLite ops. In order to use this class
-// to test the Eager delegate, implement a function that calls
+// to test the Flex delegate, implement a function that calls
// interpreter->ModifyGraphWithDelegate.
-class EagerModelTest : public ::testing::Test {
+class FlexModelTest : public ::testing::Test {
public:
- EagerModelTest() {}
- ~EagerModelTest() {}
+ FlexModelTest() {}
+ ~FlexModelTest() {}
bool Invoke();
@@ -104,7 +104,7 @@
private:
// Helper method to add a TensorFlow op. tflite_names needs to start with
- // "Eager" in order to work with the Eager delegate.
+ // "Flex" in order to work with the Flex delegate.
void AddTfOp(const char* tflite_name, const string& tf_name,
const string& nodedef_str, const std::vector<int>& inputs,
const std::vector<int>& outputs);
@@ -113,7 +113,7 @@
};
} // namespace testing
-} // namespace eager
+} // namespace flex
} // namespace tflite
-#endif // TENSORFLOW_CONTRIB_LITE_DELEGATES_EAGER_TEST_UTIL_H_
+#endif // TENSORFLOW_CONTRIB_LITE_DELEGATES_FLEX_TEST_UTIL_H_
diff --git a/tensorflow/contrib/lite/delegates/eager/util.cc b/tensorflow/contrib/lite/delegates/flex/util.cc
similarity index 96%
rename from tensorflow/contrib/lite/delegates/eager/util.cc
rename to tensorflow/contrib/lite/delegates/flex/util.cc
index 051246b..829bc38 100644
--- a/tensorflow/contrib/lite/delegates/eager/util.cc
+++ b/tensorflow/contrib/lite/delegates/flex/util.cc
@@ -12,10 +12,10 @@
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
-#include "tensorflow/contrib/lite/delegates/eager/util.h"
+#include "tensorflow/contrib/lite/delegates/flex/util.h"
namespace tflite {
-namespace eager {
+namespace flex {
TfLiteStatus ConvertStatus(TfLiteContext* context,
const tensorflow::Status& status) {
@@ -100,5 +100,5 @@
}
}
-} // namespace eager
+} // namespace flex
} // namespace tflite
diff --git a/tensorflow/contrib/lite/delegates/eager/util.h b/tensorflow/contrib/lite/delegates/flex/util.h
similarity index 89%
rename from tensorflow/contrib/lite/delegates/eager/util.h
rename to tensorflow/contrib/lite/delegates/flex/util.h
index 930cb99..7f910e7 100644
--- a/tensorflow/contrib/lite/delegates/eager/util.h
+++ b/tensorflow/contrib/lite/delegates/flex/util.h
@@ -12,8 +12,8 @@
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
-#ifndef TENSORFLOW_CONTRIB_LITE_DELEGATES_EAGER_UTIL_H_
-#define TENSORFLOW_CONTRIB_LITE_DELEGATES_EAGER_UTIL_H_
+#ifndef TENSORFLOW_CONTRIB_LITE_DELEGATES_FLEX_UTIL_H_
+#define TENSORFLOW_CONTRIB_LITE_DELEGATES_FLEX_UTIL_H_
#include "tensorflow/c/c_api_internal.h"
#include "tensorflow/contrib/lite/c/c_api_internal.h"
@@ -21,7 +21,7 @@
#include "tensorflow/core/lib/core/status.h"
namespace tflite {
-namespace eager {
+namespace flex {
// Converts a tensorflow:Status into a TfLiteStatus. If the original status
// represented an error, reports it using the given 'context'.
@@ -41,7 +41,7 @@
// Returns the TfLiteType that corresponds to the given TF C API Data type.
TfLiteType GetTensorFlowLiteType(TF_DataType);
-} // namespace eager
+} // namespace flex
} // namespace tflite
-#endif // TENSORFLOW_CONTRIB_LITE_DELEGATES_EAGER_UTIL_H_
+#endif // TENSORFLOW_CONTRIB_LITE_DELEGATES_FLEX_UTIL_H_
diff --git a/tensorflow/contrib/lite/delegates/eager/util_test.cc b/tensorflow/contrib/lite/delegates/flex/util_test.cc
similarity index 97%
rename from tensorflow/contrib/lite/delegates/eager/util_test.cc
rename to tensorflow/contrib/lite/delegates/flex/util_test.cc
index aebc911..5f049e7 100644
--- a/tensorflow/contrib/lite/delegates/eager/util_test.cc
+++ b/tensorflow/contrib/lite/delegates/flex/util_test.cc
@@ -12,7 +12,7 @@
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
-#include "tensorflow/contrib/lite/delegates/eager/util.h"
+#include "tensorflow/contrib/lite/delegates/flex/util.h"
#include <cstdarg>
@@ -22,7 +22,7 @@
#include "tensorflow/contrib/lite/testing/util.h"
namespace tflite {
-namespace eager {
+namespace flex {
namespace {
using tensorflow::DT_FLOAT;
@@ -132,7 +132,7 @@
}
} // namespace
-} // namespace eager
+} // namespace flex
} // namespace tflite
int main(int argc, char** argv) {
diff --git a/tensorflow/contrib/lite/examples/android/BUILD b/tensorflow/contrib/lite/examples/android/BUILD
index 4d2437e..d180cb4 100644
--- a/tensorflow/contrib/lite/examples/android/BUILD
+++ b/tensorflow/contrib/lite/examples/android/BUILD
@@ -28,6 +28,7 @@
srcs = glob([
"app/src/main/java/**/*.java",
]),
+ aapt_version = "aapt",
# Package assets from assets dir as well as all model targets.
# Remove undesired models (and corresponding Activities in source)
# to reduce APK size.
diff --git a/tensorflow/contrib/lite/experimental/c/c_api_experimental.cc b/tensorflow/contrib/lite/experimental/c/c_api_experimental.cc
index 0f16595..29f8701 100644
--- a/tensorflow/contrib/lite/experimental/c/c_api_experimental.cc
+++ b/tensorflow/contrib/lite/experimental/c/c_api_experimental.cc
@@ -21,9 +21,8 @@
extern "C" {
#endif // __cplusplus
-TFL_Status TFL_InterpreterResetVariableTensorsToZero(
- TFL_Interpreter* interpreter) {
- return interpreter->impl->ResetVariableTensorsToZero();
+TFL_Status TFL_InterpreterResetVariableTensors(TFL_Interpreter* interpreter) {
+ return interpreter->impl->ResetVariableTensors();
}
void TFL_InterpreterOptionsAddBuiltinOp(TFL_InterpreterOptions* options,
diff --git a/tensorflow/contrib/lite/experimental/c/c_api_experimental.h b/tensorflow/contrib/lite/experimental/c/c_api_experimental.h
index b8de7b9..fca5d92 100644
--- a/tensorflow/contrib/lite/experimental/c/c_api_experimental.h
+++ b/tensorflow/contrib/lite/experimental/c/c_api_experimental.h
@@ -25,7 +25,7 @@
typedef TfLiteBuiltinOperator TFL_BuiltinOperator;
// Resets all variable tensors to zero.
-TFL_CAPI_EXPORT extern TFL_Status TFL_InterpreterResetVariableTensorsToZero(
+TFL_CAPI_EXPORT extern TFL_Status TFL_InterpreterResetVariableTensors(
TFL_Interpreter* interpreter);
// Adds an op registration for a builtin operator.
diff --git a/tensorflow/contrib/lite/experimental/c/c_api_experimental_test.cc b/tensorflow/contrib/lite/experimental/c/c_api_experimental_test.cc
index d86ad00..1b1bedb 100644
--- a/tensorflow/contrib/lite/experimental/c/c_api_experimental_test.cc
+++ b/tensorflow/contrib/lite/experimental/c/c_api_experimental_test.cc
@@ -44,7 +44,7 @@
TFL_Interpreter* interpreter = TFL_NewInterpreter(model, options);
ASSERT_NE(interpreter, nullptr);
ASSERT_EQ(TFL_InterpreterAllocateTensors(interpreter), kTfLiteOk);
- EXPECT_EQ(TFL_InterpreterResetVariableTensorsToZero(interpreter), kTfLiteOk);
+ EXPECT_EQ(TFL_InterpreterResetVariableTensors(interpreter), kTfLiteOk);
EXPECT_EQ(TFL_InterpreterInvoke(interpreter), kTfLiteOk);
TFL_DeleteInterpreter(interpreter);
diff --git a/tensorflow/contrib/lite/g3doc/performance.md b/tensorflow/contrib/lite/g3doc/performance.md
index 0ae9400..6b7943c 100644
--- a/tensorflow/contrib/lite/g3doc/performance.md
+++ b/tensorflow/contrib/lite/g3doc/performance.md
@@ -7,12 +7,12 @@
Some models may be too large to run on embedded devices. Instead of large models it is better to use a slightly less precise but smaller model for embedded devices. Smaller models not only use less disk space and memory but are generally faster and more energy efficient. One example of models optimized for mobile devices are [MobileNets](https://arxiv.org/abs/1704.04861), which are optimized for mobile vision applications. Tensorflow Lite [models page](models.md) lists several other models that have been optimized specifically for mobile and embedded devices.
You can retrain the listed models on your own dataset by using transfer learning. Check out our transfer learning tutorial for
-[image classification] (https://codelabs.developers.google.com/codelabs/tensorflow-for-poets/#0) and
+[image classification](https://codelabs.developers.google.com/codelabs/tensorflow-for-poets/#0) and
[object detection](https://medium.com/tensorflow/training-and-serving-a-realtime-mobile-object-detector-in-30-minutes-with-cloud-tpus-b78971cf1193).
## Profile your model
-Before starting any optimization, it is a good practice to profile and benchmark your model. Tensorflow Lite [benchmarking tool](../tools/benchmark) has a built-in profiler that shows per operator profiling statistics. This can help in understanding performance bottlenecks and which operators dominate the computation time.
+Before starting any optimization, it is a good practice to profile and benchmark your model. Tensorflow Lite [benchmarking tool](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/tools/benchmark) has a built-in profiler that shows per operator profiling statistics. This can help in understanding performance bottlenecks and which operators dominate the computation time.
## Profile and optimize operators in the graph
If a particular operator appears frequently in the model and based on profiling you find the operator consuming the most amount of time, you can look into optimizing the operator.
@@ -22,7 +22,7 @@
If your model uses floating point weights or activations then it may be possible to reduce the size of model up to ~4x by using quantization and other model optimizations. Check out our [model optimization toolkit](https://www.tensorflow.org/performance/model_optimization) for details about optimizing your model. Fully quantized models can be remarkably power efficient as well.
## Tweak the number of threads
-Tensorflow Lite supports multi-threaded kernels for many operators. You can increase the number of threads and speed up execution of operators. Increasing the number of threads will however make your model use more resources and power. For some applications latency may be more important than energy efficiency. You can increase the number of threads by setting the number of [interpreter](../interpreter.h) threads.
+Tensorflow Lite supports multi-threaded kernels for many operators. You can increase the number of threads and speed up execution of operators. Increasing the number of threads will however make your model use more resources and power. For some applications latency may be more important than energy efficiency. You can increase the number of threads by setting the number of [interpreter](https://github.com/tensorflow/tensorflow/blob/1084594657a5d139102ac794f84d1427a710e39a/tensorflow/contrib/lite/interpreter.h#L337) threads.
## Eliminate redundant copies
Tensorflow Lite is optimized to reduce redundant copies. The APIs allow user to [mmap a model file](https://github.com/tensorflow/tensorflow/blob/9982fd6c8831cbd2f58954f79ea71f26660393bc/tensorflow/contrib/lite/model.h#L152) and avoid copies. If your application is not careful, there can be redundant copies when feeding the input to the model and reading output from the model. Make sure to eliminate redundant copies. If you are using higher level APIs like Java API, make sure to carefully check the documentation for performance caveats. For example, the Java API is a lot faster if ByteBuffers are used as [inputs](https://github.com/tensorflow/tensorflow/blob/6305a6d83552ba6a472cd72398b60d9241467f1f/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/Interpreter.java#L151).
@@ -31,8 +31,8 @@
Platform specific tools like [Android profiler](https://developer.android.com/studio/profile/android-profiler) and [Instruments](https://help.apple.com/instruments/mac/current/) provide a wealth of profiling information that can be used to debug your app. Sometimes the performance bug may be not in the model but in parts of application code that interact with the model. Make sure to familiarize yourself with platform specific profiling tools and best practices for your platform.
## Use hardware accelerators available on the device
-Tensorflow Lite is working on adding support for accelerators like GPU and provides acceleration through [NNAPI](https://developer.android.com/ndk/guides/neuralnetworks/) on Android.
-You can utilize these hardware accelerator backends to improve the speed and efficiency of your model. To enable NNAPI call [UseNNAPI](https://github.com/tensorflow/tensorflow/blob/6305a6d83552ba6a472cd72398b60d9241467f1f/tensorflow/contrib/lite/interpreter.h#L334) on the interpreter instance.
+Tensorflow Lite is working on adding support for accelerators like GPU and provides acceleration through [Neural Networks API](https://developer.android.com/ndk/guides/neuralnetworks/) on Android.
+You can utilize these hardware accelerator backends to improve the speed and efficiency of your model. To enable Neural Networks API call [UseNNAPI](https://github.com/tensorflow/tensorflow/blob/6305a6d83552ba6a472cd72398b60d9241467f1f/tensorflow/contrib/lite/interpreter.h#L334) on the interpreter instance.
## Need more help
The Tensorflow team is happy to help diagnose and address specific performance issues you may be facing. Please file a bug on [github](https://github.com/tensorflow/tensorflow/issues) with details of the issue.
diff --git a/tensorflow/contrib/lite/interpreter.cc b/tensorflow/contrib/lite/interpreter.cc
index 2657bcd..88e41ff 100644
--- a/tensorflow/contrib/lite/interpreter.cc
+++ b/tensorflow/contrib/lite/interpreter.cc
@@ -451,16 +451,15 @@
// Reset the variable tensors to zero after (re)allocating the tensors.
// Developers shouldn't rely on the side effect of this function to reset
- // variable tesnsors. They should call `ResetVariableTensorsToZero` directly
+ // variable tesnsors. They should call `ResetVariableTensors` directly
// instead.
- ResetVariableTensorsToZero();
+ ResetVariableTensors();
return kTfLiteOk;
}
-// TODO(ycling): Consider to provide other functions to initialize variable
-// tensors to non-zero values.
-TfLiteStatus Interpreter::ResetVariableTensorsToZero() {
+// TODO(ycling): Support non-zero default values.
+TfLiteStatus Interpreter::ResetVariableTensors() {
for (auto& tensor : tensors_) {
if (!tensor.is_variable) {
continue;
diff --git a/tensorflow/contrib/lite/interpreter.h b/tensorflow/contrib/lite/interpreter.h
index aa2bc4d..7ef736d 100644
--- a/tensorflow/contrib/lite/interpreter.h
+++ b/tensorflow/contrib/lite/interpreter.h
@@ -421,9 +421,12 @@
allow_buffer_handle_output_ = allow_buffer_handle_output;
}
- // Reset all variable tensors to zero.
+ // Reset all variable tensors to the default value.
+ // If a variable tensor doesn't have a buffer, reset it to zero.
+ // TODO(b/115961645): Implement - If a variable tensor has a buffer, reset it
+ // to the value of the buffer.
// WARNING: This is an experimental API and subject to change.
- TfLiteStatus ResetVariableTensorsToZero();
+ TfLiteStatus ResetVariableTensors();
// Retrieve an operator's description of its work, for profiling purposes.
const char* OpProfilingString(const TfLiteRegistration& op_reg,
diff --git a/tensorflow/contrib/lite/java/aar_with_jni.bzl b/tensorflow/contrib/lite/java/aar_with_jni.bzl
index db837cf..9d2aead 100644
--- a/tensorflow/contrib/lite/java/aar_with_jni.bzl
+++ b/tensorflow/contrib/lite/java/aar_with_jni.bzl
@@ -3,12 +3,12 @@
load("@build_bazel_rules_android//android:rules.bzl", "android_binary")
def aar_with_jni(name, android_library):
- # Generate dummy AndroidManifest.xml for dummy apk usage
- # (dummy apk is generated by <name>_dummy_app_for_so target below)
- native.genrule(
- name = name + "_binary_manifest_generator",
- outs = [name + "_generated_AndroidManifest.xml"],
- cmd = """
+ # Generate dummy AndroidManifest.xml for dummy apk usage
+ # (dummy apk is generated by <name>_dummy_app_for_so target below)
+ native.genrule(
+ name = name + "_binary_manifest_generator",
+ outs = [name + "_generated_AndroidManifest.xml"],
+ cmd = """
cat > $(OUTS) <<EOF
<manifest
xmlns:android="http://schemas.android.com/apk/res/android"
@@ -17,27 +17,28 @@
</manifest>
EOF
""",
- )
+ )
- # Generate dummy apk including .so files and later we extract out
- # .so files and throw away the apk.
- android_binary(
- name = name + "_dummy_app_for_so",
- manifest = name + "_generated_AndroidManifest.xml",
- custom_package = "dummy.package.for.so",
- deps = [android_library],
- # 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"],
- )
+ # Generate dummy apk including .so files and later we extract out
+ # .so files and throw away the apk.
+ android_binary(
+ name = name + "_dummy_app_for_so",
+ aapt_version = "aapt",
+ manifest = name + "_generated_AndroidManifest.xml",
+ custom_package = "dummy.package.for.so",
+ deps = [android_library],
+ # 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"],
+ )
- native.genrule(
- name = name,
- srcs = [android_library + ".aar", name + "_dummy_app_for_so_unsigned.apk"],
- outs = [name + ".aar"],
- tags = ["manual"],
- cmd = """
+ native.genrule(
+ name = name,
+ srcs = [android_library + ".aar", name + "_dummy_app_for_so_unsigned.apk"],
+ outs = [name + ".aar"],
+ tags = ["manual"],
+ cmd = """
cp $(location {}.aar) $(location :{}.aar)
chmod +w $(location :{}.aar)
origdir=$$PWD
@@ -46,4 +47,4 @@
cp -r lib jni
zip -r $$origdir/$(location :{}.aar) jni/*/*.so
""".format(android_library, name, name, name, name),
- )
+ )
diff --git a/tensorflow/contrib/lite/java/demo/app/src/main/BUILD b/tensorflow/contrib/lite/java/demo/app/src/main/BUILD
index 220d6c2..5ad7383 100644
--- a/tensorflow/contrib/lite/java/demo/app/src/main/BUILD
+++ b/tensorflow/contrib/lite/java/demo/app/src/main/BUILD
@@ -7,6 +7,7 @@
android_binary(
name = "TfLiteCameraDemo",
srcs = glob(["java/**/*.java"]),
+ aapt_version = "aapt",
assets = [
"//tensorflow/contrib/lite/java/demo/app/src/main/assets:labels_mobilenet_quant_v1_224.txt",
"@tflite_mobilenet//:mobilenet_quant_v1_224.tflite",
diff --git a/tensorflow/contrib/lite/java/ovic/BUILD b/tensorflow/contrib/lite/java/ovic/BUILD
index bb0be04..ea9b9ed 100644
--- a/tensorflow/contrib/lite/java/ovic/BUILD
+++ b/tensorflow/contrib/lite/java/ovic/BUILD
@@ -9,6 +9,7 @@
load("//tensorflow/java:build_defs.bzl", "JAVACOPTS")
+# Build targets for OVIC classification.
java_test(
name = "OvicClassifierTest",
size = "medium",
@@ -45,8 +46,9 @@
name = "ovicbenchmarkerlib",
srcs = [
"src/main/java/org/tensorflow/ovic/OvicBenchmarker.java",
+ "src/main/java/org/tensorflow/ovic/OvicClassificationResult.java",
"src/main/java/org/tensorflow/ovic/OvicClassifier.java",
- "src/main/java/org/tensorflow/ovic/OvicSingleImageResult.java",
+ "src/main/java/org/tensorflow/ovic/OvicClassifierBenchmarker.java",
],
manifest = "//tensorflow/contrib/lite/java:AndroidManifest.xml",
tags = ["no_oss"],
@@ -60,8 +62,8 @@
java_library(
name = "ovicbenchmarkerlib_java",
srcs = [
+ "src/main/java/org/tensorflow/ovic/OvicClassificationResult.java",
"src/main/java/org/tensorflow/ovic/OvicClassifier.java",
- "src/main/java/org/tensorflow/ovic/OvicSingleImageResult.java",
],
javacopts = JAVACOPTS,
tags = ["no_oss"],
@@ -73,3 +75,58 @@
"@org_checkerframework_qual",
],
)
+
+# Build targets for OVIC detection.
+java_test(
+ name = "OvicDetectorTest",
+ size = "medium",
+ srcs = ["src/test/java/org/tensorflow/ovic/OvicDetectorTest.java"],
+ data = [
+ "//tensorflow/contrib/lite/java/ovic/src/testdata:coco_labels.txt",
+ "//tensorflow/contrib/lite/java/ovic/src/testdata:ovic_testdata",
+ "@tflite_mobilenet_ssd_quant//:detect.tflite",
+ ],
+ javacopts = JAVACOPTS,
+ tags = ["no_oss"],
+ test_class = "org.tensorflow.ovic.OvicDetectorTest",
+ visibility = ["//visibility:public"],
+ deps = [
+ "//tensorflow/contrib/lite/java/ovic:ovicdetectionbenchmarkerlib_java",
+ "@com_google_truth",
+ "@junit",
+ ],
+)
+
+android_library(
+ name = "ovicdetectionbenchmarkerlib",
+ srcs = [
+ "src/main/java/org/tensorflow/ovic/BoundingBox.java",
+ "src/main/java/org/tensorflow/ovic/OvicBenchmarker.java",
+ "src/main/java/org/tensorflow/ovic/OvicDetectionResult.java",
+ "src/main/java/org/tensorflow/ovic/OvicDetector.java",
+ "src/main/java/org/tensorflow/ovic/OvicDetectorBenchmarker.java",
+ ],
+ manifest = "//tensorflow/contrib/lite/java:AndroidManifest.xml",
+ deps = [
+ "//tensorflow/contrib/lite/java:tensorflowlite",
+ "//tensorflow/contrib/lite/java/src/testhelper/java/org/tensorflow/lite:testhelper",
+ "@org_checkerframework_qual",
+ ],
+)
+
+java_library(
+ name = "ovicdetectionbenchmarkerlib_java",
+ srcs = [
+ "src/main/java/org/tensorflow/ovic/BoundingBox.java",
+ "src/main/java/org/tensorflow/ovic/OvicDetectionResult.java",
+ "src/main/java/org/tensorflow/ovic/OvicDetector.java",
+ ],
+ javacopts = JAVACOPTS,
+ deps = [
+ "//tensorflow/contrib/lite/java:libtensorflowlite_jni.so",
+ "//tensorflow/contrib/lite/java:tensorflowlite_java",
+ "//tensorflow/contrib/lite/java/src/main/native",
+ "//tensorflow/contrib/lite/java/src/testhelper/java/org/tensorflow/lite:testhelper",
+ "@org_checkerframework_qual",
+ ],
+)
diff --git a/tensorflow/contrib/lite/java/ovic/demo/app/BUILD b/tensorflow/contrib/lite/java/ovic/demo/app/BUILD
index b2e3a9bd..f567358 100644
--- a/tensorflow/contrib/lite/java/ovic/demo/app/BUILD
+++ b/tensorflow/contrib/lite/java/ovic/demo/app/BUILD
@@ -8,9 +8,12 @@
srcs = [
"OvicBenchmarkerActivity.java",
],
+ aapt_version = "aapt",
assets = [
- "//tensorflow/contrib/lite/java/ovic/src/testdata:ovic_testdata",
+ "//tensorflow/contrib/lite/java/ovic/src/testdata:coco_labels.txt",
"//tensorflow/contrib/lite/java/ovic/src/testdata:labels.txt",
+ "//tensorflow/contrib/lite/java/ovic/src/testdata:ovic_testdata",
+ "@tflite_mobilenet_ssd_quant//:detect.tflite",
],
assets_dir = "",
custom_package = "ovic.demo.app",
@@ -24,6 +27,7 @@
deps = [
"//tensorflow/contrib/lite/java:tensorflowlite",
"//tensorflow/contrib/lite/java/ovic:ovicbenchmarkerlib",
+ "//tensorflow/contrib/lite/java/ovic:ovicdetectionbenchmarkerlib",
"@androidsdk//com.android.support:support-v13-25.2.0",
"@androidsdk//com.android.support:support-v4-25.2.0",
],
diff --git a/tensorflow/contrib/lite/java/ovic/demo/app/OvicBenchmarkerActivity.java b/tensorflow/contrib/lite/java/ovic/demo/app/OvicBenchmarkerActivity.java
index 4adf94a..48c29ec 100644
--- a/tensorflow/contrib/lite/java/ovic/demo/app/OvicBenchmarkerActivity.java
+++ b/tensorflow/contrib/lite/java/ovic/demo/app/OvicBenchmarkerActivity.java
@@ -35,19 +35,18 @@
import java.nio.channels.FileChannel;
import java.text.DecimalFormat;
import org.tensorflow.ovic.OvicBenchmarker;
-import org.tensorflow.ovic.OvicSingleImageResult;
-
+import org.tensorflow.ovic.OvicClassifierBenchmarker;
+import org.tensorflow.ovic.OvicDetectorBenchmarker;
/** Class that benchmark image classifier models. */
public class OvicBenchmarkerActivity extends Activity {
/** Tag for the {@link Log}. */
private static final String TAG = "OvicBenchmarkerActivity";
- /** Name of the label file stored in Assets. */
- private static final String LABEL_PATH = "labels.txt";
-
- private static final String TEST_IMAGE_PATH = "test_image_224.jpg";
- private static final String MODEL_PATH = "float_model.lite";
+ /** Name of the task-dependent data files stored in Assets. */
+ private static String labelPath = null;
+ private static String testImagePath = null;
+ private static String modelPath = null;
/**
* Each bottom press will launch a benchmarking experiment. The experiment stops when either the
* total native latency reaches WALL_TIME or the number of iterations reaches MAX_ITERATIONS,
@@ -66,8 +65,6 @@
private MappedByteBuffer model = null;
private InputStream labelInputStream = null;
private OvicBenchmarker benchmarker;
- /** Inference result of each iteration. */
- OvicSingleImageResult iterResult = null;
private TextView textView = null;
// private Button startButton = null;
@@ -83,21 +80,31 @@
}
private Bitmap loadTestBitmap() throws IOException {
- InputStream imageStream = getAssets().open(TEST_IMAGE_PATH);
+ InputStream imageStream = getAssets().open(testImagePath);
return BitmapFactory.decodeStream(imageStream);
}
- public void initializeTest() throws IOException {
+ public void initializeTest(boolean benchmarkClassification) throws IOException {
Log.i(TAG, "Initializing benchmarker.");
- benchmarker = new OvicBenchmarker(WALL_TIME);
+ if (benchmarkClassification) {
+ benchmarker = new OvicClassifierBenchmarker(WALL_TIME);
+ labelPath = "labels.txt";
+ testImagePath = "test_image_224.jpg";
+ modelPath = "quantized_model.lite";
+ } else { // Benchmarking detection.
+ benchmarker = new OvicDetectorBenchmarker(WALL_TIME);
+ labelPath = "coco_labels.txt";
+ testImagePath = "test_image_224.jpg";
+ modelPath = "detect.tflite";
+ }
AssetManager am = getAssets();
- AssetFileDescriptor fileDescriptor = am.openFd(MODEL_PATH);
+ AssetFileDescriptor fileDescriptor = am.openFd(modelPath);
FileInputStream modelInputStream = new FileInputStream(fileDescriptor.getFileDescriptor());
FileChannel fileChannel = modelInputStream.getChannel();
long startOffset = fileDescriptor.getStartOffset();
long declaredLength = fileDescriptor.getDeclaredLength();
model = fileChannel.map(FileChannel.MapMode.READ_ONLY, startOffset, declaredLength);
- labelInputStream = am.open(LABEL_PATH);
+ labelInputStream = am.open(labelPath);
}
public Boolean doTestIteration() throws IOException, InterruptedException {
@@ -117,24 +124,44 @@
Log.i(TAG, "Going to do test iter.");
// Start testing.
Bitmap testImageBitmap = loadTestBitmap();
- iterResult = benchmarker.doTestIteration(testImageBitmap);
- testImageBitmap.recycle();
- if (iterResult == null) {
+ try {
+ if (!benchmarker.processBitmap(testImageBitmap)) {
+ throw new RuntimeException("Failed to run test.");
+ }
+ } catch (Exception e) {
+ e.printStackTrace();
+ throw e;
+ } finally {
+ testImageBitmap.recycle();
+ }
+ String iterResultString = benchmarker.getLastResultString();
+ if (iterResultString == null) {
throw new RuntimeException("Inference failed to produce a result.");
}
- Log.i(TAG, iterResult.toString());
+ Log.i(TAG, iterResultString);
return true;
}
- public void startPressed(View view) throws IOException {
- Log.i(TAG, "Start pressed");
+ public void detectPressed(View view) throws IOException {
+ benchmarkSession(false);
+ }
+ public void classifyPressed(View view) throws IOException {
+ benchmarkSession(true);
+ }
+
+ private void benchmarkSession(boolean benchmarkClassification) throws IOException {
try {
- initializeTest();
+ initializeTest(benchmarkClassification);
} catch (IOException e) {
Log.e(TAG, "Can't initialize benchmarker.", e);
throw e;
}
String displayText = "";
+ if (benchmarkClassification) {
+ displayText = "Classification benchmark: ";
+ } else {
+ displayText = "Detection benchmark: ";
+ }
try {
setProcessorAffinity(BIG_CORE_MASK);
} catch (IOException e) {
@@ -144,7 +171,6 @@
Log.i(TAG, "Successfully initialized benchmarker.");
int testIter = 0;
Boolean iterSuccess = false;
- double totalLatency = 0.0f;
while (testIter < MAX_ITERATIONS) {
try {
iterSuccess = doTestIteration();
@@ -153,23 +179,22 @@
throw e;
} catch (InterruptedException e) {
Log.e(TAG, "Interrupted at iteration " + testIter);
+ displayText += e.getMessage() + "\n";
}
if (!iterSuccess) {
break;
}
testIter++;
- totalLatency += (double) iterResult.latency;
}
- ;
Log.i(TAG, "Benchmarking finished");
if (textView != null) {
if (testIter > 0) {
textView.setText(
displayText
- + MODEL_PATH
+ + modelPath
+ ": Average latency="
- + df2.format(totalLatency / testIter)
+ + df2.format(benchmarker.getTotalRunTime() / testIter)
+ "ms after "
+ testIter
+ " runs.");
diff --git a/tensorflow/contrib/lite/java/ovic/demo/app/res/layout/activity_main.xml b/tensorflow/contrib/lite/java/ovic/demo/app/res/layout/activity_main.xml
index e9d83ba..1bce60f 100644
--- a/tensorflow/contrib/lite/java/ovic/demo/app/res/layout/activity_main.xml
+++ b/tensorflow/contrib/lite/java/ovic/demo/app/res/layout/activity_main.xml
@@ -30,14 +30,14 @@
android:layout_height="wrap_content"
android:text="@string/initial_status_msg"
android:id="@+id/textView"
- android:layout_above="@+id/button_start"
+ android:layout_above="@+id/button_clf_start"
android:layout_alignParentTop="true"/>
<Button
android:layout_width="wrap_content"
android:layout_height="wrap_content"
- android:text="@string/start_label"
- android:id="@id/button_start"
+ android:text="@string/start_clf_label"
+ android:id="@id/button_clf_start"
android:layout_alignParentBottom="true"
android:layout_alignParentLeft="true"
android:background="@drawable/start_button_color"
@@ -49,6 +49,25 @@
android:textColor="#ffffff"
android:enabled="true"
style="?android:attr/buttonBarButtonStyle"
- android:onClick="startPressed"/>
+ android:onClick="classifyPressed"/>
+
+ <Button
+ android:layout_width="wrap_content"
+ android:layout_height="wrap_content"
+ android:text="@string/start_det_label"
+ android:id="@+id/button_det_start"
+ android:layout_alignParentBottom="true"
+ android:layout_alignParentRight="true"
+ android:layout_toRightOf="@id/button_clf_start"
+ android:background="@drawable/start_button_color"
+ android:padding="10dp"
+ android:layout_marginRight="100dp"
+ android:layout_marginLeft="30dp"
+ android:layout_marginTop="10dp"
+ android:foreground="#000000"
+ android:textColor="#ffffff"
+ android:enabled="true"
+ style="?android:attr/buttonBarButtonStyle"
+ android:onClick="detectPressed"/>
</RelativeLayout>
diff --git a/tensorflow/contrib/lite/java/ovic/demo/app/res/values/strings.xml b/tensorflow/contrib/lite/java/ovic/demo/app/res/values/strings.xml
index d26beb1..5352590 100644
--- a/tensorflow/contrib/lite/java/ovic/demo/app/res/values/strings.xml
+++ b/tensorflow/contrib/lite/java/ovic/demo/app/res/values/strings.xml
@@ -17,6 +17,7 @@
<resources>
<string name="app_name" translatable="false">Benchmarker</string>
- <string name="start_label" translatable="false">Start</string>
+ <string name="start_clf_label" translatable="false">Clf</string>
+ <string name="start_det_label" translatable="false">Det</string>
<string name="initial_status_msg" translatable="false"> Press start to run the benchmarks.</string>
</resources>
diff --git a/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/BoundingBox.java b/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/BoundingBox.java
new file mode 100644
index 0000000..9bf7d00
--- /dev/null
+++ b/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/BoundingBox.java
@@ -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
+
+ https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+package org.tensorflow.ovic;
+
+/** Class for holding a detection bounding box with category and confidence. */
+public class BoundingBox {
+ // Upper left point.
+ public float x1;
+ public float y1;
+
+ // Lower right point.
+ public float x2;
+ public float y2;
+
+ // The area of the box
+ public float area;
+
+ // The object category
+ public int category;
+
+ // The confidence of the detection
+ public float score;
+
+ public BoundingBox(float x1, float y1, float x2, float y2, int category, float score) {
+ this.x1 = x1;
+ this.y1 = y1;
+ this.x2 = x2;
+ this.y2 = y2;
+ this.category = category;
+ this.score = score;
+ // -1 stands for area not initialized
+ this.area = -1;
+ }
+
+ // The intersection area of two bounding boxes
+ public float intersect(BoundingBox bbx) {
+ return Math.max(0, Math.min(x2, bbx.x2) - Math.max(x1, bbx.x1))
+ * Math.max(0, Math.min(y2, bbx.y2) - Math.max(y1, bbx.y1));
+ }
+
+ // The union area of two bounding boxes
+ public float union(BoundingBox bbx) {
+ return bbx.getArea() + this.getArea() - this.intersect(bbx);
+ }
+
+ public float getArea() {
+ if (area < 0) {
+ area = (x2 - x1) * (y2 - y1);
+ }
+ return area;
+ }
+
+ public float computeIoU(BoundingBox bbx) {
+ return (float) (this.intersect(bbx) * 1.0 / this.union(bbx));
+ }
+}
diff --git a/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicBenchmarker.java b/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicBenchmarker.java
index 4cda258..15d9511 100644
--- a/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicBenchmarker.java
+++ b/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicBenchmarker.java
@@ -20,11 +20,10 @@
import java.io.IOException;
import java.io.InputStream;
import java.nio.ByteBuffer;
-import java.nio.ByteOrder;
import java.nio.MappedByteBuffer;
/**
- * Class that benchmarks image classifier models.
+ * Base class that benchmarks image models.
*
* <p>===================== General workflow =======================
*
@@ -33,37 +32,40 @@
* benchmarker.getReadyToTest(labelInputStream, model);
* while (!benchmarker.shouldStop()) {
* Bitmap bitmap = ...
- * benchmarker.doTestIteration(bitmap);
+ * imgId = ...
+ * benchmarker.processBitmap(bitmap, imgId);
* }
* }</pre>
*/
-public class OvicBenchmarker {
+public abstract class OvicBenchmarker {
/** Tag for the {@link Log}. */
private static final String TAG = "OvicBenchmarker";
- /** Evaluation transformation parameters. */
- private static final float CENTRAL_FRACTION = 0.875f;
-
/** Dimensions of inputs. */
- private static final int DIM_BATCH_SIZE = 1;
- private static final int DIM_PIXEL_SIZE = 3;
- private int imgHeight = 224;
- private int imgWidth = 224;
+ protected static final int DIM_BATCH_SIZE = 1;
+ protected static final int DIM_PIXEL_SIZE = 3;
+ protected int imgHeight = 224;
+ protected int imgWidth = 224;
+
+ /** Preprocess parameters (only used when input is float). */
+ protected static final float IMAGE_MEAN = 127.5f;
+ protected static final float IMAGE_STD = 127.5f;
+
+ /** Whether input is float or quantized. */
+ protected Boolean quantizedInput = null;
/* Preallocated buffers for storing image data in. */
- private int[] intValues = null;
+ protected int[] intValues = null;
/** A ByteBuffer to hold image data, to be feed into classifier as inputs. */
- private ByteBuffer imgData = null;
-
- private OvicClassifier classifier;
+ protected ByteBuffer imgData = null;
/** Total runtime in ms. */
- private double totalRuntime = 0.0;
+ protected double totalRuntime = 0.0;
/** Total allowed runtime in ms. */
- private double wallTime = 20000 * 30.0;
-
- private Boolean benchmarkStarted = null;
+ protected double wallTime = 20000 * 30.0;
+ /** Record whether benchmark has started (used to skip the first image). */
+ protected boolean benchmarkStarted = false;
/**
* Initializes an {@link OvicBenchmarker}
@@ -76,6 +78,11 @@
this.wallTime = wallTime;
}
+ /** Return the cumulative latency of all runs so far. */
+ public double getTotalRunTime() {
+ return totalRuntime;
+ }
+
/** Check whether the benchmarker should stop. */
public Boolean shouldStop() {
if (totalRuntime >= wallTime) {
@@ -90,105 +97,62 @@
return false;
}
- /** Check whether the benchmarker is ready to start classifying images. */
- public Boolean readyToTest() {
- return (classifier != null);
- }
+ /** Abstract class for checking whether the benchmarker is ready to start processing images */
+ public abstract boolean readyToTest();
/**
- * Getting the benchmarker ready for classifying images.
+ * Abstract class for getting the benchmarker ready.
*
* @param labelInputStream: an {@link InputStream} specifying where the list of labels should be
* read from.
* @param model: a {@link MappedByteBuffer} model to benchmark.
*/
- public void getReadyToTest(InputStream labelInputStream, MappedByteBuffer model) {
- try {
- Log.i(TAG, "Creating classifier.");
- classifier = new OvicClassifier(labelInputStream, model);
- int [] inputDims = classifier.getInputDims();
- imgHeight = inputDims[1];
- imgWidth = inputDims[2];
- // Only accept QUANTIZED_UINT8 input.
- imgData = ByteBuffer.allocateDirect(DIM_BATCH_SIZE * imgHeight * imgWidth * DIM_PIXEL_SIZE);
- imgData.order(ByteOrder.nativeOrder());
- intValues = new int[imgHeight * imgWidth];
- } catch (Exception e) {
- Log.e(TAG, e.getMessage());
- Log.e(TAG, "Failed to initialize ImageNet classifier for the benchmarker.");
- }
- }
-
- /** Return how many classes are predicted per image. */
- public int getNumPredictions() {
- return classifier.getNumPredictions();
- }
+ public abstract void getReadyToTest(InputStream labelInputStream, MappedByteBuffer model);
/**
* Perform test on a single bitmap image.
*
- * @param bitmap: a {@link Bitmap} image to classify.
+ * @param bitmap: a {@link Bitmap} image to process.
+ * @param imageId: an ID uniquely representing the image.
*/
- public OvicSingleImageResult doTestIteration(Bitmap bitmap)
- throws IOException, InterruptedException {
- if (shouldStop() || !readyToTest()) {
- return null;
- }
- OvicSingleImageResult iterResult = null;
- try {
- Log.i(TAG, "Converting bitmap.");
- convertBitmapToInput(bitmap);
- Log.i(TAG, "Classifying image.");
- iterResult = classifier.classifyByteBuffer(imgData);
- } catch (RuntimeException e) {
- Log.e(TAG, e.getMessage());
- Log.e(TAG, "Failed to classify image.");
- }
- if (iterResult == null || iterResult.latency == null) {
- throw new RuntimeException("Classification result or timing is invalid.");
- }
- Log.d(TAG, "Native inference latency: " + iterResult.latency);
- Log.i(TAG, iterResult.toString());
+ public abstract boolean processBitmap(Bitmap bitmap, int imageId)
+ throws IOException, InterruptedException;
- if (!benchmarkStarted) { // Skip the first image to discount warming-up time.
- benchmarkStarted = true;
- } else {
- totalRuntime += (double) iterResult.latency;
- }
- return iterResult;
+ /** Perform test on a single bitmap image without an image ID. */
+ public boolean processBitmap(Bitmap bitmap) throws IOException, InterruptedException {
+ return processBitmap(bitmap, /* imageId = */ 0);
}
+ /** Returns the last inference results as string. */
+ public abstract String getLastResultString();
+
/**
- * Writes Image data into a {@link ByteBuffer}.
- *
- * @param bitmap: a {@link Bitmap} source image.
- */
- private void convertBitmapToInput(Bitmap bitmap) throws RuntimeException {
- if (imgData == null) {
+ * Loads input buffer from intValues into ByteBuffer for the interpreter.
+ * Input buffer must be loaded in intValues and output will be placed in imgData.
+ */
+ protected void loadsInputToByteBuffer() {
+ if (imgData == null || intValues == null || quantizedInput == null) {
throw new RuntimeException("Benchmarker is not yet ready to test.");
}
- imgData.rewind();
- // Perform transformations corresponding to evaluation mode.
- float width = (float) bitmap.getWidth();
- float height = (float) bitmap.getHeight();
- int stWidth = Math.round((width - width * CENTRAL_FRACTION) / 2);
- int stHeight = Math.round((height - height * CENTRAL_FRACTION) / 2);
- int newWidth = Math.round(width - stWidth * 2);
- int newHeight = Math.round(height - stHeight * 2);
- bitmap = Bitmap.createBitmap(bitmap, stWidth, stHeight, newWidth, newHeight);
- bitmap = Bitmap.createScaledBitmap(bitmap, imgWidth, imgHeight, true);
- bitmap.getPixels(intValues, 0, bitmap.getWidth(), 0, 0, bitmap.getWidth(), bitmap.getHeight());
-
// Convert the image to ByteBuffer.
+ imgData.rewind();
int pixel = 0;
long startTime = SystemClock.uptimeMillis();
for (int i = 0; i < imgHeight; ++i) {
for (int j = 0; j < imgWidth; ++j) {
- final int val = intValues[pixel++];
- imgData.put((byte) ((val >> 16) & 0xFF));
- imgData.put((byte) ((val >> 8) & 0xFF));
- imgData.put((byte) (val & 0xFF));
+ final int pixelValue = intValues[pixel++];
+ if (quantizedInput) {
+ // Quantized model
+ imgData.put((byte) ((pixelValue >> 16) & 0xFF));
+ imgData.put((byte) ((pixelValue >> 8) & 0xFF));
+ imgData.put((byte) (pixelValue & 0xFF));
+ } else {
+ // Float model
+ imgData.putFloat((((pixelValue >> 16) & 0xFF) - IMAGE_MEAN) / IMAGE_STD);
+ imgData.putFloat((((pixelValue >> 8) & 0xFF) - IMAGE_MEAN) / IMAGE_STD);
+ imgData.putFloat(((pixelValue & 0xFF) - IMAGE_MEAN) / IMAGE_STD);
+ }
}
}
long endTime = SystemClock.uptimeMillis();
diff --git a/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicSingleImageResult.java b/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicClassificationResult.java
similarity index 83%
rename from tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicSingleImageResult.java
rename to tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicClassificationResult.java
index 4af9a65..5ab804e 100644
--- a/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicSingleImageResult.java
+++ b/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicClassificationResult.java
@@ -1,4 +1,4 @@
-/*Copyright 2018 Google LLC
+/* 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.
@@ -17,17 +17,17 @@
import java.util.ArrayList;
/** Result class for inference run on a single image. */
-public class OvicSingleImageResult {
+public class OvicClassificationResult {
/** Top K classes and probabilities. */
- public ArrayList<String> topKClasses;
- public ArrayList<Float> topKProbs;
- public ArrayList<Integer> topKIndices;
+ public final ArrayList<String> topKClasses;
+ public final ArrayList<Float> topKProbs;
+ public final ArrayList<Integer> topKIndices;
/** Latency (ms). */
public Long latency;
- OvicSingleImageResult() {
+ OvicClassificationResult() {
topKClasses = new ArrayList<>();
topKProbs = new ArrayList<>();
topKIndices = new ArrayList<>();
diff --git a/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicClassifier.java b/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicClassifier.java
index fd610b0..d8a54c1 100644
--- a/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicClassifier.java
+++ b/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicClassifier.java
@@ -31,7 +31,7 @@
import org.tensorflow.lite.Interpreter;
import org.tensorflow.lite.TestHelper;
-/** Benchmark ImageNet Classifier with Tensorflow Lite. */
+/** Class for running ImageNet classification with a TfLite model. */
public class OvicClassifier {
/** Tag for the {@link Log}. */
@@ -106,7 +106,7 @@
/** Classifies a {@link ByteBuffer} image. */
// @throws RuntimeException if model is uninitialized.
- public OvicSingleImageResult classifyByteBuffer(ByteBuffer imgData) {
+ public OvicClassificationResult classifyByteBuffer(ByteBuffer imgData) {
if (tflite == null) {
throw new RuntimeException(TAG + ": ImageNet classifier has not been initialized; Failed.");
}
@@ -122,7 +122,7 @@
labelProbArray[0][i] = (inferenceOutputArray[0][i] & 0xff) / 255.0f;
}
}
- OvicSingleImageResult iterResult = computeTopKLabels();
+ OvicClassificationResult iterResult = computeTopKLabels();
iterResult.latency = getLastNativeInferenceLatencyMilliseconds();
return iterResult;
}
@@ -174,7 +174,7 @@
}
/** Computes top-K labels. */
- private OvicSingleImageResult computeTopKLabels() {
+ private OvicClassificationResult computeTopKLabels() {
if (labelList == null) {
throw new RuntimeException("Label file has not been loaded.");
}
@@ -184,7 +184,7 @@
sortedLabels.poll();
}
}
- OvicSingleImageResult singleImageResult = new OvicSingleImageResult();
+ OvicClassificationResult singleImageResult = new OvicClassificationResult();
if (sortedLabels.size() != RESULTS_TO_SHOW) {
throw new RuntimeException(
"Number of returned labels does not match requirement: "
diff --git a/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicClassifierBenchmarker.java b/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicClassifierBenchmarker.java
new file mode 100644
index 0000000..0cdd0f7
--- /dev/null
+++ b/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicClassifierBenchmarker.java
@@ -0,0 +1,142 @@
+/* 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
+
+ https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+package org.tensorflow.ovic;
+
+import android.graphics.Bitmap;
+import android.util.Log;
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.nio.MappedByteBuffer;
+
+/** Class that benchmarks image classifier models. */
+public final class OvicClassifierBenchmarker extends OvicBenchmarker {
+ /** Tag for the {@link Log}. */
+ private static final String TAG = "OvicClassifierBenchmarker";
+
+ /** ImageNet preprocessing parameters. */
+ private static final float CENTRAL_FRACTION = 0.875f;
+ private OvicClassifier classifier;
+ private OvicClassificationResult iterResult = null;
+
+ public OvicClassifierBenchmarker(double wallTime) {
+ super(wallTime);
+ }
+
+ /** Test if the classifier is ready for benchmarking. */
+ @Override
+ public boolean readyToTest() {
+ return (classifier != null);
+ }
+
+ /**
+ * Getting the benchmarker ready for classifying images.
+ *
+ * @param labelInputStream: an {@link InputStream} specifying where the list of labels should be
+ * read from.
+ * @param model: a {@link MappedByteBuffer} model to benchmark.
+ */
+ @Override
+ public void getReadyToTest(InputStream labelInputStream, MappedByteBuffer model) {
+ try {
+ Log.i(TAG, "Creating classifier.");
+ classifier = new OvicClassifier(labelInputStream, model);
+ int [] inputDims = classifier.getInputDims();
+ imgHeight = inputDims[1];
+ imgWidth = inputDims[2];
+ quantizedInput = true;
+ // Only accept QUANTIZED_UINT8 input.
+ imgData = ByteBuffer.allocateDirect(DIM_BATCH_SIZE * imgHeight * imgWidth * DIM_PIXEL_SIZE);
+ imgData.order(ByteOrder.nativeOrder());
+ intValues = new int[imgHeight * imgWidth];
+ } catch (Exception e) {
+ Log.e(TAG, e.getMessage());
+ Log.e(TAG, "Failed to initialize ImageNet classifier for the benchmarker.");
+ }
+ }
+
+ /**
+ * Perform classification on a single bitmap image.
+ *
+ * @param bitmap: a {@link Bitmap} image to process.
+ * @param imageId: an ID uniquely representing the image.
+ */
+ @Override
+ public boolean processBitmap(Bitmap bitmap, int imageId)
+ throws IOException, InterruptedException {
+ if (shouldStop() || !readyToTest()) {
+ return false;
+ }
+ try {
+ Log.i(TAG, "Converting bitmap.");
+ convertBitmapToInput(bitmap);
+ Log.i(TAG, "Classifying image: " + imageId);
+ iterResult = classifier.classifyByteBuffer(imgData);
+ } catch (RuntimeException e) {
+ Log.e(TAG, e.getMessage());
+ Log.e(TAG, "Failed to classify image.");
+ }
+ if (iterResult == null || iterResult.latency == null) {
+ throw new RuntimeException("Classification result or timing is invalid.");
+ }
+ Log.d(TAG, "Native inference latency: " + iterResult.latency);
+ Log.i(TAG, iterResult.toString());
+
+ if (!benchmarkStarted) { // Skip the first image to discount warming-up time.
+ benchmarkStarted = true;
+ } else {
+ totalRuntime += ((double) iterResult.latency);
+ }
+ return true;
+ }
+
+ /** Return how many classes are predicted per image. */
+ public int getNumPredictions() {
+ return classifier.getNumPredictions();
+ }
+
+ public OvicClassificationResult getLastClassificationResult() {
+ return iterResult;
+ }
+
+ @Override
+ public String getLastResultString() {
+ if (iterResult == null) {
+ return null;
+ } else {
+ return iterResult.toString();
+ }
+ }
+
+ /**
+ * Preprocess bitmap according to ImageNet protocol then writes result into a {@link ByteBuffer}.
+ *
+ * @param bitmap: a {@link Bitmap} source image.
+ */
+ private void convertBitmapToInput(Bitmap bitmap) {
+ // Perform transformations corresponding to evaluation mode.
+ float width = (float) bitmap.getWidth();
+ float height = (float) bitmap.getHeight();
+ int stWidth = Math.round((width - width * CENTRAL_FRACTION) / 2);
+ int stHeight = Math.round((height - height * CENTRAL_FRACTION) / 2);
+ int newWidth = Math.round(width - stWidth * 2);
+ int newHeight = Math.round(height - stHeight * 2);
+ bitmap = Bitmap.createBitmap(bitmap, stWidth, stHeight, newWidth, newHeight);
+ bitmap = Bitmap.createScaledBitmap(bitmap, imgWidth, imgHeight, true);
+ bitmap.getPixels(intValues, 0, bitmap.getWidth(), 0, 0, bitmap.getWidth(), bitmap.getHeight());
+ loadsInputToByteBuffer();
+ }
+}
diff --git a/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicDetectionResult.java b/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicDetectionResult.java
new file mode 100644
index 0000000..cf2902a
--- /dev/null
+++ b/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicDetectionResult.java
@@ -0,0 +1,91 @@
+/* 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
+
+ https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+package org.tensorflow.ovic;
+
+import java.util.ArrayList;
+
+/** Result class for inference run on a single image. */
+public class OvicDetectionResult {
+
+ // Top K classes and probabilities.
+ public final ArrayList<BoundingBox> detections;
+ // Latency (ms).
+ public Long latency = -1L;
+ // id of the image.
+ public int id = -1;
+ // Number of valid detections (separately maintained, maybe different from detections.size()).
+ public int count = 0;
+
+ // Create OvicDetectionResult object with pre-filled capacity. Note that detections.size() will
+ // be equal to capacity after this call.
+ OvicDetectionResult(int capacity) {
+ detections = new ArrayList<BoundingBox>(capacity);
+ for (int i = 0; i < capacity; i++) {
+ detections.add(new BoundingBox(-1.0f, -1.0f, -1.0f, -1.0f, -1, -1.0f));
+ }
+ }
+
+ public void resetTo(Long latency, int id) {
+ count = 0;
+ this.latency = latency;
+ this.id = id;
+ }
+
+ public void addBox(float x1, float y1, float x2, float y2, int category, float score) {
+ detections.get(count).x1 = x1;
+ detections.get(count).y1 = y1;
+ detections.get(count).x2 = x2;
+ detections.get(count).y2 = y2;
+ detections.get(count).category = category;
+ detections.get(count).score = score;
+ count += 1;
+ }
+
+ public void scaleUp(double scaleFactorWidth, double scaleFactorHeight) {
+ for (BoundingBox box : detections) {
+ box.x1 = (float) (box.x1 * scaleFactorWidth);
+ box.y1 = (float) (box.y1 * scaleFactorHeight);
+ box.x2 = (float) (box.x2 * scaleFactorWidth);
+ box.y2 = (float) (box.y2 * scaleFactorHeight);
+ }
+ }
+
+ @Override
+ public String toString() {
+ String textToShow = latency + "ms";
+ int k = 0;
+ for (BoundingBox box : detections) {
+ textToShow +=
+ "\nPrediction ["
+ + k
+ + "] = Class "
+ + box.category
+ + " ("
+ + box.x1
+ + ", "
+ + box.y1
+ + ", "
+ + box.x2
+ + ", "
+ + box.y2
+ + ") : "
+ + box.score;
+ k++;
+ }
+
+
+ return textToShow;
+ }
+}
diff --git a/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicDetector.java b/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicDetector.java
new file mode 100644
index 0000000..56836a7
--- /dev/null
+++ b/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicDetector.java
@@ -0,0 +1,184 @@
+/* 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
+
+ https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+package org.tensorflow.ovic;
+
+import java.io.BufferedReader;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.InputStreamReader;
+import java.nio.ByteBuffer;
+import java.nio.MappedByteBuffer;
+import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.tensorflow.lite.Interpreter;
+import org.tensorflow.lite.TestHelper;
+
+/** Class for running COCO detection with a TfLite model. */
+public class OvicDetector implements AutoCloseable {
+
+ /** Tag for the {@link Log}. */
+ private static final String TAG = "OvicDetector";
+
+ /** An instance of the driver class to run model inference with Tensorflow Lite. */
+ private Interpreter tflite;
+
+ /** Labels corresponding to the output of the vision model. */
+ private final List<String> labelList;
+
+ /** Define the output format. */
+ private final Boolean inputIsFloat;
+
+ /** Number of detections per image. 10 for demo, 100 for the actual competition. */
+ private static final int NUM_RESULTS = 10;
+
+ /** The output arrays for the mobilenet SSD. */
+ private float[][][] outputLocations;
+ private float[][] outputClasses;
+ private float[][] outputScores;
+ private float[] numDetections;
+ private Map<Integer, Object> outputMap;
+
+ /** Input resolution. */
+ private final int[] inputDims;
+
+ /** Final result. */
+ public OvicDetectionResult result = null;
+
+ OvicDetector(InputStream labelInputStream, MappedByteBuffer model) throws IOException {
+ // Load the label list.
+ labelList = loadLabelList(labelInputStream);
+
+ // Create the TfLite interpreter.
+ tflite = new Interpreter(model, new Interpreter.Options().setNumThreads(1));
+ inputDims = TestHelper.getInputDims(tflite, 0);
+ inputIsFloat = TestHelper.getInputDataType(tflite, 0).equals("float");
+ if (inputDims.length != 4) {
+ throw new RuntimeException("The model's input dimensions must be 4 (BWHC).");
+ }
+ if (inputDims[0] != 1) {
+ throw new RuntimeException(
+ "The model must have a batch size of 1, got " + inputDims[0] + " instead.");
+ }
+ if (inputDims[3] != 3) {
+ throw new RuntimeException(
+ "The model must have three color channels, got " + inputDims[3] + " instead.");
+ }
+ // Check the resolution.
+ int minSide = Math.min(inputDims[1], inputDims[2]);
+ int maxSide = Math.max(inputDims[1], inputDims[2]);
+ if (minSide <= 0 || maxSide > 1000) {
+ throw new RuntimeException("The model's resolution must be between (0, 1000].");
+ }
+
+ // Initialize the input array and result arrays. The input images are stored in a list of
+ // Object. Since this function anaylzed one image per time, there is only 1 item.
+ // The output is fomulated as a map of int -> Object. The output arrays are added to the map.
+ outputLocations = new float[1][NUM_RESULTS][4];
+ outputClasses = new float[1][NUM_RESULTS];
+ outputScores = new float[1][NUM_RESULTS];
+ numDetections = new float[1];
+ outputMap = new HashMap<>();
+ outputMap.put(0, outputLocations);
+ outputMap.put(1, outputClasses);
+ outputMap.put(2, outputScores);
+ outputMap.put(3, numDetections);
+ // Preallocate the result. This will be where inference result is stored after each
+ // detectByteBuffer call.
+ result = new OvicDetectionResult(NUM_RESULTS);
+ }
+
+ public Boolean quantizedInput() {
+ return !inputIsFloat;
+ }
+
+ /** Reads label list from Assets. */
+ private static List<String> loadLabelList(InputStream labelInputStream) throws IOException {
+ List<String> labelList = new ArrayList<>();
+ try (BufferedReader reader =
+ new BufferedReader(new InputStreamReader(labelInputStream, StandardCharsets.UTF_8))) {
+ String line;
+ while ((line = reader.readLine()) != null) {
+ labelList.add(line);
+ }
+ }
+ return labelList;
+ }
+
+ /**
+ * The interface to run the detection. This method currently only support float mobilenet_ssd
+ * model. The quantized models will be added in the future.
+ *
+ * @param imgData The image buffer in ByteBuffer format.
+ * @return boolean indicator of whether detection was a success. If success, the detection results
+ * is available in the result member variable.
+ * See OvicDetectionResult.java for details.
+ */
+ boolean detectByteBuffer(ByteBuffer imgData, int imageId) {
+ if (tflite == null) {
+ throw new RuntimeException(TAG + ": Detector has not been initialized; Failed.");
+ }
+ if (inputIsFloat == null) {
+ throw new RuntimeException(TAG + ": Detector input type has not been resolved.");
+ }
+
+ Object[] inputArray = {imgData};
+ tflite.runForMultipleInputsOutputs(inputArray, outputMap);
+
+ Long latency = getLastNativeInferenceLatencyMilliseconds();
+
+ // Update the results.
+ result.resetTo(latency, imageId);
+ for (int i = 0; i < NUM_RESULTS; i++) {
+ result.addBox(outputLocations[0][i][1] * inputDims[1],
+ outputLocations[0][i][0] * inputDims[1],
+ outputLocations[0][i][3] * inputDims[2],
+ outputLocations[0][i][2] * inputDims[2],
+ Math.round(outputClasses[0][i] + 1 /* Label offset */),
+ outputScores[0][i]);
+ }
+ return true; // Marks that the result is available.
+ }
+
+ /*
+ * Get native inference latency of last image detection run.
+ * @throws RuntimeException if model is uninitialized.
+ * @return The inference latency in millisecond.
+ */
+ public Long getLastNativeInferenceLatencyMilliseconds() {
+ if (tflite == null) {
+ throw new RuntimeException(TAG + ": ImageNet classifier has not been initialized; Failed.");
+ }
+ Long latency = tflite.getLastNativeInferenceDurationNanoseconds();
+ return (latency == null) ? null : (Long) (latency / 1000000);
+ }
+
+ public int[] getInputDims() {
+ return inputDims;
+ }
+
+ public List<String> getLabels() {
+ return labelList;
+ }
+
+ /** Closes tflite to release resources. */
+ @Override
+ public void close() {
+ tflite.close();
+ tflite = null;
+ }
+}
diff --git a/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicDetectorBenchmarker.java b/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicDetectorBenchmarker.java
new file mode 100644
index 0000000..1a4e193
--- /dev/null
+++ b/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicDetectorBenchmarker.java
@@ -0,0 +1,160 @@
+/* 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
+
+ https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+package org.tensorflow.ovic;
+
+import android.graphics.Bitmap;
+import android.util.Log;
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.nio.MappedByteBuffer;
+
+/**
+ * Class that benchmarks object detection models.
+ */
+public final class OvicDetectorBenchmarker extends OvicBenchmarker {
+ /** Tag for the {@link Log}. */
+ private static final String TAG = "OvicDetectorBenchmarker";
+
+ public double scaleFactorWidth = 1.0f;
+ public double scaleFactorHeight = 1.0f;
+ private Bitmap scaledBitmap = null; // Preallocate bitmap for scaling.
+
+ private OvicDetector detector;
+
+ /**
+ * Initializes an {@link OvicDetectionBenchmarker}
+ *
+ * @param wallTime: a double number specifying the total amount of time to benchmark.
+ */
+ public OvicDetectorBenchmarker(double wallTime) {
+ super(wallTime);
+ }
+
+ /** Check to see if the detector is ready to test. */
+ @Override
+ public boolean readyToTest() {
+ return (detector != null);
+ }
+
+ /**
+ * Getting the benchmarker ready for detecting images.
+ *
+ * @param labelInputStream: an {@link InputStream} specifying where the list of labels should be
+ * read from.
+ * @param model: a {@link MappedByteBuffer} model to benchmark.
+ */
+ @Override
+ public void getReadyToTest(InputStream labelInputStream, MappedByteBuffer model) {
+ try {
+ Log.i(TAG, "Creating detector.");
+ detector = new OvicDetector(labelInputStream, model);
+ quantizedInput = detector.quantizedInput();
+ int[] inputDims = detector.getInputDims();
+ imgHeight = inputDims[1];
+ imgWidth = inputDims[2];
+ if (quantizedInput) {
+ imgData = ByteBuffer.allocateDirect(DIM_BATCH_SIZE * imgHeight * imgWidth * DIM_PIXEL_SIZE);
+ } else {
+ imgData =
+ ByteBuffer.allocateDirect(DIM_BATCH_SIZE * imgHeight * imgWidth * DIM_PIXEL_SIZE * 4);
+ }
+ imgData.order(ByteOrder.nativeOrder());
+ intValues = new int[imgHeight * imgWidth];
+ benchmarkStarted = false;
+ } catch (Exception e) {
+ Log.e(TAG, e.getMessage());
+ Log.e(TAG, "Failed to initialize COCO detector for the benchmarker.", e);
+ }
+ }
+
+ /**
+ * Perform detection on a single ByteBuffer {@link ByteBuffer} image. The image must have the
+ * same dimension that the model expects.
+ *
+ * @param image: a {@link ByteBuffer} image to process.
+ * @param imageId: an ID uniquely representing the image.
+ */
+ public boolean processBuffer(ByteBuffer image, int imageId) {
+ if (!readyToTest()) {
+ return false;
+ }
+ try {
+ if (!detector.detectByteBuffer(image, imageId)) {
+ return false;
+ }
+ } catch (RuntimeException e) {
+ Log.e(TAG, e.getMessage());
+ return false;
+ }
+
+ if (!benchmarkStarted) { // Skip the first image to discount warming-up time.
+ benchmarkStarted = true;
+ } else {
+ totalRuntime += ((double) detector.result.latency);
+ }
+ return true; // Indicating that result is ready.
+ }
+
+ /**
+ * Perform detection on a single bitmap image.
+ *
+ * @param bitmap: a {@link Bitmap} image to process.
+ * @param imageId: an ID uniquely representing the image.
+ */
+ @Override
+ public boolean processBitmap(Bitmap bitmap, int imageId)
+ throws IOException, InterruptedException {
+ if (shouldStop() || !readyToTest()) {
+ return false;
+ }
+ convertBitmapToInput(bitmap); // Scale bitmap if needed, store result in imgData.
+ if (!processBuffer(imgData, imageId)) {
+ return false;
+ }
+ // Scale results back to original image coordinates.
+ detector.result.scaleUp(scaleFactorWidth, scaleFactorHeight);
+ return true; // Indicating that result is ready.
+ }
+
+ public OvicDetectionResult getLastDetectionResult() {
+ return detector.result;
+ }
+
+ @Override
+ public String getLastResultString() {
+ if (detector.result == null) {
+ return null;
+ }
+ return detector.result.toString();
+ }
+
+ /**
+ * Preprocess bitmap image into {@link ByteBuffer} format for the detector.
+ *
+ * @param bitmap: a {@link Bitmap} source image.
+ */
+ private void convertBitmapToInput(Bitmap bitmap) {
+ int originalWidth = bitmap.getWidth();
+ int originalHeight = bitmap.getHeight();
+ scaledBitmap = Bitmap.createScaledBitmap(bitmap, imgWidth, imgHeight, true);
+ scaleFactorWidth = originalWidth * 1.0 / imgWidth;
+ scaleFactorHeight = originalHeight * 1.0 / imgHeight;
+ scaledBitmap.getPixels(intValues, 0, imgWidth, 0, 0, imgWidth, imgHeight);
+ scaledBitmap.recycle();
+ loadsInputToByteBuffer();
+ }
+}
diff --git a/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicValidator.java b/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicValidator.java
index a504ec7..baa14ba 100644
--- a/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicValidator.java
+++ b/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicValidator.java
@@ -51,7 +51,7 @@
MappedByteBuffer model = loadModelFile(modelFile);
OvicClassifier classifier = new OvicClassifier(labelsInputStream, model);
ByteBuffer imgData = createByteBufferForClassifier(classifier);
- OvicSingleImageResult testResult = classifier.classifyByteBuffer(imgData);
+ OvicClassificationResult testResult = classifier.classifyByteBuffer(imgData);
if (testResult.topKClasses.isEmpty()) {
throw new RuntimeException("Failed to return top K predictions.");
}
diff --git a/tensorflow/contrib/lite/java/ovic/src/test/java/org/tensorflow/ovic/OvicClassifierTest.java b/tensorflow/contrib/lite/java/ovic/src/test/java/org/tensorflow/ovic/OvicClassifierTest.java
index 1587c3c..99e874c 100644
--- a/tensorflow/contrib/lite/java/ovic/src/test/java/org/tensorflow/ovic/OvicClassifierTest.java
+++ b/tensorflow/contrib/lite/java/ovic/src/test/java/org/tensorflow/ovic/OvicClassifierTest.java
@@ -1,4 +1,4 @@
-/*Copyright 2018 Google LLC
+/* 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.
@@ -43,7 +43,7 @@
private MappedByteBuffer lowResModel = null;
private ByteBuffer testImage = null;
private ByteBuffer lowResTestImage = null;
- private OvicSingleImageResult testResult = null;
+ private OvicClassificationResult testResult = null;
private static final String LABELS_PATH =
"tensorflow/contrib/lite/java/ovic/src/testdata/labels.txt";
private static final String QUANTIZED_MODEL_PATH =
@@ -147,7 +147,7 @@
return imgData;
}
- private static void assertCorrectTopK(OvicSingleImageResult testResult) {
+ private static void assertCorrectTopK(OvicClassificationResult testResult) {
assertThat(testResult.topKClasses.size() > 0).isTrue();
Boolean topKAccurate = false;
// Assert that the correct class is in the top K.
diff --git a/tensorflow/contrib/lite/java/ovic/src/test/java/org/tensorflow/ovic/OvicDetectorTest.java b/tensorflow/contrib/lite/java/ovic/src/test/java/org/tensorflow/ovic/OvicDetectorTest.java
new file mode 100644
index 0000000..4681e26
--- /dev/null
+++ b/tensorflow/contrib/lite/java/ovic/src/test/java/org/tensorflow/ovic/OvicDetectorTest.java
@@ -0,0 +1,149 @@
+/* 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
+
+ https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+package org.tensorflow.ovic;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import java.awt.Graphics2D;
+import java.awt.image.BufferedImage;
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.nio.MappedByteBuffer;
+import java.nio.channels.FileChannel;
+import javax.imageio.ImageIO;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Unit test for {@link org.tensorflow.ovic.OvicDetector}. */
+@RunWith(JUnit4.class)
+public final class OvicDetectorTest {
+ private OvicDetector detector = null;
+ private InputStream labelsInputStream = null;
+ private MappedByteBuffer model = null;
+ private ByteBuffer testImage = null;
+
+ private static final float IMAGE_MEAN = 128f;
+ private static final float IMAGE_STD = 128f;
+
+ private Boolean quantizedInput = null;
+ private static final String LABELS_PATH =
+ "tensorflow/contrib/lite/java/ovic/src/testdata/coco_labels.txt";
+ private static final String MODEL_PATH =
+ "external/tflite_mobilenet_ssd_quant/detect.tflite";
+ private static final String TEST_IMAGE_PATH =
+ "external/tflite_ovic_testdata/test_image_224.jpg";
+ private static final int GROUNDTRUTH = 1 /* Person */;
+
+ @Before
+ public void setUp() {
+ try {
+ // load models.
+ model = loadModelFile(MODEL_PATH);
+
+ // Load label files;
+ File labelsfile = new File(LABELS_PATH);
+ labelsInputStream = new FileInputStream(labelsfile);
+
+ // Create detector.
+ detector = new OvicDetector(labelsInputStream, model);
+ quantizedInput = detector.quantizedInput();
+
+ // Load test image and convert into byte buffer.
+ File imageFile = new File(TEST_IMAGE_PATH);
+ BufferedImage rawimg = ImageIO.read(imageFile);
+ int[] inputDims = detector.getInputDims();
+ BufferedImage img = new BufferedImage(inputDims[1], inputDims[2], rawimg.getType());
+ Graphics2D g = img.createGraphics();
+ g.drawImage(rawimg, 0, 0, inputDims[1], inputDims[2], null);
+ g.dispose();
+ testImage = toByteBuffer(img);
+ } catch (IOException e) {
+ System.out.println(e.getMessage());
+ }
+
+ System.out.println("Successfully setup");
+ }
+
+ private static MappedByteBuffer loadModelFile(String modelFilePath) throws IOException {
+ File modelfile = new File(modelFilePath);
+ FileInputStream inputStream = new FileInputStream(modelfile);
+ FileChannel fileChannel = inputStream.getChannel();
+ long startOffset = 0L;
+ long declaredLength = fileChannel.size();
+ return fileChannel.map(FileChannel.MapMode.READ_ONLY, startOffset, declaredLength);
+ }
+
+ private ByteBuffer toByteBuffer(BufferedImage image) {
+ ByteBuffer imgData;
+ if (quantizedInput) {
+ imgData = ByteBuffer.allocateDirect(image.getHeight() * image.getWidth() * 3);
+ } else {
+ imgData = ByteBuffer.allocateDirect(image.getHeight() * image.getWidth() * 12);
+ }
+ imgData.order(ByteOrder.nativeOrder());
+ for (int y = 0; y < image.getHeight(); y++) {
+ for (int x = 0; x < image.getWidth(); x++) {
+ int pixelValue = image.getRGB(x, y);
+ if (quantizedInput) {
+ // Quantized model
+ imgData.put((byte) ((pixelValue >> 16) & 0xFF));
+ imgData.put((byte) ((pixelValue >> 8) & 0xFF));
+ imgData.put((byte) (pixelValue & 0xFF));
+ } else {
+ // Float model
+ imgData.putFloat((((pixelValue >> 16) & 0xFF) - IMAGE_MEAN) / IMAGE_STD);
+ imgData.putFloat((((pixelValue >> 8) & 0xFF) - IMAGE_MEAN) / IMAGE_STD);
+ imgData.putFloat(((pixelValue & 0xFF) - IMAGE_MEAN) / IMAGE_STD);
+ }
+ }
+ }
+ return imgData;
+ }
+
+ @Test
+ public void ovicDetector_detectSuccess() throws Exception {
+ assertThat(detector.detectByteBuffer(testImage, 1)).isTrue();
+ assertThat(detector.result != null).isTrue();
+ }
+
+ @Test
+ public void ovicDetector_simpleBatchTest() throws Exception {
+ final int numRepeats = 5;
+ for (int i = 0; i < numRepeats; i++) {
+ assertThat(detector.detectByteBuffer(testImage, 1)).isTrue();
+ OvicDetectionResult result = detector.result;
+ Boolean detectWithinTop5 = false;
+ for (int j = 0; j < Math.min(5, result.count); j++) {
+ if (result.detections.get(j).category == GROUNDTRUTH) {
+ detectWithinTop5 = true;
+ break;
+ }
+ }
+ if (!detectWithinTop5) {
+ System.out.println("---------------- Image " + i + " ---------------------");
+ System.out.println("Expect category " + GROUNDTRUTH);
+ System.out.println("Detection results: ");
+ System.out.println(result.toString());
+ }
+ assertThat(detectWithinTop5).isTrue();
+ }
+ }
+}
diff --git a/tensorflow/contrib/lite/java/ovic/src/testdata/BUILD b/tensorflow/contrib/lite/java/ovic/src/testdata/BUILD
index 1021ea3..051aa22 100644
--- a/tensorflow/contrib/lite/java/ovic/src/testdata/BUILD
+++ b/tensorflow/contrib/lite/java/ovic/src/testdata/BUILD
@@ -14,6 +14,9 @@
)
exports_files(
- ["labels.txt"],
+ [
+ "labels.txt",
+ "coco_labels.txt",
+ ],
visibility = ["//visibility:public"],
)
diff --git a/tensorflow/contrib/lite/java/ovic/src/testdata/coco_labels.txt b/tensorflow/contrib/lite/java/ovic/src/testdata/coco_labels.txt
new file mode 100644
index 0000000..d91f535
--- /dev/null
+++ b/tensorflow/contrib/lite/java/ovic/src/testdata/coco_labels.txt
@@ -0,0 +1,91 @@
+person
+bicycle
+car
+motorcycle
+airplane
+bus
+train
+truck
+boat
+traffic light
+fire hydrant
+empty
+stop sign
+parking meter
+bench
+bird
+cat
+dog
+horse
+sheep
+cow
+elephant
+bear
+zebra
+giraffe
+empty
+backpack
+umbrella
+empty
+empty
+handbag
+tie
+suitcase
+frisbee
+skis
+snowboard
+sports ball
+kite
+baseball bat
+baseball glove
+skateboard
+surfboard
+tennis racket
+bottle
+empty
+wine glasses
+cup
+fork
+knife
+spoon
+bowl
+banana
+apple
+sandwich
+orange
+broccoli
+carrot
+hot dog
+pizza
+donut
+cake
+chair
+couch
+potted plant
+bed
+empty
+dining table
+empty
+empty
+toilet
+empty
+tv
+laptop
+mouse
+remote
+keyboard
+cell phone
+microwave
+oven
+toaster
+sink
+refrigerator
+empty
+book
+clock
+vase
+scissors
+teddy bear
+hair drier
+toothbrush
+empty
diff --git a/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/Interpreter.java b/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/Interpreter.java
index ffb0449..5cc6e75 100644
--- a/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/Interpreter.java
+++ b/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/Interpreter.java
@@ -17,6 +17,7 @@
import java.io.File;
import java.nio.ByteBuffer;
+import java.nio.MappedByteBuffer;
import java.util.HashMap;
import java.util.Map;
import org.checkerframework.checker.nullness.qual.NonNull;
@@ -74,8 +75,19 @@
return this;
}
+ /**
+ * Sets whether to allow float16 precision for FP32 calculation when possible. Defaults to false
+ * (disallow).
+ * WARNING: This is an experimental API and subject to change.
+ */
+ public Options setAllowFp16PrecisionForFp32(boolean allow) {
+ this.allowFp16PrecisionForFp32 = allow;
+ return this;
+ }
+
int numThreads = -1;
boolean useNNAPI = false;
+ boolean allowFp16PrecisionForFp32 = false;
}
/**
@@ -138,6 +150,20 @@
}
/**
+ * Initializes a {@code Interpreter} with a {@code MappedByteBuffer} to the model file.
+ *
+ * <p>The {@code MappedByteBuffer} should remain unchanged after the construction of a {@code
+ * Interpreter}.
+ *
+ * @deprecated Prefer using the {@link #Interpreter(ByteBuffer,Options)} constructor. This method
+ * will be removed in a future release.
+ */
+ @Deprecated
+ public Interpreter(@NonNull MappedByteBuffer mappedByteBuffer) {
+ this(mappedByteBuffer, /* options= */ null);
+ }
+
+ /**
* Initializes a {@code Interpreter} with a {@code ByteBuffer} of a model file and a set of custom
* {@link #Options}.
*
@@ -256,8 +282,9 @@
/**
* Returns native inference timing.
- * <p>IllegalArgumentException will be thrown if the model is not initialized by the
- * {@link Interpreter}.
+ *
+ * <p>IllegalArgumentException will be thrown if the model is not initialized by the {@link
+ * Interpreter}.
*/
public Long getLastNativeInferenceDurationNanoseconds() {
checkNotClosed();
diff --git a/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/NativeInterpreterWrapper.java b/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/NativeInterpreterWrapper.java
index 6feff9a..9bc44bf 100644
--- a/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/NativeInterpreterWrapper.java
+++ b/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/NativeInterpreterWrapper.java
@@ -45,6 +45,9 @@
isMemoryAllocated = true;
inputTensors = new Tensor[getInputCount(interpreterHandle)];
outputTensors = new Tensor[getOutputCount(interpreterHandle)];
+ if (options.allowFp16PrecisionForFp32) {
+ setAllowFp16PrecisionForFp32(options.allowFp16PrecisionForFp32);
+ }
}
NativeInterpreterWrapper(ByteBuffer byteBuffer) {
@@ -72,6 +75,9 @@
if (options.useNNAPI) {
setUseNNAPI(options.useNNAPI);
}
+ if (options.allowFp16PrecisionForFp32) {
+ setAllowFp16PrecisionForFp32(options.allowFp16PrecisionForFp32);
+ }
}
/** Releases resources associated with this {@code NativeInterpreterWrapper}. */
@@ -159,6 +165,10 @@
useNNAPI(interpreterHandle, useNNAPI);
}
+ void setAllowFp16PrecisionForFp32(boolean allow) {
+ allowFp16PrecisionForFp32(interpreterHandle, allow);
+ }
+
void setNumThreads(int numThreads) {
numThreads(interpreterHandle, numThreads);
}
@@ -323,6 +333,8 @@
private static native void numThreads(long interpreterHandle, int numThreads);
+ private static native void allowFp16PrecisionForFp32(long interpreterHandle, boolean allow);
+
private static native long createErrorReporter(int size);
private static native long createModel(String modelPathOrBuffer, long errorHandle);
diff --git a/tensorflow/contrib/lite/java/src/main/native/nativeinterpreterwrapper_jni.cc b/tensorflow/contrib/lite/java/src/main/native/nativeinterpreterwrapper_jni.cc
index fdcf00a..abb7320 100644
--- a/tensorflow/contrib/lite/java/src/main/native/nativeinterpreterwrapper_jni.cc
+++ b/tensorflow/contrib/lite/java/src/main/native/nativeinterpreterwrapper_jni.cc
@@ -59,7 +59,6 @@
return outputs;
}
-
int getDataType(TfLiteType data_type) {
switch (data_type) {
case kTfLiteFloat32:
@@ -234,10 +233,18 @@
}
JNIEXPORT void JNICALL
+Java_org_tensorflow_lite_NativeInterpreterWrapper_allowFp16PrecisionForFp32(
+ JNIEnv* env, jclass clazz, jlong handle, jboolean allow) {
+ tflite::Interpreter* interpreter = convertLongToInterpreter(env, handle);
+ if (interpreter == nullptr) return;
+ interpreter->SetAllowFp16PrecisionForFp32(static_cast<bool>(allow));
+}
+
+JNIEXPORT void JNICALL
Java_org_tensorflow_lite_NativeInterpreterWrapper_numThreads(JNIEnv* env,
- jclass clazz,
- jlong handle,
- jint num_threads) {
+ jclass clazz,
+ jlong handle,
+ jint num_threads) {
tflite::Interpreter* interpreter = convertLongToInterpreter(env, handle);
if (interpreter == nullptr) return;
interpreter->SetNumThreads(static_cast<int>(num_threads));
diff --git a/tensorflow/contrib/lite/java/src/main/native/nativeinterpreterwrapper_jni.h b/tensorflow/contrib/lite/java/src/main/native/nativeinterpreterwrapper_jni.h
index 06b35d7..aa809df 100644
--- a/tensorflow/contrib/lite/java/src/main/native/nativeinterpreterwrapper_jni.h
+++ b/tensorflow/contrib/lite/java/src/main/native/nativeinterpreterwrapper_jni.h
@@ -120,6 +120,15 @@
/*
* Class: org_tensorflow_lite_NativeInterpreterWrapper
* Method:
+ * Signature: (JZ)V
+ */
+JNIEXPORT void JNICALL
+Java_org_tensorflow_lite_NativeInterpreterWrapper_allowFp16PrecisionForFp32(
+ JNIEnv* env, jclass clazz, jlong handle, jboolean allow);
+
+/*
+ * Class: org_tensorflow_lite_NativeInterpreterWrapper
+ * Method:
* Signature: (JI)V
*/
JNIEXPORT void JNICALL
diff --git a/tensorflow/contrib/lite/java/src/test/java/org/tensorflow/lite/InterpreterTest.java b/tensorflow/contrib/lite/java/src/test/java/org/tensorflow/lite/InterpreterTest.java
index dfdd7d2..a98fca0 100644
--- a/tensorflow/contrib/lite/java/src/test/java/org/tensorflow/lite/InterpreterTest.java
+++ b/tensorflow/contrib/lite/java/src/test/java/org/tensorflow/lite/InterpreterTest.java
@@ -71,7 +71,7 @@
Path path = MODEL_FILE.toPath();
FileChannel fileChannel =
(FileChannel) Files.newByteChannel(path, EnumSet.of(StandardOpenOption.READ));
- MappedByteBuffer mappedByteBuffer =
+ ByteBuffer mappedByteBuffer =
fileChannel.map(FileChannel.MapMode.READ_ONLY, 0, fileChannel.size());
Interpreter interpreter = new Interpreter(mappedByteBuffer);
float[] oneD = {1.23f, 6.54f, 7.81f};
@@ -118,7 +118,7 @@
byteBuffer.order(ByteOrder.nativeOrder());
fileChannel.read(byteBuffer);
try {
- Interpreter interpreter = new Interpreter(byteBuffer);
+ new Interpreter(byteBuffer);
fail();
} catch (IllegalArgumentException e) {
assertThat(e)
@@ -323,7 +323,9 @@
MappedByteBuffer mappedByteBuffer =
fileChannel.map(FileChannel.MapMode.READ_ONLY, 0, fileChannel.size());
Interpreter interpreter =
- new Interpreter(mappedByteBuffer, new Interpreter.Options().setUseNNAPI(true));
+ new Interpreter(
+ mappedByteBuffer,
+ new Interpreter.Options().setUseNNAPI(true).setAllowFp16PrecisionForFp32(true));
float[] oneD = {1.23f, 6.54f, 7.81f};
float[][] twoD = {oneD, oneD, oneD, oneD, oneD, oneD, oneD, oneD};
float[][][] threeD = {twoD, twoD, twoD, twoD, twoD, twoD, twoD, twoD};
diff --git a/tensorflow/contrib/lite/kernels/activations.cc b/tensorflow/contrib/lite/kernels/activations.cc
index b2d9b84..cf9441a 100644
--- a/tensorflow/contrib/lite/kernels/activations.cc
+++ b/tensorflow/contrib/lite/kernels/activations.cc
@@ -348,18 +348,22 @@
return kTfLiteOk;
} break;
case kTfLiteInt16: {
- optimized_ops::Tanh(GetTensorData<int16_t>(input), GetTensorShape(input),
- data->input_left_shift,
- GetTensorData<int16_t>(output),
- GetTensorShape(output));
+ TanhParams params;
+ params.input_left_shift = data->input_left_shift;
+ optimized_ops::Tanh(params, GetTensorShape(input),
+ GetTensorData<int16_t>(input), GetTensorShape(output),
+ GetTensorData<int16_t>(output));
return kTfLiteOk;
} break;
case kTfLiteUInt8: {
- optimized_ops::Tanh(GetTensorData<uint8_t>(input), GetTensorShape(input),
- input->params.zero_point, data->input_range_radius,
- data->input_multiplier, data->input_left_shift,
- GetTensorData<uint8_t>(output),
- GetTensorShape(output));
+ TanhParams params;
+ params.input_zero_point = input->params.zero_point;
+ params.input_range_radius = data->input_range_radius;
+ params.input_multiplier = data->input_multiplier;
+ params.input_left_shift = data->input_left_shift;
+ optimized_ops::Tanh(params, GetTensorShape(input),
+ GetTensorData<uint8_t>(input), GetTensorShape(output),
+ GetTensorData<uint8_t>(output));
return kTfLiteOk;
} break;
default:
@@ -385,17 +389,21 @@
break;
}
case kTfLiteInt16: {
+ LogisticParams params;
optimized_ops::Logistic(
- GetTensorData<int16>(input), GetTensorShape(input),
- GetTensorData<int16_t>(output), GetTensorShape(output));
+ params, GetTensorShape(input), GetTensorData<int16_t>(input),
+ GetTensorShape(output), GetTensorData<int16_t>(output));
break;
}
case kTfLiteUInt8: {
+ LogisticParams params;
+ params.input_zero_point = input->params.zero_point;
+ params.input_range_radius = data->input_range_radius;
+ params.input_multiplier = data->input_multiplier;
+ params.input_left_shift = data->input_left_shift;
optimized_ops::Logistic(
- GetTensorData<uint8_t>(input), GetTensorShape(input),
- input->params.zero_point, data->input_range_radius,
- data->input_multiplier, data->input_left_shift,
- GetTensorData<uint8_t>(output), GetTensorShape(output));
+ params, GetTensorShape(input), GetTensorData<uint8_t>(input),
+ GetTensorShape(output), GetTensorData<uint8_t>(output));
break;
}
default:
@@ -459,11 +467,13 @@
const int batch_size = input->dims->data[0];
const int intermediate_size = input->dims->data[1];
const int input_size = input->dims->data[2];
+ SoftmaxParams op_params;
+ op_params.beta = params->beta;
optimized_ops::Softmax(
+ op_params, GetTensorShape({batch_size, intermediate_size, 1, input_size}),
GetTensorData<float>(input),
GetTensorShape({batch_size, intermediate_size, 1, input_size}),
- params->beta, GetTensorData<float>(output),
- GetTensorShape({batch_size, intermediate_size, 1, input_size}));
+ GetTensorData<float>(output));
}
void Softmax1DQuantized(const TfLiteTensor* input, TfLiteTensor* output,
@@ -473,10 +483,14 @@
// tensor is 4D in a special way. We will convert a (Y) shape into a (1,
// 1, 1, Y) shape.
const int input_size = input->dims->data[0];
- optimized_ops::Softmax(
- GetTensorData<uint8_t>(input), GetTensorShape({1, 1, 1, input_size}),
- data->input_multiplier, data->input_left_shift, data->diff_min,
- GetTensorData<uint8_t>(output), GetTensorShape({1, 1, 1, input_size}));
+ SoftmaxParams op_params;
+ op_params.input_multiplier = data->input_multiplier;
+ op_params.input_left_shift = data->input_left_shift;
+ op_params.diff_min = data->diff_min;
+ optimized_ops::Softmax(op_params, GetTensorShape({1, 1, 1, input_size}),
+ GetTensorData<uint8_t>(input),
+ GetTensorShape({1, 1, 1, input_size}),
+ GetTensorData<uint8_t>(output));
}
void Softmax2DQuantized(const TfLiteTensor* input, TfLiteTensor* output,
TfLiteSoftmaxParams* params, OpData* data) {
@@ -486,11 +500,15 @@
// 1, 1, Y) shape.
const int batch_size = input->dims->data[0];
const int input_size = input->dims->data[1];
- optimized_ops::Softmax(GetTensorData<uint8_t>(input),
+ SoftmaxParams op_params;
+ op_params.input_multiplier = data->input_multiplier;
+ op_params.input_left_shift = data->input_left_shift;
+ op_params.diff_min = data->diff_min;
+ optimized_ops::Softmax(op_params,
GetTensorShape({batch_size, 1, 1, input_size}),
- data->input_multiplier, data->input_left_shift,
- data->diff_min, GetTensorData<uint8_t>(output),
- GetTensorShape({batch_size, 1, 1, input_size}));
+ GetTensorData<uint8_t>(input),
+ GetTensorShape({batch_size, 1, 1, input_size}),
+ GetTensorData<uint8_t>(output));
}
void Softmax3DQuantized(const TfLiteTensor* input, TfLiteTensor* output,
@@ -498,28 +516,36 @@
const int batch_size = input->dims->data[0];
const int intermediate_size = input->dims->data[1];
const int input_size = input->dims->data[2];
+ SoftmaxParams op_params;
+ op_params.input_multiplier = data->input_multiplier;
+ op_params.input_left_shift = data->input_left_shift;
+ op_params.diff_min = data->diff_min;
optimized_ops::Softmax(
+ op_params, GetTensorShape({batch_size, intermediate_size, 1, input_size}),
GetTensorData<uint8_t>(input),
GetTensorShape({batch_size, intermediate_size, 1, input_size}),
- data->input_multiplier, data->input_left_shift, data->diff_min,
- GetTensorData<uint8_t>(output),
- GetTensorShape({batch_size, intermediate_size, 1, input_size}));
+ GetTensorData<uint8_t>(output));
}
// Takes a 4D tensor and perform softmax along the forth dimension.
void Softmax4DFloat(const TfLiteTensor* input, TfLiteTensor* output,
TfLiteSoftmaxParams* params) {
- optimized_ops::Softmax(GetTensorData<float>(input), GetTensorShape(input),
- params->beta, GetTensorData<float>(output),
- GetTensorShape(output));
+ SoftmaxParams op_params;
+ op_params.beta = params->beta;
+ optimized_ops::Softmax(op_params, GetTensorShape(input),
+ GetTensorData<float>(input), GetTensorShape(output),
+ GetTensorData<float>(output));
}
void Softmax4DQuantized(const TfLiteTensor* input, TfLiteTensor* output,
TfLiteSoftmaxParams* params, OpData* data) {
- optimized_ops::Softmax(GetTensorData<uint8_t>(input), GetTensorShape(input),
- data->input_multiplier, data->input_left_shift,
- data->diff_min, GetTensorData<uint8_t>(output),
- GetTensorShape(output));
+ SoftmaxParams op_params;
+ op_params.input_multiplier = data->input_multiplier;
+ op_params.input_left_shift = data->input_left_shift;
+ op_params.diff_min = data->diff_min;
+ optimized_ops::Softmax(op_params, GetTensorShape(input),
+ GetTensorData<uint8_t>(input), GetTensorShape(output),
+ GetTensorData<uint8_t>(output));
}
TfLiteStatus SoftmaxEval(TfLiteContext* context, TfLiteNode* node) {
@@ -591,17 +617,20 @@
TfLiteTensor* output = GetOutput(context, node, 0);
switch (input->type) {
case kTfLiteFloat32:
+ SoftmaxParams op_params;
optimized_ops::LogSoftmax(
- GetTensorData<float>(input), GetTensorShape(input),
- GetTensorData<float>(output), GetTensorShape(output));
+ op_params, GetTensorShape(input), GetTensorData<float>(input),
+ GetTensorShape(output), GetTensorData<float>(output));
return kTfLiteOk;
case kTfLiteUInt8:
+ op_params.input_multiplier = data->input_multiplier;
+ op_params.input_left_shift = data->input_left_shift;
+ op_params.reverse_scaling_divisor = data->reverse_scaling_divisor;
+ op_params.reverse_scaling_right_shift = data->reverse_scaling_right_shift;
+ op_params.diff_min = data->diff_min;
optimized_ops::LogSoftmax(
- GetTensorData<uint8_t>(input), GetTensorShape(input),
- data->input_multiplier, data->input_left_shift,
- data->reverse_scaling_divisor, data->reverse_scaling_right_shift,
- data->diff_min, GetTensorData<uint8_t>(output),
- GetTensorShape(output));
+ op_params, GetTensorShape(input), GetTensorData<uint8_t>(input),
+ GetTensorShape(output), GetTensorData<uint8_t>(output));
return kTfLiteOk;
default:
context->ReportError(context, "Only float32 supported currently., got %d",
diff --git a/tensorflow/contrib/lite/kernels/bidirectional_sequence_lstm.cc b/tensorflow/contrib/lite/kernels/bidirectional_sequence_lstm.cc
index 541f320..66b9477 100644
--- a/tensorflow/contrib/lite/kernels/bidirectional_sequence_lstm.cc
+++ b/tensorflow/contrib/lite/kernels/bidirectional_sequence_lstm.cc
@@ -770,51 +770,29 @@
}
// Loop through the sequence.
- if (forward_sequence) {
- for (int t = 0; t < max_time; t++) {
- const float* input_ptr = input->data.f + t * n_batch * n_input;
- float* output_ptr_time = output->data.f + t * n_batch * n_output;
+ const int input_step = n_batch * n_input;
+ const int output_step = n_batch * n_output;
+ for (int t = 0; t < max_time; t++) {
+ // If this is the forward_sequence, step forward, otherwise step backwards.
+ const int t_rel = forward_sequence ? t : max_time - t - 1;
+ const float* input_ptr = input->data.f + t_rel * input_step;
+ float* output_ptr_time = output->data.f + t_rel * output_step;
- kernel_utils::LstmStepWithAuxInput(
- input_ptr, input_to_input_weights_ptr,
- input_to_forget_weights->data.f, input_to_cell_weights->data.f,
- input_to_output_weights->data.f, aux_input_ptr,
- aux_input_to_input_weights_ptr, aux_input_to_forget_weights_ptr,
- aux_input_to_cell_weights_ptr, aux_input_to_output_weights_ptr,
- recurrent_to_input_weights_ptr, recurrent_to_forget_weights->data.f,
- recurrent_to_cell_weights->data.f,
- recurrent_to_output_weights->data.f, cell_to_input_weights_ptr,
- cell_to_forget_weights_ptr, cell_to_output_weights_ptr,
- input_gate_bias_ptr, forget_gate_bias->data.f, cell_bias->data.f,
- output_gate_bias->data.f, projection_weights_ptr, projection_bias_ptr,
- params, n_batch, n_cell, n_input, aux_input_size, n_output,
- activation_state->data.f, cell_state->data.f, input_gate_scratch,
- forget_gate_scratch, cell_scratch, output_gate_scratch,
- output_ptr_time);
- }
- } else {
- // Loop through the sequence backwards.
- for (int t = max_time - 1; t >= 0; t--) {
- const float* input_ptr = input->data.f + t * n_batch * n_input;
- float* output_ptr_time = output->data.f + t * n_batch * n_output;
-
- kernel_utils::LstmStepWithAuxInput(
- input_ptr, input_to_input_weights_ptr,
- input_to_forget_weights->data.f, input_to_cell_weights->data.f,
- input_to_output_weights->data.f, aux_input_ptr,
- aux_input_to_input_weights_ptr, aux_input_to_forget_weights_ptr,
- aux_input_to_cell_weights_ptr, aux_input_to_output_weights_ptr,
- recurrent_to_input_weights_ptr, recurrent_to_forget_weights->data.f,
- recurrent_to_cell_weights->data.f,
- recurrent_to_output_weights->data.f, cell_to_input_weights_ptr,
- cell_to_forget_weights_ptr, cell_to_output_weights_ptr,
- input_gate_bias_ptr, forget_gate_bias->data.f, cell_bias->data.f,
- output_gate_bias->data.f, projection_weights_ptr, projection_bias_ptr,
- params, n_batch, n_cell, n_input, aux_input_size, n_output,
- activation_state->data.f, cell_state->data.f, input_gate_scratch,
- forget_gate_scratch, cell_scratch, output_gate_scratch,
- output_ptr_time);
- }
+ kernel_utils::LstmStepWithAuxInput(
+ input_ptr, input_to_input_weights_ptr, input_to_forget_weights->data.f,
+ input_to_cell_weights->data.f, input_to_output_weights->data.f,
+ aux_input_ptr, aux_input_to_input_weights_ptr,
+ aux_input_to_forget_weights_ptr, aux_input_to_cell_weights_ptr,
+ aux_input_to_output_weights_ptr, recurrent_to_input_weights_ptr,
+ recurrent_to_forget_weights->data.f, recurrent_to_cell_weights->data.f,
+ recurrent_to_output_weights->data.f, cell_to_input_weights_ptr,
+ cell_to_forget_weights_ptr, cell_to_output_weights_ptr,
+ input_gate_bias_ptr, forget_gate_bias->data.f, cell_bias->data.f,
+ output_gate_bias->data.f, projection_weights_ptr, projection_bias_ptr,
+ params, n_batch, n_cell, n_input, aux_input_size, n_output,
+ activation_state->data.f, cell_state->data.f, input_gate_scratch,
+ forget_gate_scratch, cell_scratch, output_gate_scratch,
+ output_ptr_time);
}
return kTfLiteOk;
}
@@ -991,72 +969,41 @@
aux_input_to_output_weights_scale =
aux_input_to_output_weights->params.scale;
}
- if (forward_sequence) {
- // Feed the sequence into the LSTM step-by-step.
- for (int t = 0; t < max_time; t++) {
- const float* input_ptr = input->data.f + t * n_batch * n_input;
- float* output_ptr = output->data.f + t * n_batch * n_output;
- kernel_utils::LstmStepWithAuxInput(
- input_ptr, input_to_input_weights_ptr, input_to_input_weights_scale,
- input_to_forget_weights_ptr, input_to_forget_weights_scale,
- input_to_cell_weights_ptr, input_to_cell_weights_scale,
- input_to_output_weights_ptr, input_to_output_weights_scale,
- aux_input_ptr, aux_input_to_input_weights_ptr,
- aux_input_to_input_weights_scale, aux_input_to_forget_weights_ptr,
- aux_input_to_forget_weights_scale, aux_input_to_cell_weights_ptr,
- aux_input_to_cell_weights_scale, aux_input_to_output_weights_ptr,
- aux_input_to_output_weights_scale, recurrent_to_input_weights_ptr,
- recurrent_to_input_weights_scale, recurrent_to_forget_weights_ptr,
- recurrent_to_forget_weights_scale, recurrent_to_cell_weights_ptr,
- recurrent_to_cell_weights_scale, recurrent_to_output_weights_ptr,
- recurrent_to_output_weights_scale, cell_to_input_weights_ptr,
- cell_to_input_weights_scale, cell_to_forget_weights_ptr,
- cell_to_forget_weights_scale, cell_to_output_weights_ptr,
- 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, aux_input_size, n_output,
- input_gate_scratch, forget_gate_scratch, cell_scratch,
- output_gate_scratch, scaling_factors_ptr, prod_scaling_factors_ptr,
- recovered_cell_weights_ptr, quantized_input_ptr,
- quantized_aux_input_ptr, quantized_output_state_ptr,
- quantized_cell_state_ptr, output_state_ptr, cell_state_ptr,
- output_ptr);
- }
- } else {
- // Loop through the sequence backwards.
- for (int t = max_time - 1; t >= 0; t--) {
- const float* input_ptr = input->data.f + t * n_batch * n_input;
- float* output_ptr = output->data.f + t * n_batch * n_output;
+ // Feed the sequence into the LSTM step-by-step.
+ const int input_step = n_batch * n_input;
+ const int output_step = n_batch * n_output;
+ for (int t = 0; t < max_time; t++) {
+ // If this is the forward_sequence, step forward, otherwise step backwards.
+ const int t_rel = forward_sequence ? t : max_time - t - 1;
+ const float* input_ptr = input->data.f + t_rel * input_step;
+ float* output_ptr = output->data.f + t_rel * output_step;
- kernel_utils::LstmStepWithAuxInput(
- input_ptr, input_to_input_weights_ptr, input_to_input_weights_scale,
- input_to_forget_weights_ptr, input_to_forget_weights_scale,
- input_to_cell_weights_ptr, input_to_cell_weights_scale,
- input_to_output_weights_ptr, input_to_output_weights_scale,
- aux_input_ptr, aux_input_to_input_weights_ptr,
- aux_input_to_input_weights_scale, aux_input_to_forget_weights_ptr,
- aux_input_to_forget_weights_scale, aux_input_to_cell_weights_ptr,
- aux_input_to_cell_weights_scale, aux_input_to_output_weights_ptr,
- aux_input_to_output_weights_scale, recurrent_to_input_weights_ptr,
- recurrent_to_input_weights_scale, recurrent_to_forget_weights_ptr,
- recurrent_to_forget_weights_scale, recurrent_to_cell_weights_ptr,
- recurrent_to_cell_weights_scale, recurrent_to_output_weights_ptr,
- recurrent_to_output_weights_scale, cell_to_input_weights_ptr,
- cell_to_input_weights_scale, cell_to_forget_weights_ptr,
- cell_to_forget_weights_scale, cell_to_output_weights_ptr,
- 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, aux_input_size, n_output,
- input_gate_scratch, forget_gate_scratch, cell_scratch,
- output_gate_scratch, scaling_factors_ptr, prod_scaling_factors_ptr,
- recovered_cell_weights_ptr, quantized_input_ptr,
- quantized_aux_input_ptr, quantized_output_state_ptr,
- quantized_cell_state_ptr, output_state_ptr, cell_state_ptr,
- output_ptr);
- }
+ kernel_utils::LstmStepWithAuxInput(
+ input_ptr, input_to_input_weights_ptr, input_to_input_weights_scale,
+ input_to_forget_weights_ptr, input_to_forget_weights_scale,
+ input_to_cell_weights_ptr, input_to_cell_weights_scale,
+ input_to_output_weights_ptr, input_to_output_weights_scale,
+ aux_input_ptr, aux_input_to_input_weights_ptr,
+ aux_input_to_input_weights_scale, aux_input_to_forget_weights_ptr,
+ aux_input_to_forget_weights_scale, aux_input_to_cell_weights_ptr,
+ aux_input_to_cell_weights_scale, aux_input_to_output_weights_ptr,
+ aux_input_to_output_weights_scale, recurrent_to_input_weights_ptr,
+ recurrent_to_input_weights_scale, recurrent_to_forget_weights_ptr,
+ recurrent_to_forget_weights_scale, recurrent_to_cell_weights_ptr,
+ recurrent_to_cell_weights_scale, recurrent_to_output_weights_ptr,
+ recurrent_to_output_weights_scale, cell_to_input_weights_ptr,
+ cell_to_input_weights_scale, cell_to_forget_weights_ptr,
+ cell_to_forget_weights_scale, cell_to_output_weights_ptr,
+ 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, aux_input_size, n_output, input_gate_scratch,
+ forget_gate_scratch, cell_scratch, output_gate_scratch,
+ scaling_factors_ptr, prod_scaling_factors_ptr,
+ recovered_cell_weights_ptr, quantized_input_ptr,
+ quantized_aux_input_ptr, quantized_output_state_ptr,
+ quantized_cell_state_ptr, output_state_ptr, cell_state_ptr, output_ptr);
}
return kTfLiteOk;
diff --git a/tensorflow/contrib/lite/kernels/comparisons.cc b/tensorflow/contrib/lite/kernels/comparisons.cc
index 4cd9634..f765235 100644
--- a/tensorflow/contrib/lite/kernels/comparisons.cc
+++ b/tensorflow/contrib/lite/kernels/comparisons.cc
@@ -83,20 +83,24 @@
QuantizeMultiplierSmallerThanOneExp(real_input2_multiplier, \
&input2_multiplier, &input2_shift); \
\
+ ComparisonParams op_params; \
+ op_params.left_shift = left_shift; \
+ op_params.input1_offset = input1_offset; \
+ op_params.input1_multiplier = input1_multiplier; \
+ op_params.input1_shift = -input1_shift; \
+ op_params.input2_offset = input2_offset; \
+ op_params.input2_multiplier = input2_multiplier; \
+ op_params.input2_shift = -input2_shift; \
if (requires_broadcast) { \
- reference_ops::Broadcast##opname( \
- left_shift, GetTensorData<uint8_t>(input1), GetTensorDims(input1), \
- input1_offset, input1_multiplier, input1_shift, \
- GetTensorData<uint8_t>(input2), GetTensorDims(input2), \
- input2_offset, input2_multiplier, input2_shift, \
- GetTensorData<bool>(output), GetTensorDims(output)); \
+ reference_ops::Broadcast4DSlow##opname##WithScaling( \
+ op_params, GetTensorShape(input1), GetTensorData<uint8_t>(input1), \
+ GetTensorShape(input2), GetTensorData<uint8_t>(input2), \
+ GetTensorShape(output), GetTensorData<bool>(output)); \
} else { \
- reference_ops::opname( \
- left_shift, GetTensorData<uint8_t>(input1), GetTensorDims(input1), \
- input1_offset, input1_multiplier, input1_shift, \
- GetTensorData<uint8_t>(input2), GetTensorDims(input2), \
- input2_offset, input2_multiplier, input2_shift, \
- GetTensorData<bool>(output), GetTensorDims(output)); \
+ reference_ops::opname##WithScaling( \
+ op_params, GetTensorShape(input1), GetTensorData<uint8_t>(input1), \
+ GetTensorShape(input2), GetTensorData<uint8_t>(input2), \
+ GetTensorShape(output), GetTensorData<bool>(output)); \
} \
} \
}
@@ -108,16 +112,19 @@
TF_LITE_QUANTIZE_COMPARISON(LessEqual);
#undef TF_LITE_QUANTIZE_COMPARISON
-#define TF_LITE_COMPARISON(type, opname, requires_broadcast) \
- requires_broadcast \
- ? reference_ops::Broadcast##opname( \
- GetTensorData<type>(input1), GetTensorDims(input1), \
- GetTensorData<type>(input2), GetTensorDims(input2), \
- GetTensorData<bool>(output), GetTensorDims(output)) \
- : reference_ops::opname( \
- GetTensorData<type>(input1), GetTensorDims(input1), \
- GetTensorData<type>(input2), GetTensorDims(input2), \
- GetTensorData<bool>(output), GetTensorDims(output));
+#define TF_LITE_COMPARISON(type, opname, requires_broadcast) \
+ { \
+ ComparisonParams op_params; \
+ requires_broadcast \
+ ? reference_ops::Broadcast4DSlow##opname##NoScaling( \
+ op_params, GetTensorShape(input1), GetTensorData<type>(input1), \
+ GetTensorShape(input2), GetTensorData<type>(input2), \
+ GetTensorShape(output), GetTensorData<bool>(output)) \
+ : reference_ops::opname##NoScaling( \
+ op_params, GetTensorShape(input1), GetTensorData<type>(input1), \
+ GetTensorShape(input2), GetTensorData<type>(input2), \
+ GetTensorShape(output), GetTensorData<bool>(output)); \
+ }
TfLiteStatus EqualEval(TfLiteContext* context, TfLiteNode* node) {
const TfLiteTensor* input1 = GetInput(context, node, kInputTensor1);
diff --git a/tensorflow/contrib/lite/kernels/concatenation.cc b/tensorflow/contrib/lite/kernels/concatenation.cc
index 25ea556..7ad3399 100644
--- a/tensorflow/contrib/lite/kernels/concatenation.cc
+++ b/tensorflow/contrib/lite/kernels/concatenation.cc
@@ -100,20 +100,31 @@
// allocate and populate these during Prepare().
// TODO(ycling): Activation function parameter is ignored. For now we dont have
// a model with a Concatenation with fused activation function.
-#define TF_LITE_CONCATENATION(type, scalar) \
- VectorOfTensors<scalar> all_inputs(*context, *node->inputs); \
- type::Concatenation<FusedActivationFunctionType::kNone, scalar>( \
- RemapDim(NumDimensions(output), axis), all_inputs.data(), \
- all_inputs.dims(), node->inputs->size, GetTensorData<scalar>(output), \
- GetTensorDims(output))
+#define TF_LITE_CONCATENATION(type, scalar) \
+ { \
+ VectorOfTensors<scalar> all_inputs(*context, *node->inputs); \
+ tflite::ConcatenationParams op_params; \
+ op_params.axis = axis; \
+ op_params.inputs_count = node->inputs->size; \
+ type::Concatenation(op_params, all_inputs.shapes(), all_inputs.data(), \
+ GetTensorShape(output), \
+ GetTensorData<scalar>(output)); \
+ }
-#define TF_LITE_CONCATENATION_QUANTIZED(type) \
- VectorOfQuantizedTensors all_inputs(*context, *node->inputs); \
- type::Concatenation( \
- RemapDim(NumDimensions(output), axis), all_inputs.data(), \
- all_inputs.dims(), all_inputs.zero_point(), all_inputs.scale(), \
- node->inputs->size, GetTensorData<uint8>(output), GetTensorDims(output), \
- output->params.zero_point, output->params.scale)
+#define TF_LITE_CONCATENATION_QUANTIZED(type) \
+ { \
+ VectorOfQuantizedTensors all_inputs(*context, *node->inputs); \
+ tflite::ConcatenationParams op_params; \
+ op_params.axis = axis; \
+ op_params.input_zeropoint = all_inputs.zero_point(); \
+ op_params.input_scale = all_inputs.scale(); \
+ op_params.inputs_count = node->inputs->size; \
+ op_params.output_zeropoint = output->params.zero_point; \
+ op_params.output_scale = output->params.scale; \
+ type::ConcatenationWithScaling(op_params, all_inputs.shapes(), \
+ all_inputs.data(), GetTensorShape(output), \
+ GetTensorData<uint8>(output)); \
+ }
switch (output->type) { // Already know in/outtypes are same.
case kTfLiteFloat32:
diff --git a/tensorflow/contrib/lite/kernels/conv.cc b/tensorflow/contrib/lite/kernels/conv.cc
index 101b4fc..dbcadbe 100644
--- a/tensorflow/contrib/lite/kernels/conv.cc
+++ b/tensorflow/contrib/lite/kernels/conv.cc
@@ -86,6 +86,18 @@
bool run_multithreaded_kernel;
};
+inline PaddingType RuntimePaddingType(TfLitePadding padding) {
+ switch (padding) {
+ case TfLitePadding::kTfLitePaddingSame:
+ return PaddingType::kSame;
+ case TfLitePadding::kTfLitePaddingValid:
+ return PaddingType::kValid;
+ case TfLitePadding::kTfLitePaddingUnknown:
+ default:
+ return PaddingType::kNone;
+ }
+}
+
void* Init(TfLiteContext* context, const char* buffer, size_t length) {
// This is a builtin op, so we don't use the contents in 'buffer', if any.
// Instead, we allocate a new object to use as scratch space for im2col, and
@@ -487,18 +499,18 @@
} else {
effective_kernel_type = kernel_type;
}
+ ConvParams op_params;
+ op_params.padding_type = RuntimePaddingType(params->padding);
+ op_params.padding_values.width = data->padding.width;
+ op_params.padding_values.height = data->padding.height;
+ op_params.stride_width = params->stride_width;
+ op_params.stride_height = params->stride_height;
+ op_params.dilation_width_factor = params->dilation_width_factor;
+ op_params.dilation_height_factor = params->dilation_height_factor;
+ op_params.float_activation_min = output_activation_min;
+ op_params.float_activation_max = output_activation_max;
switch (effective_kernel_type) {
case kReference: {
- ConvParams op_params;
- op_params.padding_type = PaddingType::kSame;
- op_params.padding_values.width = data->padding.width;
- op_params.padding_values.height = data->padding.height;
- op_params.stride_width = params->stride_width;
- op_params.stride_height = params->stride_height;
- op_params.dilation_width_factor = params->dilation_width_factor;
- op_params.dilation_height_factor = params->dilation_height_factor;
- op_params.float_activation_min = output_activation_min;
- op_params.float_activation_max = output_activation_max;
reference_ops::Conv(op_params, GetTensorShape(input),
GetTensorData<float>(input), GetTensorShape(filter),
GetTensorData<float>(filter), GetTensorShape(bias),
@@ -508,16 +520,6 @@
break;
}
case kGenericOptimized: {
- ConvParams op_params;
- op_params.padding_type = PaddingType::kSame;
- op_params.padding_values.width = data->padding.width;
- op_params.padding_values.height = data->padding.height;
- op_params.stride_width = params->stride_width;
- op_params.stride_height = params->stride_height;
- op_params.dilation_width_factor = params->dilation_width_factor;
- op_params.dilation_height_factor = params->dilation_height_factor;
- op_params.float_activation_min = output_activation_min;
- op_params.float_activation_max = output_activation_max;
optimized_ops::Conv(op_params, GetTensorShape(input),
GetTensorData<float>(input), GetTensorShape(filter),
GetTensorData<float>(filter), GetTensorShape(bias),
@@ -534,25 +536,21 @@
filter_data = GetTensorData<float>(filter);
}
multithreaded_ops::Conv(
- *eigen_support::GetThreadPoolDevice(context),
- GetTensorData<float>(input), GetTensorDims(input), filter_data,
- GetTensorDims(filter), GetTensorData<float>(bias),
- GetTensorDims(bias), params->stride_width, params->stride_height,
- data->padding.width, data->padding.height, params->padding,
- output_activation_min, output_activation_max,
- GetTensorData<float>(output), GetTensorDims(output),
- GetTensorData<float>(im2col), GetTensorDims(im2col));
+ *eigen_support::GetThreadPoolDevice(context), op_params,
+ GetTensorShape(input), GetTensorData<float>(input),
+ GetTensorShape(filter), filter_data, GetTensorShape(bias),
+ GetTensorData<float>(bias), GetTensorShape(output),
+ GetTensorData<float>(output), GetTensorShape(im2col),
+ GetTensorData<float>(im2col));
break;
}
case kCblasOptimized: {
- cblas_ops::Conv(GetTensorData<float>(input), GetTensorDims(input),
- GetTensorData<float>(filter), GetTensorDims(filter),
- GetTensorData<float>(bias), GetTensorDims(bias),
- params->stride_width, params->stride_height,
- data->padding.width, data->padding.height,
- output_activation_min, output_activation_max,
- GetTensorData<float>(output), GetTensorDims(output),
- GetTensorData<float>(im2col), GetTensorDims(im2col));
+ cblas_ops::Conv(op_params, GetTensorShape(input),
+ GetTensorData<float>(input), GetTensorShape(filter),
+ GetTensorData<float>(filter), GetTensorShape(bias),
+ GetTensorData<float>(bias), GetTensorShape(output),
+ GetTensorData<float>(output), GetTensorShape(im2col),
+ GetTensorData<float>(im2col));
break;
}
}
diff --git a/tensorflow/contrib/lite/kernels/dequantize.cc b/tensorflow/contrib/lite/kernels/dequantize.cc
index 3a08f48..59bf64e 100644
--- a/tensorflow/contrib/lite/kernels/dequantize.cc
+++ b/tensorflow/contrib/lite/kernels/dequantize.cc
@@ -77,13 +77,13 @@
return kTfLiteOk;
}
- auto zero_point = op_context.input->params.zero_point;
- auto scale = op_context.input->params.scale;
-
- optimized_ops::Dequantize(GetTensorData<uint8_t>(op_context.input),
- GetTensorDims(op_context.input), zero_point, scale,
- GetTensorData<float>(op_context.output),
- GetTensorDims(op_context.output));
+ 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));
if (IsConstantTensor(op_context.input)) {
op_data->float_dequantized_weights_initialized = true;
diff --git a/tensorflow/contrib/lite/kernels/div.cc b/tensorflow/contrib/lite/kernels/div.cc
index 7945c09..8d4bb51 100644
--- a/tensorflow/contrib/lite/kernels/div.cc
+++ b/tensorflow/contrib/lite/kernels/div.cc
@@ -81,24 +81,27 @@
void EvalDiv(TfLiteContext* context, TfLiteNode* node, TfLiteDivParams* params,
const OpData* data, const TfLiteTensor* input1,
const TfLiteTensor* input2, TfLiteTensor* output) {
-#define TF_LITE_DIV(type, opname, data_type) \
- data_type output_activation_min, output_activation_max; \
- CalculateActivationRange(params->activation, &output_activation_min, \
- &output_activation_max); \
- type::opname(GetTensorData<data_type>(input1), GetTensorDims(input1), \
- GetTensorData<data_type>(input2), GetTensorDims(input2), \
- output_activation_min, output_activation_max, \
- GetTensorData<data_type>(output), GetTensorDims(output))
+#define TF_LITE_DIV(type, opname, data_type) \
+ tflite::ArithmeticParams op_params; \
+ data_type output_activation_min, output_activation_max; \
+ CalculateActivationRange(params->activation, &output_activation_min, \
+ &output_activation_max); \
+ SetActivationParams(output_activation_min, output_activation_max, \
+ &op_params); \
+ type::opname(op_params, GetTensorShape(input1), \
+ GetTensorData<data_type>(input1), GetTensorShape(input2), \
+ GetTensorData<data_type>(input2), GetTensorShape(output), \
+ GetTensorData<data_type>(output))
if (output->type == kTfLiteInt32) {
if (kernel_type == kReference) {
if (data->requires_broadcast) {
- TF_LITE_DIV(reference_ops, BroadcastDiv, int32_t);
+ TF_LITE_DIV(reference_ops, BroadcastDiv4DSlow, int32_t);
} else {
TF_LITE_DIV(reference_ops, Div, int32_t);
}
} else {
if (data->requires_broadcast) {
- TF_LITE_DIV(optimized_ops, BroadcastDiv, int32_t);
+ TF_LITE_DIV(optimized_ops, BroadcastDiv4DSlow, int32_t);
} else {
TF_LITE_DIV(optimized_ops, Div, int32_t);
}
@@ -106,13 +109,13 @@
} else if (output->type == kTfLiteFloat32) {
if (kernel_type == kReference) {
if (data->requires_broadcast) {
- TF_LITE_DIV(reference_ops, BroadcastDiv, float);
+ TF_LITE_DIV(reference_ops, BroadcastDiv4DSlow, float);
} else {
TF_LITE_DIV(reference_ops, Div, float);
}
} else {
if (data->requires_broadcast) {
- TF_LITE_DIV(optimized_ops, BroadcastDiv, float);
+ TF_LITE_DIV(optimized_ops, BroadcastDiv4DSlow, float);
} else {
TF_LITE_DIV(optimized_ops, Div, float);
}
diff --git a/tensorflow/contrib/lite/kernels/fake_quant.cc b/tensorflow/contrib/lite/kernels/fake_quant.cc
index f9bc374..b51af72 100644
--- a/tensorflow/contrib/lite/kernels/fake_quant.cc
+++ b/tensorflow/contrib/lite/kernels/fake_quant.cc
@@ -68,11 +68,14 @@
const auto* params =
reinterpret_cast<TfLiteFakeQuantParams*>(node->builtin_data);
- reference_ops::FakeQuant(GetTensorData<float>(op_context.input),
- GetTensorDims(op_context.input), params->min,
- params->max, params->num_bits,
- GetTensorData<float>(op_context.output),
- GetTensorDims(op_context.output));
+ tflite::FakeQuantParams op_params;
+ op_params.num_bits = params->num_bits;
+ op_params.minmax.min = params->min;
+ op_params.minmax.max = params->max;
+ reference_ops::FakeQuant(op_params, GetTensorShape(op_context.input),
+ GetTensorData<float>(op_context.input),
+ GetTensorShape(op_context.output),
+ GetTensorData<float>(op_context.output));
return kTfLiteOk;
}
diff --git a/tensorflow/contrib/lite/kernels/gather.cc b/tensorflow/contrib/lite/kernels/gather.cc
index badd2de..b5afeb1 100644
--- a/tensorflow/contrib/lite/kernels/gather.cc
+++ b/tensorflow/contrib/lite/kernels/gather.cc
@@ -84,11 +84,15 @@
const TfLiteTensor* positions = GetInput(context, node, kInputPositions);
TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
const int input_rank = NumDimensions(input);
-#define TF_LITE_GATHER(data_type, index_type) \
- optimized_ops::Gather( \
- GetTensorData<data_type>(input), GetTensorDims(input), input_rank, \
- GetTensorData<index_type>(positions), GetTensorDims(positions), \
- GetTensorData<data_type>(output), GetTensorDims(output));
+#define TF_LITE_GATHER(data_type, index_type) \
+ { \
+ tflite::GatherParams op_params; \
+ op_params.input_rank = input_rank; \
+ optimized_ops::Gather( \
+ op_params, GetTensorShape(input), GetTensorData<data_type>(input), \
+ GetTensorShape(positions), GetTensorData<index_type>(positions), \
+ GetTensorShape(output), GetTensorData<data_type>(output)); \
+ }
switch (input->type) {
case kTfLiteFloat32:
TF_LITE_GATHER(float, int32_t);
diff --git a/tensorflow/contrib/lite/kernels/internal/logsoftmax_quantized_test.cc b/tensorflow/contrib/lite/kernels/internal/logsoftmax_quantized_test.cc
index 3624c20..2252ca1 100644
--- a/tensorflow/contrib/lite/kernels/internal/logsoftmax_quantized_test.cc
+++ b/tensorflow/contrib/lite/kernels/internal/logsoftmax_quantized_test.cc
@@ -43,11 +43,15 @@
// Reference data generated via Dequant of input into float, and then applying
// float LogSoftmax.
- reference_ops::Dequantize(
- input_data, ToRuntimeDims(shape_common), input_offset, input_scale,
- reference_dequant_data.data(), ToRuntimeDims(shape_common));
- optimized_ops::LogSoftmax(reference_dequant_data.data(), shape_common,
- reference_output_float_data.data(), shape_common);
+ DequantizationParams dq_params;
+ dq_params.zero_point = input_offset;
+ dq_params.scale = input_scale;
+ reference_ops::Dequantize(dq_params, shape_common, input_data, shape_common,
+ reference_dequant_data.data());
+ SoftmaxParams sm_params;
+ optimized_ops::LogSoftmax(sm_params, shape_common,
+ reference_dequant_data.data(), shape_common,
+ reference_output_float_data.data());
// Work with quantized scaling for LogSoftmax, under which 255 represents 0,
// and -16 gets nudged up to 0.
for (int i = 0; i < ref_buffer_size; i++) {
@@ -129,14 +133,16 @@
const int diff_min = -tflite::CalculateInputRadius(kScaledDiffIntegerBits,
input_beta_left_shift);
- optimized_ops::LogSoftmax(input_data, shape_common, input_beta_multiplier,
- input_beta_left_shift, reverse_scaling_divisor,
- reverse_scaling_right_shift, diff_min,
- optimized_logsoftmax_output.data(), shape_common);
- reference_ops::LogSoftmax(
- input_data, shape_common, input_beta_multiplier, input_beta_left_shift,
- reverse_scaling_divisor, reverse_scaling_right_shift, diff_min,
- reference_quant_logsoftmax_output.data(), shape_common);
+ SoftmaxParams params;
+ params.input_multiplier = input_beta_multiplier;
+ params.input_left_shift = input_beta_left_shift;
+ params.reverse_scaling_divisor = reverse_scaling_divisor;
+ params.reverse_scaling_right_shift = reverse_scaling_right_shift;
+ params.diff_min = diff_min;
+ optimized_ops::LogSoftmax(params, shape_common, input_data, shape_common,
+ optimized_logsoftmax_output.data());
+ reference_ops::LogSoftmax(params, shape_common, input_data, shape_common,
+ reference_quant_logsoftmax_output.data());
CheckOutputData(optimized_logsoftmax_output.data(),
reference_float_logsoftmax_output.data(), shape_common,
diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/cblas_conv.h b/tensorflow/contrib/lite/kernels/internal/optimized/cblas_conv.h
index 40d42bb..2d96da6 100644
--- a/tensorflow/contrib/lite/kernels/internal/optimized/cblas_conv.h
+++ b/tensorflow/contrib/lite/kernels/internal/optimized/cblas_conv.h
@@ -31,20 +31,29 @@
namespace tflite {
namespace cblas_ops {
-inline void Conv(const float* input_data, const Dims<4>& input_dims,
- const float* filter_data, const Dims<4>& filter_dims,
- const float* bias_data, const Dims<4>& bias_dims,
- int stride_width, int stride_height, int pad_width,
- int pad_height, float output_activation_min,
- float output_activation_max, float* output_data,
- const Dims<4>& output_dims, float* im2col_data,
- const Dims<4>& im2col_dims) {
+inline void Conv(const ConvParams& params, const RuntimeShape& input_shape,
+ const float* input_data, const RuntimeShape& filter_shape,
+ const float* filter_data, const RuntimeShape& bias_shape,
+ const float* bias_data, const RuntimeShape& output_shape,
+ float* output_data, const RuntimeShape& im2col_shape,
+ float* im2col_data) {
+ const int stride_width = params.stride_width;
+ const int stride_height = params.stride_height;
+ const int pad_width = params.padding_values.width;
+ const int pad_height = params.padding_values.height;
+ const int dilation_width_factor = params.dilation_width_factor;
+ const int dilation_height_factor = params.dilation_height_factor;
+ const float output_activation_min = params.float_activation_min;
+ const float output_activation_max = params.float_activation_max;
+ TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
+ TFLITE_DCHECK_EQ(filter_shape.DimensionsCount(), 4);
+ TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
gemmlowp::ScopedProfilingLabel label("Conv/cblas");
const float* gemm_input_data = nullptr;
- const Dims<4>* gemm_input_dims = nullptr;
- const int filter_width = ArraySize(filter_dims, 1);
- const int filter_height = ArraySize(filter_dims, 2);
+ const RuntimeShape* gemm_input_shape = nullptr;
+ const int filter_width = filter_shape.Dims(2);
+ const int filter_height = filter_shape.Dims(1);
const bool need_im2col = stride_width != 1 || stride_height != 1 ||
filter_width != 1 || filter_height != 1;
if (need_im2col) {
@@ -55,18 +64,17 @@
op_params.padding_values.height = pad_height;
op_params.stride_width = stride_width;
op_params.stride_height = stride_height;
- op_params.dilation_width_factor = 1;
- op_params.dilation_height_factor = 1;
+ op_params.dilation_width_factor = dilation_width_factor;
+ op_params.dilation_height_factor = dilation_height_factor;
optimized_ops::Im2col(op_params, filter_height, filter_width, 0,
- DimsToShape(input_dims), input_data,
- DimsToShape(im2col_dims), im2col_data);
+ input_shape, input_data, im2col_shape, im2col_data);
gemm_input_data = im2col_data;
- gemm_input_dims = &im2col_dims;
+ gemm_input_shape = &im2col_shape;
} else {
TFLITE_DCHECK(!im2col_data);
gemm_input_data = input_data;
- gemm_input_dims = &input_dims;
+ gemm_input_shape = &input_shape;
}
// The following code computes matrix multiplication c = a * transponse(b)
@@ -78,10 +86,10 @@
const float* a = gemm_input_data;
const float* b = filter_data;
float* c = output_data;
- int m = gemm_input_dims->sizes[1] * gemm_input_dims->sizes[2] *
- gemm_input_dims->sizes[3];
- int n = output_dims.sizes[0];
- int k = gemm_input_dims->sizes[0];
+ const int gemm_input_dims = gemm_input_shape->DimensionsCount();
+ int m = FlatSizeSkipDim(*gemm_input_shape, gemm_input_dims - 1);
+ int n = output_shape.Dims(3);
+ int k = gemm_input_shape->Dims(gemm_input_dims - 1);
// The stride of matrix a, b and c respectively.
int stride_a = k;
int stride_b = k;
@@ -91,8 +99,8 @@
stride_a, b, stride_b, 0.0f, c, stride_c);
optimized_ops::AddBiasAndEvalActivationFunction(
- output_activation_min, output_activation_max, DimsToShape(bias_dims),
- bias_data, DimsToShape(output_dims), output_data);
+ output_activation_min, output_activation_max, bias_shape, bias_data,
+ output_shape, output_data);
}
} // namespace cblas_ops
diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_float.h b/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_float.h
index 114575a..d8dd7bb 100644
--- a/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_float.h
+++ b/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_float.h
@@ -1092,80 +1092,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void DepthwiseConv(const float* input_data, const Dims<4>& input_dims,
- const float* filter_data, const Dims<4>& filter_dims,
- const float* bias_data, const Dims<4>& bias_dims,
- int stride_width, int stride_height,
- int dilation_width_factor, int dilation_height_factor,
- int pad_width, int pad_height, int depth_multiplier,
- float output_activation_min,
- float output_activation_max, float* output_data,
- const Dims<4>& output_dims) {
- tflite::DepthwiseParams op_params;
- // Padding type is ignored, but still set.
- op_params.padding_type = PaddingType::kSame;
- op_params.padding_values.width = pad_width;
- op_params.padding_values.height = pad_height;
- op_params.stride_width = stride_width;
- op_params.stride_height = stride_height;
- op_params.dilation_width_factor = dilation_width_factor;
- op_params.dilation_height_factor = dilation_height_factor;
- op_params.depth_multiplier = depth_multiplier;
- op_params.float_activation_min = output_activation_min;
- op_params.float_activation_max = output_activation_max;
-
- DepthwiseConv(op_params, DimsToShape(input_dims), input_data,
- DimsToShape(filter_dims), filter_data, DimsToShape(bias_dims),
- bias_data, DimsToShape(output_dims), output_data);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-inline void DepthwiseConv(const float* input_data, const Dims<4>& input_dims,
- const float* filter_data, const Dims<4>& filter_dims,
- const float* bias_data, const Dims<4>& bias_dims,
- int stride_width, int stride_height, int pad_width,
- int pad_height, int depth_multiplier,
- float output_activation_min,
- float output_activation_max, float* output_data,
- const Dims<4>& output_dims) {
- DepthwiseConv(input_data, input_dims, filter_data, filter_dims, bias_data,
- bias_dims, stride_width, stride_height, 1, 1, pad_width,
- pad_height, depth_multiplier, output_activation_min,
- output_activation_max, output_data, output_dims);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// legacy, for compatibility with old checked-in code
-template <FusedActivationFunctionType Ac>
-void DepthwiseConv(const float* input_data, const Dims<4>& input_dims,
- const float* filter_data, const Dims<4>& filter_dims,
- const float* bias_data, const Dims<4>& bias_dims,
- int stride_width, int stride_height, int pad_width,
- int pad_height, int depth_multiplier, float* output_data,
- const Dims<4>& output_dims) {
- float output_activation_min, output_activation_max;
- GetActivationMinMax(Ac, &output_activation_min, &output_activation_max);
- DepthwiseConv(input_data, input_dims, filter_data, filter_dims, bias_data,
- bias_dims, stride_width, stride_height, pad_width, pad_height,
- depth_multiplier, output_activation_min, output_activation_max,
- output_data, output_dims);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// legacy, for compatibility with old checked-in code
-template <FusedActivationFunctionType Ac>
-void DepthwiseConv(const float* input_data, const Dims<4>& input_dims,
- const float* filter_data, const Dims<4>& filter_dims,
- const float* bias_data, const Dims<4>& bias_dims, int stride,
- int pad_width, int pad_height, int depth_multiplier,
- float* output_data, const Dims<4>& output_dims) {
- DepthwiseConv<Ac>(input_data, input_dims, filter_data, filter_dims, bias_data,
- bias_dims, stride, stride, pad_width, pad_height,
- depth_multiplier, output_data, output_dims);
-}
-
} // namespace optimized_ops
} // namespace tflite
diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_uint8.h b/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_uint8.h
index f892b8f..803eff2 100644
--- a/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_uint8.h
+++ b/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_uint8.h
@@ -24,9 +24,6 @@
namespace tflite {
namespace optimized_ops {
-// TODO(b/80418076): Move to legacy ops file, along with invocations.
-static constexpr int kDepthwiseReverseShift = -1;
-
// Implementation of quantized DepthwiseConv
template <bool kAllowStrided, int kFixedInputDepth, int kFixedDepthMultiplier>
@@ -1701,6 +1698,8 @@
const int output_shift = params.output_shift;
const int dilation_width_factor = params.dilation_width_factor;
const int dilation_height_factor = params.dilation_height_factor;
+ TFLITE_DCHECK_GE(dilation_width_factor, 1);
+ TFLITE_DCHECK_GE(dilation_height_factor, 1);
TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
TFLITE_DCHECK_EQ(filter_shape.DimensionsCount(), 4);
TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
@@ -1994,105 +1993,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void DepthwiseConv(const uint8* input_data, const Dims<4>& input_dims,
- int32 input_offset, const uint8* filter_data,
- const Dims<4>& filter_dims, int32 filter_offset,
- const int32* bias_data, const Dims<4>& bias_dims,
- int stride_width, int stride_height,
- int dilation_width_factor, int dilation_height_factor,
- int pad_width, int pad_height, int depth_multiplier,
- int32 output_offset, int32 output_multiplier,
- int output_shift, int32 output_activation_min,
- int32 output_activation_max, uint8* output_data,
- const Dims<4>& output_dims) {
- tflite::DepthwiseParams op_params;
- // Padding type is ignored, but still set.
- op_params.padding_type = PaddingType::kSame;
- op_params.padding_values.width = pad_width;
- op_params.padding_values.height = pad_height;
- op_params.stride_width = stride_width;
- op_params.stride_height = stride_height;
- op_params.dilation_width_factor = dilation_width_factor;
- op_params.dilation_height_factor = dilation_height_factor;
- op_params.depth_multiplier = depth_multiplier;
- op_params.quantized_activation_min = output_activation_min;
- op_params.quantized_activation_max = output_activation_max;
- op_params.input_offset = input_offset;
- op_params.weights_offset = filter_offset;
- op_params.output_offset = output_offset;
- op_params.output_multiplier = output_multiplier;
- // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
- op_params.output_shift = kDepthwiseReverseShift * output_shift;
-
- DepthwiseConv(op_params, DimsToShape(input_dims), input_data,
- DimsToShape(filter_dims), filter_data, DimsToShape(bias_dims),
- bias_data, DimsToShape(output_dims), output_data);
-}
-
-inline void DepthwiseConv(const uint8* input_data, const Dims<4>& input_dims,
- int32 input_offset, const uint8* filter_data,
- const Dims<4>& filter_dims, int32 filter_offset,
- const int32* bias_data, const Dims<4>& bias_dims,
- int stride_width, int stride_height, int pad_width,
- int pad_height, int depth_multiplier,
- int32 output_offset, int32 output_multiplier,
- int output_shift, int32 output_activation_min,
- int32 output_activation_max, uint8* output_data,
- const Dims<4>& output_dims) {
- DepthwiseConv(input_data, input_dims, input_offset, filter_data, filter_dims,
- filter_offset, bias_data, bias_dims, stride_width,
- stride_height, 1, 1, pad_width, pad_height, depth_multiplier,
- output_offset, output_multiplier, output_shift,
- output_activation_min, output_activation_max, output_data,
- output_dims);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy, for compatibility with old checked-in code.
-template <FusedActivationFunctionType Ac>
-void DepthwiseConv(const uint8* input_data, const Dims<4>& input_dims,
- int32 input_offset, const uint8* filter_data,
- const Dims<4>& filter_dims, int32 filter_offset,
- const int32* bias_data, const Dims<4>& bias_dims,
- int stride_width, int stride_height, int pad_width,
- int pad_height, int depth_multiplier, int32 output_offset,
- int32 output_multiplier, int output_shift,
- int32 output_activation_min, int32 output_activation_max,
- uint8* output_data, const Dims<4>& output_dims) {
- if (Ac == FusedActivationFunctionType::kNone) {
- TFLITE_DCHECK_EQ(output_activation_min, 0);
- TFLITE_DCHECK_EQ(output_activation_max, 255);
- }
- DepthwiseConv(input_data, input_dims, input_offset, filter_data, filter_dims,
- filter_offset, bias_data, bias_dims, stride_width,
- stride_height, pad_width, pad_height, depth_multiplier,
- output_offset, output_multiplier, output_shift,
- output_activation_min, output_activation_max, output_data,
- output_dims);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy, for compatibility with old checked-in code.
-template <FusedActivationFunctionType Ac>
-void DepthwiseConv(const uint8* input_data, const Dims<4>& input_dims,
- int32 input_offset, const uint8* filter_data,
- const Dims<4>& filter_dims, int32 filter_offset,
- const int32* bias_data, const Dims<4>& bias_dims, int stride,
- int pad_width, int pad_height, int depth_multiplier,
- int32 output_offset, int32 output_multiplier,
- int output_shift, int32 output_activation_min,
- int32 output_activation_max, uint8* output_data,
- const Dims<4>& output_dims) {
- DepthwiseConv<Ac>(input_data, input_dims, input_offset, filter_data,
- filter_dims, filter_offset, bias_data, bias_dims, stride,
- stride, pad_width, pad_height, depth_multiplier,
- output_offset, output_multiplier, output_shift,
- output_activation_min, output_activation_max, output_data,
- output_dims);
-}
-
} // namespace optimized_ops
} // namespace tflite
diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/legacy_optimized_ops.h b/tensorflow/contrib/lite/kernels/internal/optimized/legacy_optimized_ops.h
index b6151c4..4218be2 100644
--- a/tensorflow/contrib/lite/kernels/internal/optimized/legacy_optimized_ops.h
+++ b/tensorflow/contrib/lite/kernels/internal/optimized/legacy_optimized_ops.h
@@ -19,6 +19,8 @@
#include <sys/types.h>
#include "tensorflow/contrib/lite/kernels/internal/common.h"
+#include "tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_float.h"
+#include "tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_uint8.h"
#include "tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h"
#include "tensorflow/contrib/lite/kernels/internal/reference/legacy_reference_ops.h"
#include "tensorflow/contrib/lite/kernels/internal/types.h"
@@ -28,9 +30,857 @@
// Unoptimized reference ops:
using reference_ops::ArgMax;
+using reference_ops::ArgMinMax;
+using reference_ops::Broadcast4DSlowGreater;
+using reference_ops::Broadcast4DSlowGreaterEqual;
+using reference_ops::Broadcast4DSlowGreaterEqualWithScaling;
+using reference_ops::Broadcast4DSlowGreaterWithScaling;
+using reference_ops::Broadcast4DSlowLess;
+using reference_ops::Broadcast4DSlowLessEqual;
+using reference_ops::Broadcast4DSlowLessEqualWithScaling;
+using reference_ops::Broadcast4DSlowLessWithScaling;
+using reference_ops::BroadcastAdd4DSlow;
+using reference_ops::BroadcastGreater;
+using reference_ops::BroadcastGreaterEqual;
+using reference_ops::BroadcastLess;
+using reference_ops::BroadcastLessEqual;
+using reference_ops::BroadcastMul4DSlow;
+using reference_ops::BroadcastSub4DSlow;
+using reference_ops::Concatenation;
+using reference_ops::ConcatenationWithScaling;
+using reference_ops::DepthConcatenation;
+using reference_ops::Dequantize;
+using reference_ops::Div;
+using reference_ops::FakeQuant;
+using reference_ops::Gather;
+using reference_ops::Greater;
+using reference_ops::GreaterEqual;
+using reference_ops::GreaterEqualWithScaling;
+using reference_ops::GreaterWithScaling;
+using reference_ops::Less;
+using reference_ops::LessEqual;
+using reference_ops::LessEqualWithScaling;
+using reference_ops::LessWithScaling;
+using reference_ops::Mean;
+using reference_ops::RankOneSelect;
using reference_ops::Relu1;
using reference_ops::Relu6;
+using reference_ops::ReluX;
+using reference_ops::Select;
using reference_ops::SpaceToBatchND;
+using reference_ops::Split;
+using reference_ops::StridedSlice;
+using reference_ops::TensorFlowSplit;
+using reference_ops::Transpose;
+
+static constexpr int kDepthwiseReverseShift = -1;
+
+template <typename Scalar, int N>
+VectorMap<Scalar> MapAsVector(Scalar* data, const Dims<N>& dims) {
+ const int size = FlatSize(dims);
+ return VectorMap<Scalar>(data, size, 1);
+}
+
+template <typename Scalar, int N>
+MatrixMap<Scalar> MapAsMatrixWithFirstDimAsRows(Scalar* data,
+ const Dims<N>& dims) {
+ const int rows = dims.sizes[0];
+ int cols = 1;
+ for (int d = 1; d < N; d++) {
+ cols *= dims.sizes[d];
+ }
+ return MatrixMap<Scalar>(data, rows, cols);
+}
+
+template <typename Scalar, int N>
+MatrixMap<Scalar> MapAsMatrixWithLastDimAsCols(Scalar* data,
+ const Dims<N>& dims) {
+ const int cols = dims.sizes[N - 1];
+ int rows = 1;
+ for (int d = 0; d < N - 1; d++) {
+ rows *= dims.sizes[d];
+ }
+ return MatrixMap<Scalar>(data, rows, cols);
+}
+
+template <typename Scalar, int N>
+ArrayMap<Scalar> MapAsArrayWithFirstDimAsRows(Scalar* data,
+ const Dims<N>& dims) {
+ const int rows = dims.sizes[0];
+ int cols = 1;
+ for (int d = 1; d < N; d++) {
+ cols *= dims.sizes[d];
+ }
+ return ArrayMap<Scalar>(data, rows, cols);
+}
+
+// TODO(b/62193649): this function is only needed as long
+// as we have the --variable_batch hack.
+template <typename Scalar, int N>
+MatrixMap<Scalar> MapAsMatrixWithGivenNumberOfRows(Scalar* data,
+ const Dims<N>& dims,
+ int rows) {
+ const int flatsize = FlatSize(dims);
+ TFLITE_DCHECK((flatsize % rows) == 0);
+ const int cols = flatsize / rows;
+ return MatrixMap<Scalar>(data, rows, cols);
+}
+
+inline bool AreSameDims(const Dims<4>& dims1, const Dims<4>& dims2) {
+ for (int i = 0; i < 4; i++) {
+ if (dims1.sizes[i] != dims2.sizes[i]) {
+ return false;
+ }
+ }
+ return true;
+}
+
+inline void DepthwiseConv(const float* input_data, const Dims<4>& input_dims,
+ const float* filter_data, const Dims<4>& filter_dims,
+ const float* bias_data, const Dims<4>& bias_dims,
+ int stride_width, int stride_height,
+ int dilation_width_factor, int dilation_height_factor,
+ int pad_width, int pad_height, int depth_multiplier,
+ float output_activation_min,
+ float output_activation_max, float* output_data,
+ const Dims<4>& output_dims) {
+ tflite::DepthwiseParams op_params;
+ // Padding type is ignored, but still set.
+ op_params.padding_type = PaddingType::kSame;
+ op_params.padding_values.width = pad_width;
+ op_params.padding_values.height = pad_height;
+ op_params.stride_width = stride_width;
+ op_params.stride_height = stride_height;
+ op_params.dilation_width_factor = dilation_width_factor;
+ op_params.dilation_height_factor = dilation_height_factor;
+ op_params.depth_multiplier = depth_multiplier;
+ op_params.float_activation_min = output_activation_min;
+ op_params.float_activation_max = output_activation_max;
+
+ DepthwiseConv(op_params, DimsToShape(input_dims), input_data,
+ DimsToShape(filter_dims), filter_data, DimsToShape(bias_dims),
+ bias_data, DimsToShape(output_dims), output_data);
+}
+
+inline void DepthwiseConv(const float* input_data, const Dims<4>& input_dims,
+ const float* filter_data, const Dims<4>& filter_dims,
+ const float* bias_data, const Dims<4>& bias_dims,
+ int stride_width, int stride_height, int pad_width,
+ int pad_height, int depth_multiplier,
+ float output_activation_min,
+ float output_activation_max, float* output_data,
+ const Dims<4>& output_dims) {
+ DepthwiseConv(input_data, input_dims, filter_data, filter_dims, bias_data,
+ bias_dims, stride_width, stride_height, 1, 1, pad_width,
+ pad_height, depth_multiplier, output_activation_min,
+ output_activation_max, output_data, output_dims);
+}
+
+// legacy, for compatibility with old checked-in code
+template <FusedActivationFunctionType Ac>
+void DepthwiseConv(const float* input_data, const Dims<4>& input_dims,
+ const float* filter_data, const Dims<4>& filter_dims,
+ const float* bias_data, const Dims<4>& bias_dims,
+ int stride_width, int stride_height, int pad_width,
+ int pad_height, int depth_multiplier, float* output_data,
+ const Dims<4>& output_dims) {
+ float output_activation_min, output_activation_max;
+ GetActivationMinMax(Ac, &output_activation_min, &output_activation_max);
+ DepthwiseConv(input_data, input_dims, filter_data, filter_dims, bias_data,
+ bias_dims, stride_width, stride_height, pad_width, pad_height,
+ depth_multiplier, output_activation_min, output_activation_max,
+ output_data, output_dims);
+}
+
+// legacy, for compatibility with old checked-in code
+template <FusedActivationFunctionType Ac>
+void DepthwiseConv(const float* input_data, const Dims<4>& input_dims,
+ const float* filter_data, const Dims<4>& filter_dims,
+ const float* bias_data, const Dims<4>& bias_dims, int stride,
+ int pad_width, int pad_height, int depth_multiplier,
+ float* output_data, const Dims<4>& output_dims) {
+ DepthwiseConv<Ac>(input_data, input_dims, filter_data, filter_dims, bias_data,
+ bias_dims, stride, stride, pad_width, pad_height,
+ depth_multiplier, output_data, output_dims);
+}
+
+inline void DepthwiseConv(const uint8* input_data, const Dims<4>& input_dims,
+ int32 input_offset, const uint8* filter_data,
+ const Dims<4>& filter_dims, int32 filter_offset,
+ const int32* bias_data, const Dims<4>& bias_dims,
+ int stride_width, int stride_height,
+ int dilation_width_factor, int dilation_height_factor,
+ int pad_width, int pad_height, int depth_multiplier,
+ int32 output_offset, int32 output_multiplier,
+ int output_shift, int32 output_activation_min,
+ int32 output_activation_max, uint8* output_data,
+ const Dims<4>& output_dims) {
+ tflite::DepthwiseParams op_params;
+ // Padding type is ignored, but still set.
+ op_params.padding_type = PaddingType::kSame;
+ op_params.padding_values.width = pad_width;
+ op_params.padding_values.height = pad_height;
+ op_params.stride_width = stride_width;
+ op_params.stride_height = stride_height;
+ op_params.dilation_width_factor = dilation_width_factor;
+ op_params.dilation_height_factor = dilation_height_factor;
+ op_params.depth_multiplier = depth_multiplier;
+ op_params.quantized_activation_min = output_activation_min;
+ op_params.quantized_activation_max = output_activation_max;
+ op_params.input_offset = input_offset;
+ op_params.weights_offset = filter_offset;
+ op_params.output_offset = output_offset;
+ op_params.output_multiplier = output_multiplier;
+ // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
+ op_params.output_shift = kDepthwiseReverseShift * output_shift;
+
+ DepthwiseConv(op_params, DimsToShape(input_dims), input_data,
+ DimsToShape(filter_dims), filter_data, DimsToShape(bias_dims),
+ bias_data, DimsToShape(output_dims), output_data);
+}
+
+inline void DepthwiseConv(const uint8* input_data, const Dims<4>& input_dims,
+ int32 input_offset, const uint8* filter_data,
+ const Dims<4>& filter_dims, int32 filter_offset,
+ const int32* bias_data, const Dims<4>& bias_dims,
+ int stride_width, int stride_height, int pad_width,
+ int pad_height, int depth_multiplier,
+ int32 output_offset, int32 output_multiplier,
+ int output_shift, int32 output_activation_min,
+ int32 output_activation_max, uint8* output_data,
+ const Dims<4>& output_dims) {
+ DepthwiseConv(input_data, input_dims, input_offset, filter_data, filter_dims,
+ filter_offset, bias_data, bias_dims, stride_width,
+ stride_height, 1, 1, pad_width, pad_height, depth_multiplier,
+ output_offset, output_multiplier, output_shift,
+ output_activation_min, output_activation_max, output_data,
+ output_dims);
+}
+
+// Legacy, for compatibility with old checked-in code.
+template <FusedActivationFunctionType Ac>
+void DepthwiseConv(const uint8* input_data, const Dims<4>& input_dims,
+ int32 input_offset, const uint8* filter_data,
+ const Dims<4>& filter_dims, int32 filter_offset,
+ const int32* bias_data, const Dims<4>& bias_dims,
+ int stride_width, int stride_height, int pad_width,
+ int pad_height, int depth_multiplier, int32 output_offset,
+ int32 output_multiplier, int output_shift,
+ int32 output_activation_min, int32 output_activation_max,
+ uint8* output_data, const Dims<4>& output_dims) {
+ if (Ac == FusedActivationFunctionType::kNone) {
+ TFLITE_DCHECK_EQ(output_activation_min, 0);
+ TFLITE_DCHECK_EQ(output_activation_max, 255);
+ }
+ DepthwiseConv(input_data, input_dims, input_offset, filter_data, filter_dims,
+ filter_offset, bias_data, bias_dims, stride_width,
+ stride_height, pad_width, pad_height, depth_multiplier,
+ output_offset, output_multiplier, output_shift,
+ output_activation_min, output_activation_max, output_data,
+ output_dims);
+}
+
+// Legacy, for compatibility with old checked-in code.
+template <FusedActivationFunctionType Ac>
+void DepthwiseConv(const uint8* input_data, const Dims<4>& input_dims,
+ int32 input_offset, const uint8* filter_data,
+ const Dims<4>& filter_dims, int32 filter_offset,
+ const int32* bias_data, const Dims<4>& bias_dims, int stride,
+ int pad_width, int pad_height, int depth_multiplier,
+ int32 output_offset, int32 output_multiplier,
+ int output_shift, int32 output_activation_min,
+ int32 output_activation_max, uint8* output_data,
+ const Dims<4>& output_dims) {
+ DepthwiseConv<Ac>(input_data, input_dims, input_offset, filter_data,
+ filter_dims, filter_offset, bias_data, bias_dims, stride,
+ stride, pad_width, pad_height, depth_multiplier,
+ output_offset, output_multiplier, output_shift,
+ output_activation_min, output_activation_max, output_data,
+ output_dims);
+}
+
+inline void AddBiasAndEvalActivationFunction(const float* bias_data,
+ const Dims<4>& bias_dims,
+ float* array_data,
+ const Dims<4>& array_dims,
+ float output_activation_min,
+ float output_activation_max) {
+ AddBiasAndEvalActivationFunction(output_activation_min, output_activation_max,
+ DimsToShape(bias_dims), bias_data,
+ DimsToShape(array_dims), array_data);
+}
+
+// legacy, for compatibility with old checked-in code
+template <FusedActivationFunctionType Ac>
+void AddBiasAndEvalActivationFunction(const float* bias_data,
+ const Dims<4>& bias_dims,
+ float* array_data,
+ const Dims<4>& array_dims) {
+ float output_activation_min, output_activation_max;
+ GetActivationMinMax(Ac, &output_activation_min, &output_activation_max);
+ AddBiasAndEvalActivationFunction(bias_data, bias_dims, array_data, array_dims,
+ output_activation_min,
+ output_activation_max);
+}
+
+inline void FullyConnected(const float* input_data, const Dims<4>& input_dims,
+ const float* weights_data,
+ const Dims<4>& weights_dims, const float* bias_data,
+ const Dims<4>& bias_dims,
+ float output_activation_min,
+ float output_activation_max, float* output_data,
+ const Dims<4>& output_dims) {
+ tflite::FullyConnectedParams op_params;
+ op_params.float_activation_min = output_activation_min;
+ op_params.float_activation_max = output_activation_max;
+
+ FullyConnected(op_params, DimsToShape(input_dims), input_data,
+ DimsToShape(weights_dims), weights_data,
+ DimsToShape(bias_dims), bias_data, DimsToShape(output_dims),
+ output_data);
+}
+
+// legacy, for compatibility with old checked-in code
+template <FusedActivationFunctionType Ac>
+void FullyConnected(const float* input_data, const Dims<4>& input_dims,
+ const float* weights_data, const Dims<4>& weights_dims,
+ const float* bias_data, const Dims<4>& bias_dims,
+ float* output_data, const Dims<4>& output_dims) {
+ float output_activation_min, output_activation_max;
+ GetActivationMinMax(Ac, &output_activation_min, &output_activation_max);
+ FullyConnected(input_data, input_dims, weights_data, weights_dims, bias_data,
+ bias_dims, output_activation_min, output_activation_max,
+ output_data, output_dims);
+}
+
+inline void FullyConnected(const uint8* input_data, const Dims<4>& input_dims,
+ int32 input_offset, const uint8* filter_data,
+ const Dims<4>& filter_dims, int32 filter_offset,
+ const int32* bias_data, const Dims<4>& bias_dims,
+ int32 output_offset, int32 output_multiplier,
+ int output_shift, int32 output_activation_min,
+ int32 output_activation_max, uint8* output_data,
+ const Dims<4>& output_dims,
+ gemmlowp::GemmContext* gemm_context) {
+ tflite::FullyConnectedParams op_params;
+ op_params.input_offset = input_offset;
+ op_params.weights_offset = filter_offset;
+ op_params.output_offset = output_offset;
+ op_params.output_multiplier = output_multiplier;
+ // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
+ op_params.output_shift = kReverseShift * output_shift;
+ op_params.quantized_activation_min = output_activation_min;
+ op_params.quantized_activation_max = output_activation_max;
+
+ FullyConnected(op_params, DimsToShape(input_dims), input_data,
+ DimsToShape(filter_dims), filter_data, DimsToShape(bias_dims),
+ bias_data, DimsToShape(output_dims), output_data,
+ gemm_context);
+}
+
+inline void FullyConnected(
+ const uint8* input_data, const Dims<4>& input_dims, int32 input_offset,
+ const uint8* filter_data, const Dims<4>& filter_dims, int32 filter_offset,
+ const int32* bias_data_int32, const Dims<4>& bias_dims, int32 output_offset,
+ int32 output_multiplier, int output_shift, int32 output_activation_min,
+ int32 output_activation_max, int16* output_data, const Dims<4>& output_dims,
+ gemmlowp::GemmContext* gemm_context) {
+ tflite::FullyConnectedParams op_params;
+ op_params.input_offset = input_offset;
+ op_params.weights_offset = filter_offset;
+ op_params.output_offset = output_offset;
+ op_params.output_multiplier = output_multiplier;
+ // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
+ op_params.output_shift = kReverseShift * output_shift;
+ op_params.quantized_activation_min = output_activation_min;
+ op_params.quantized_activation_max = output_activation_max;
+
+ FullyConnected(op_params, DimsToShape(input_dims), input_data,
+ DimsToShape(filter_dims), filter_data, DimsToShape(bias_dims),
+ bias_data_int32, DimsToShape(output_dims), output_data,
+ gemm_context);
+}
+
+// legacy, for compatibility with old checked-in code
+template <FusedActivationFunctionType Ac>
+void FullyConnected(const uint8* input_data, const Dims<4>& input_dims,
+ int32 input_offset, const uint8* filter_data,
+ const Dims<4>& filter_dims, int32 filter_offset,
+ const int32* bias_data, const Dims<4>& bias_dims,
+ int32 output_offset, int32 output_multiplier,
+ int output_shift, int32 output_activation_min,
+ int32 output_activation_max, uint8* output_data,
+ const Dims<4>& output_dims,
+ gemmlowp::GemmContext* gemm_context) {
+ static_assert(Ac == FusedActivationFunctionType::kNone ||
+ Ac == FusedActivationFunctionType::kRelu ||
+ Ac == FusedActivationFunctionType::kRelu6 ||
+ Ac == FusedActivationFunctionType::kRelu1,
+ "");
+ FullyConnected(input_data, input_dims, input_offset, filter_data, filter_dims,
+ filter_offset, bias_data, bias_dims, output_offset,
+ output_multiplier, output_shift, output_activation_min,
+ output_activation_max, output_data, output_dims, gemm_context);
+}
+
+inline void ShuffledFullyConnected(
+ const uint8* input_data, const Dims<4>& input_dims,
+ const uint8* shuffled_weights_data, const Dims<4>& weights_dims,
+ const int32* bias_data, const Dims<4>& bias_dims, int32 output_multiplier,
+ int output_shift, int32 output_activation_min, int32 output_activation_max,
+ int16* output_data, const Dims<4>& output_dims,
+ uint8* shuffled_input_workspace_data, gemmlowp::GemmContext* gemm_context) {
+ tflite::FullyConnectedParams op_params;
+ op_params.output_multiplier = output_multiplier;
+ // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
+ op_params.output_shift = kReverseShift * output_shift;
+ op_params.quantized_activation_min = output_activation_min;
+ op_params.quantized_activation_max = output_activation_max;
+
+ ShuffledFullyConnected(op_params, DimsToShape(input_dims), input_data,
+ DimsToShape(weights_dims), shuffled_weights_data,
+ DimsToShape(bias_dims), bias_data,
+ DimsToShape(output_dims), output_data,
+ shuffled_input_workspace_data, gemm_context);
+}
+
+template <typename T>
+inline void ExtractPatchIntoBufferColumn(
+ const Dims<4>& input_dims, int w, int h, int b, int kheight, int kwidth,
+ int stride_width, int stride_height, int pad_width, int pad_height,
+ int in_width, int in_height, int in_depth, int single_buffer_length,
+ int buffer_id, const T* in_data, T* conv_buffer_data, uint8 zero_byte) {
+ ExtractPatchIntoBufferColumn(
+ DimsToShape(input_dims), w, h, b, kheight, kwidth, stride_width,
+ stride_height, pad_width, pad_height, in_width, in_height, in_depth,
+ single_buffer_length, buffer_id, in_data, conv_buffer_data, zero_byte);
+}
+
+template <typename T>
+void DilatedIm2col(const T* input_data, const Dims<4>& input_dims,
+ const Dims<4>& filter_dims, int stride_width,
+ int stride_height, int dilation_width_factor,
+ int dilation_height_factor, int pad_width, int pad_height,
+ const Dims<4>& output_dims, uint8 zero_byte,
+ T* im2col_data) {
+ tflite::ConvParams op_params;
+ // Padding type is ignored, but still set.
+ op_params.padding_type = PaddingType::kSame;
+ op_params.padding_values.width = pad_width;
+ op_params.padding_values.height = pad_height;
+ op_params.stride_width = stride_width;
+ op_params.stride_height = stride_height;
+ op_params.dilation_width_factor = dilation_width_factor;
+ op_params.dilation_height_factor = dilation_height_factor;
+
+ DilatedIm2col(op_params, zero_byte, DimsToShape(input_dims), input_data,
+ DimsToShape(filter_dims), DimsToShape(output_dims),
+ im2col_data);
+}
+
+template <typename T>
+void Im2col(const T* input_data, const Dims<4>& input_dims, int stride_width,
+ int stride_height, int pad_width, int pad_height, int kheight,
+ int kwidth, uint8 zero_byte, T* output_data,
+ const Dims<4>& output_dims) {
+ tflite::ConvParams op_params;
+ // Padding type is ignored, but still set.
+ op_params.padding_type = PaddingType::kSame;
+ op_params.padding_values.width = pad_width;
+ op_params.padding_values.height = pad_height;
+ op_params.stride_width = stride_width;
+ op_params.stride_height = stride_height;
+ op_params.dilation_width_factor = 1;
+ op_params.dilation_height_factor = 1;
+
+ Im2col(op_params, kheight, kwidth, zero_byte, DimsToShape(input_dims),
+ input_data, DimsToShape(output_dims), output_data);
+}
+
+// legacy, for compatibility with old checked-in code
+template <typename T>
+void Im2col(const T* input_data, const Dims<4>& input_dims, int stride,
+ int pad_width, int pad_height, int kheight, int kwidth,
+ uint8 zero_byte, T* output_data, const Dims<4>& output_dims) {
+ Im2col(input_data, input_dims, stride, stride, pad_width, pad_height, kheight,
+ kwidth, zero_byte, output_data, output_dims);
+}
+
+inline void Conv(const float* input_data, const Dims<4>& input_dims,
+ const float* filter_data, const Dims<4>& filter_dims,
+ const float* bias_data, const Dims<4>& bias_dims,
+ int stride_width, int stride_height, int dilation_width_factor,
+ int dilation_height_factor, int pad_width, int pad_height,
+ float output_activation_min, float output_activation_max,
+ float* output_data, const Dims<4>& output_dims,
+ float* im2col_data, const Dims<4>& im2col_dims) {
+ tflite::ConvParams op_params;
+ // Padding type is ignored, but still set.
+ op_params.padding_type = PaddingType::kSame;
+ op_params.padding_values.width = pad_width;
+ op_params.padding_values.height = pad_height;
+ op_params.stride_width = stride_width;
+ op_params.stride_height = stride_height;
+ op_params.dilation_width_factor = dilation_width_factor;
+ op_params.dilation_height_factor = dilation_height_factor;
+ op_params.float_activation_min = output_activation_min;
+ op_params.float_activation_max = output_activation_max;
+
+ Conv(op_params, DimsToShape(input_dims), input_data, DimsToShape(filter_dims),
+ filter_data, DimsToShape(bias_dims), bias_data, DimsToShape(output_dims),
+ output_data, DimsToShape(im2col_dims), im2col_data);
+}
+
+inline void HybridConv(const int8_t* input_data, const Dims<4>& input_dims,
+ const int8_t* filter_data, const Dims<4>& filter_dims,
+ const float* bias_data, const Dims<4>& bias_dims,
+ int stride_width, int stride_height, int pad_width,
+ int pad_height, float* scaling_factors_ptr,
+ float output_activation_min, float output_activation_max,
+ float* output_data, const Dims<4>& output_dims,
+ int8_t* im2col_data, const Dims<4>& im2col_dims) {
+ tflite::ConvParams op_params;
+ // Padding type is ignored, but still set.
+ op_params.padding_type = PaddingType::kSame;
+ op_params.padding_values.width = pad_width;
+ op_params.padding_values.height = pad_height;
+ op_params.stride_width = stride_width;
+ op_params.stride_height = stride_height;
+ op_params.float_activation_min = output_activation_min;
+ op_params.float_activation_max = output_activation_max;
+
+ HybridConv(op_params, scaling_factors_ptr, DimsToShape(input_dims),
+ input_data, DimsToShape(filter_dims), filter_data,
+ DimsToShape(bias_dims), bias_data, DimsToShape(output_dims),
+ output_data, DimsToShape(im2col_dims), im2col_data);
+}
+
+template <FusedActivationFunctionType Ac>
+void Conv(const float* input_data, const Dims<4>& input_dims,
+ const float* filter_data, const Dims<4>& filter_dims,
+ const float* bias_data, const Dims<4>& bias_dims, int stride_width,
+ int stride_height, int dilation_width_factor,
+ int dilation_height_factor, int pad_width, int pad_height,
+ float* output_data, const Dims<4>& output_dims, float* im2col_data,
+ const Dims<4>& im2col_dims) {
+ float output_activation_min, output_activation_max;
+ GetActivationMinMax(Ac, &output_activation_min, &output_activation_max);
+ Conv(input_data, input_dims, filter_data, filter_dims, bias_data, bias_dims,
+ stride_width, stride_height, dilation_width_factor,
+ dilation_height_factor, pad_width, pad_height, output_activation_min,
+ output_activation_max, output_data, output_dims, im2col_data,
+ im2col_dims);
+}
+
+// legacy, for compatibility with old checked-in code
+template <FusedActivationFunctionType Ac>
+void Conv(const float* input_data, const Dims<4>& input_dims,
+ const float* filter_data, const Dims<4>& filter_dims,
+ const float* bias_data, const Dims<4>& bias_dims, int stride_width,
+ int stride_height, int pad_width, int pad_height, float* output_data,
+ const Dims<4>& output_dims, float* im2col_data,
+ const Dims<4>& im2col_dims) {
+ float output_activation_min, output_activation_max;
+ GetActivationMinMax(Ac, &output_activation_min, &output_activation_max);
+ Conv(input_data, input_dims, filter_data, filter_dims, bias_data, bias_dims,
+ stride_width, stride_height, 1, 1, pad_width, pad_height,
+ output_activation_min, output_activation_max, output_data, output_dims,
+ im2col_data, im2col_dims);
+}
+
+// legacy, for compatibility with old checked-in code
+template <FusedActivationFunctionType Ac>
+void Conv(const float* input_data, const Dims<4>& input_dims,
+ const float* filter_data, const Dims<4>& filter_dims,
+ const float* bias_data, const Dims<4>& bias_dims, int stride,
+ int pad_width, int pad_height, float* output_data,
+ const Dims<4>& output_dims, float* im2col_data,
+ const Dims<4>& im2col_dims) {
+ Conv<Ac>(input_data, input_dims, filter_data, filter_dims, bias_data,
+ bias_dims, stride, stride, 1, 1, pad_width, pad_height, output_data,
+ output_dims, im2col_data, im2col_dims);
+}
+
+inline void Conv(const uint8* input_data, const Dims<4>& input_dims,
+ int32 input_offset, const uint8* filter_data,
+ const Dims<4>& filter_dims, int32 filter_offset,
+ const int32* bias_data, const Dims<4>& bias_dims,
+ int stride_width, int stride_height, int dilation_width_factor,
+ int dilation_height_factor, int pad_width, int pad_height,
+ int32 output_offset, int32 output_multiplier, int output_shift,
+ int32 output_activation_min, int32 output_activation_max,
+ uint8* output_data, const Dims<4>& output_dims,
+ uint8* im2col_data, const Dims<4>& im2col_dims,
+ gemmlowp::GemmContext* gemm_context) {
+ tflite::ConvParams op_params;
+ // Padding type is ignored, but still set.
+ op_params.padding_type = PaddingType::kSame;
+ op_params.padding_values.width = pad_width;
+ op_params.padding_values.height = pad_height;
+ op_params.stride_width = stride_width;
+ op_params.stride_height = stride_height;
+ op_params.dilation_width_factor = dilation_width_factor;
+ op_params.dilation_height_factor = dilation_height_factor;
+ op_params.input_offset = input_offset;
+ op_params.weights_offset = filter_offset;
+ op_params.output_offset = output_offset;
+ op_params.output_multiplier = output_multiplier;
+ // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
+ op_params.output_shift = kReverseShift * output_shift;
+ op_params.quantized_activation_min = output_activation_min;
+ op_params.quantized_activation_max = output_activation_max;
+
+ Conv(op_params, DimsToShape(input_dims), input_data, DimsToShape(filter_dims),
+ filter_data, DimsToShape(bias_dims), bias_data, DimsToShape(output_dims),
+ output_data, DimsToShape(im2col_dims), im2col_data, gemm_context);
+}
+
+inline void Conv(const uint8* input_data, const Dims<4>& input_dims,
+ int32 input_offset, const uint8* filter_data,
+ const Dims<4>& filter_dims, int32 filter_offset,
+ const int32* bias_data, const Dims<4>& bias_dims,
+ int stride_width, int stride_height, int pad_width,
+ int pad_height, int32 output_offset, int32 output_multiplier,
+ int output_shift, int32 output_activation_min,
+ int32 output_activation_max, uint8* output_data,
+ const Dims<4>& output_dims, uint8* im2col_data,
+ const Dims<4>& im2col_dims,
+ gemmlowp::GemmContext* gemm_context) {
+ Conv(input_data, input_dims, input_offset, filter_data, filter_dims,
+ filter_offset, bias_data, bias_dims, stride_width, stride_height, 1, 1,
+ pad_width, pad_height, output_offset, output_multiplier, output_shift,
+ output_activation_min, output_activation_max, output_data, output_dims,
+ im2col_data, im2col_dims, gemm_context);
+}
+
+// legacy, for compatibility with old checked-in code
+template <FusedActivationFunctionType Ac>
+inline void Conv(const uint8* input_data, const Dims<4>& input_dims,
+ int32 input_offset, const uint8* filter_data,
+ const Dims<4>& filter_dims, int32 filter_offset,
+ const int32* bias_data, const Dims<4>& bias_dims,
+ int stride_width, int stride_height, int pad_width,
+ int pad_height, int32 output_offset, int32 output_multiplier,
+ int output_shift, int32 output_activation_min,
+ int32 output_activation_max, uint8* output_data,
+ const Dims<4>& output_dims, uint8* im2col_data,
+ const Dims<4>& im2col_dims,
+ gemmlowp::GemmContext* gemm_context) {
+ static_assert(Ac == FusedActivationFunctionType::kNone ||
+ Ac == FusedActivationFunctionType::kRelu ||
+ Ac == FusedActivationFunctionType::kRelu6 ||
+ Ac == FusedActivationFunctionType::kRelu1,
+ "");
+ if (Ac == FusedActivationFunctionType::kNone) {
+ TFLITE_DCHECK_EQ(output_activation_min, 0);
+ TFLITE_DCHECK_EQ(output_activation_max, 255);
+ }
+ Conv(input_data, input_dims, input_offset, filter_data, filter_dims,
+ filter_offset, bias_data, bias_dims, stride_width, stride_height,
+ pad_width, pad_height, output_offset, output_multiplier, output_shift,
+ output_activation_min, output_activation_max, output_data, output_dims,
+ im2col_data, im2col_dims, gemm_context);
+}
+
+// legacy, for compatibility with old checked-in code
+template <FusedActivationFunctionType Ac>
+void Conv(const uint8* input_data, const Dims<4>& input_dims,
+ int32 input_offset, const uint8* filter_data,
+ const Dims<4>& filter_dims, int32 filter_offset,
+ const int32* bias_data, const Dims<4>& bias_dims, int stride,
+ int pad_width, int pad_height, int32 output_offset,
+ int32 output_multiplier, int output_shift,
+ int32 output_activation_min, int32 output_activation_max,
+ uint8* output_data, const Dims<4>& output_dims, uint8* im2col_data,
+ const Dims<4>& im2col_dims, gemmlowp::GemmContext* gemm_context) {
+ static_assert(Ac == FusedActivationFunctionType::kNone ||
+ Ac == FusedActivationFunctionType::kRelu ||
+ Ac == FusedActivationFunctionType::kRelu6 ||
+ Ac == FusedActivationFunctionType::kRelu1,
+ "");
+ Conv(input_data, input_dims, input_offset, filter_data, filter_dims,
+ filter_offset, bias_data, bias_dims, stride, stride, pad_width,
+ pad_height, output_offset, output_multiplier, output_shift,
+ output_activation_min, output_activation_max, output_data, output_dims,
+ im2col_data, im2col_dims, gemm_context);
+}
+
+// legacy, for compatibility with old checked-in code
+template <FusedActivationFunctionType Ac, typename T>
+void Im2col(const T* input_data, const Dims<4>& input_dims, int stride,
+ int pad_width, int pad_height, int kheight, int kwidth,
+ uint8 zero_byte, T* output_data, const Dims<4>& output_dims) {
+ Im2col(input_data, input_dims, stride, stride, pad_width, pad_height, kheight,
+ kwidth, zero_byte, output_data, output_dims);
+}
+
+// legacy, for compatibility with old checked-in code
+template <FusedActivationFunctionType Ac>
+void ConvAsGemm(const float* input_data, const Dims<4>& input_dims,
+ const float* filter_data, const Dims<4>& filter_dims,
+ const float* bias_data, const Dims<4>& bias_dims,
+ float* output_data, const Dims<4>& output_dims) {
+ gemmlowp::ScopedProfilingLabel label("ConvAsGemm");
+
+ const auto input_matrix_map =
+ MapAsMatrixWithFirstDimAsRows(input_data, input_dims);
+ const auto filter_matrix_map =
+ MapAsMatrixWithLastDimAsCols(filter_data, filter_dims);
+ auto output_matrix_map =
+ MapAsMatrixWithFirstDimAsRows(output_data, output_dims);
+
+ Gemm(filter_matrix_map.transpose(), input_matrix_map, &output_matrix_map);
+
+ AddBiasAndEvalActivationFunction<Ac>(bias_data, bias_dims, output_data,
+ output_dims);
+}
+
+// legacy, for compatibility with old checked-in code
+template <FusedActivationFunctionType Ac>
+void ConvAsGemm(const uint8* input_data, const Dims<4>& input_dims,
+ int32 input_offset, const uint8* filter_data,
+ const Dims<4>& filter_dims, int32 filter_offset,
+ const int32* bias_data, const Dims<4>& bias_dims,
+ int32 output_offset, int32 output_multiplier, int output_shift,
+ int32 output_activation_min, int32 output_activation_max,
+ uint8* output_data, const Dims<4>& output_dims,
+ gemmlowp::GemmContext* gemm_context) {
+ gemmlowp::ScopedProfilingLabel label("ConvAsGemm/8bit");
+ static_assert(Ac == FusedActivationFunctionType::kNone ||
+ Ac == FusedActivationFunctionType::kRelu ||
+ Ac == FusedActivationFunctionType::kRelu6 ||
+ Ac == FusedActivationFunctionType::kRelu1,
+ "");
+ const int input_rows = input_dims.sizes[0];
+ const int input_cols = FlatSizeSkipDim(input_dims, 0);
+ const int filter_rows = filter_dims.sizes[3];
+ const int filter_cols = FlatSizeSkipDim(filter_dims, 3);
+ const int output_rows = output_dims.sizes[0];
+ const int output_cols = FlatSizeSkipDim(output_dims, 0);
+ TFLITE_DCHECK_EQ(output_rows, filter_rows);
+ TFLITE_DCHECK_EQ(output_cols, input_cols);
+ TFLITE_DCHECK_EQ(filter_cols, input_rows);
+ TFLITE_DCHECK_EQ(bias_dims.sizes[0], output_rows);
+ TFLITE_DCHECK_EQ(bias_dims.sizes[1], 1);
+ TFLITE_DCHECK_EQ(bias_dims.sizes[2], 1);
+ TFLITE_DCHECK_EQ(bias_dims.sizes[3], 1);
+ gemmlowp::MatrixMap<const uint8, gemmlowp::MapOrder::RowMajor> filter_matrix(
+ filter_data, output_rows, filter_cols, filter_cols);
+ gemmlowp::MatrixMap<const uint8, gemmlowp::MapOrder::ColMajor> input_matrix(
+ input_data, filter_cols, output_cols, filter_cols);
+ gemmlowp::MatrixMap<uint8, gemmlowp::MapOrder::ColMajor> output_matrix(
+ output_data, output_rows, output_cols, output_rows);
+ const auto& output_pipeline = GemmlowpOutputPipeline::MakeExp(
+ bias_data, output_rows, output_offset, output_multiplier, -output_shift,
+ output_activation_min, output_activation_max);
+ gemmlowp::GemmWithOutputPipeline<uint8, uint8,
+ gemmlowp::L8R8WithLhsNonzeroBitDepthParams>(
+ gemm_context, filter_matrix, input_matrix, &output_matrix, filter_offset,
+ input_offset, output_pipeline);
+}
+
+inline void TransposeConv(const float* input_data, const Dims<4>& input_dims,
+ const float* filter_data, const Dims<4>& filter_dims,
+ int stride_width, int stride_height, int pad_width,
+ int pad_height, float* output_data,
+ const Dims<4>& output_dims, float* im2col_data,
+ const Dims<4>& im2col_dims) {
+ tflite::ConvParams op_params;
+ // Padding type is ignored, but still set.
+ op_params.padding_type = PaddingType::kSame;
+ op_params.padding_values.width = pad_width;
+ op_params.padding_values.height = pad_height;
+ op_params.stride_width = stride_width;
+ op_params.stride_height = stride_height;
+
+ TransposeConv(op_params, DimsToShape(input_dims), input_data,
+ DimsToShape(filter_dims), filter_data, DimsToShape(output_dims),
+ output_data, DimsToShape(im2col_dims), im2col_data);
+}
+
+template <typename T>
+void TransposeIm2col(const T* input_data, const Dims<4>& input_dims,
+ const Dims<4>& filter_dims, int stride_width,
+ int stride_height, int pad_width, int pad_height,
+ const Dims<4>& output_dims, uint8 zero_byte,
+ T* im2col_data) {
+ tflite::ConvParams op_params;
+ // Padding type is ignored, but still set.
+ op_params.padding_type = PaddingType::kSame;
+ op_params.padding_values.width = pad_width;
+ op_params.padding_values.height = pad_height;
+ op_params.stride_width = stride_width;
+ op_params.stride_height = stride_height;
+
+ TransposeIm2col(op_params, zero_byte, DimsToShape(input_dims), input_data,
+ DimsToShape(filter_dims), DimsToShape(output_dims),
+ im2col_data);
+}
+
+inline void LstmCell(const float* input_data, const Dims<4>& input_dims,
+ const float* prev_activ_data,
+ const Dims<4>& prev_activ_dims, const float* weights_data,
+ const Dims<4>& weights_dims, const float* bias_data,
+ const Dims<4>& bias_dims, const float* prev_state_data,
+ const Dims<4>& prev_state_dims, float* output_state_data,
+ const Dims<4>& output_state_dims, float* output_activ_data,
+ const Dims<4>& output_activ_dims, float* concat_temp_data,
+ const Dims<4>& concat_temp_dims, float* activ_temp_data,
+ const Dims<4>& activ_temp_dims) {
+ tflite::LstmCellParams op_params;
+ // Float LSTM cell does not need parameters to be set: leave untouched.
+
+ LstmCell(op_params, DimsToShape(input_dims), input_data,
+ DimsToShape(prev_activ_dims), prev_activ_data,
+ DimsToShape(weights_dims), weights_data, DimsToShape(bias_dims),
+ bias_data, DimsToShape(prev_state_dims), prev_state_data,
+ DimsToShape(output_state_dims), output_state_data,
+ DimsToShape(output_activ_dims), output_activ_data,
+ DimsToShape(concat_temp_dims), concat_temp_data,
+ DimsToShape(activ_temp_dims), activ_temp_data);
+}
+
+template <int StateIntegerBits>
+void LstmCell(const uint8* input_data_uint8, const Dims<4>& input_dims,
+ const uint8* prev_activ_data_uint8,
+ const Dims<4>& prev_activ_dims, const uint8* weights_data_uint8,
+ const Dims<4>& weights_dims, const int32* bias_data_int32,
+ const Dims<4>& bias_dims, const int16* prev_state_data_int16,
+ const Dims<4>& prev_state_dims, int16* output_state_data_int16,
+ const Dims<4>& output_state_dims, uint8* output_activ_data_uint8,
+ const Dims<4>& output_activ_dims, uint8* concat_temp_data_uint8,
+ const Dims<4>& concat_temp_dims, int16* activ_temp_data_int16,
+ const Dims<4>& activ_temp_dims, int32 weights_zero_point,
+ int32 accum_multiplier, int accum_shift,
+ gemmlowp::GemmContext* gemm_context) {
+ tflite::LstmCellParams op_params;
+ op_params.weights_zero_point = weights_zero_point;
+ op_params.accum_multiplier = accum_multiplier;
+ op_params.accum_shift = accum_shift;
+
+ LstmCell<StateIntegerBits>(
+ op_params, DimsToShape(input_dims), input_data_uint8,
+ DimsToShape(prev_activ_dims), prev_activ_data_uint8,
+ DimsToShape(weights_dims), weights_data_uint8, DimsToShape(bias_dims),
+ bias_data_int32, DimsToShape(prev_state_dims), prev_state_data_int16,
+ DimsToShape(output_state_dims), output_state_data_int16,
+ DimsToShape(output_activ_dims), output_activ_data_uint8,
+ DimsToShape(concat_temp_dims), concat_temp_data_uint8,
+ DimsToShape(activ_temp_dims), activ_temp_data_int16, gemm_context);
+}
+
+template <typename T>
+void BroadcastDiv(const T* input1_data, const Dims<4>& input1_dims,
+ const T* input2_data, const Dims<4>& input2_dims,
+ T output_activation_min, T output_activation_max,
+ T* output_data, const Dims<4>& output_dims) {
+ tflite::ArithmeticParams op_params;
+ SetActivationParams(output_activation_min, output_activation_max, &op_params);
+
+ BroadcastDiv4DSlow(op_params, DimsToShape(input1_dims), input1_data,
+ DimsToShape(input2_dims), input2_data,
+ DimsToShape(output_dims), output_data);
+}
template <FusedActivationFunctionType Ac>
void L2Normalization(const float* input_data, const RuntimeShape& input_shape,
@@ -574,6 +1424,14 @@
filter_width, filter_height, output_data, output_dims);
}
+inline void Softmax(const float* input_data, const RuntimeShape& input_shape,
+ float beta, float* output_data,
+ const RuntimeShape& output_shape) {
+ SoftmaxParams params;
+ params.beta = beta;
+ Softmax(params, input_shape, input_data, output_shape, output_data);
+}
+
inline void Softmax(const float* input_data, const Dims<4>& input_dims,
float beta, float* output_data,
const Dims<4>& output_dims) {
@@ -581,6 +1439,16 @@
DimsToShape(output_dims));
}
+inline void Softmax(const uint8* input_data, const RuntimeShape& input_shape,
+ int32 input_beta_multiplier, int32 input_beta_left_shift,
+ int diff_min, uint8* output_data,
+ const RuntimeShape& output_shape) {
+ SoftmaxParams params;
+ params.input_multiplier = input_beta_multiplier;
+ params.input_left_shift = input_beta_left_shift;
+ params.diff_min = diff_min;
+ Softmax(params, input_shape, input_data, output_shape, output_data);
+}
inline void Softmax(const uint8* input_data, const Dims<4>& input_dims,
int32 input_beta_multiplier, int32 input_beta_left_shift,
int diff_min, uint8* output_data,
@@ -590,12 +1458,33 @@
DimsToShape(output_dims));
}
+inline void LogSoftmax(const float* input_data, const RuntimeShape& input_shape,
+ float* output_data, const RuntimeShape& output_shape) {
+ SoftmaxParams params;
+ // No params currently used for float LogSoftmax.
+ LogSoftmax(params, input_shape, input_data, output_shape, output_data);
+}
+
inline void LogSoftmax(const float* input_data, const Dims<4>& input_dims,
float* output_data, const Dims<4>& output_dims) {
LogSoftmax(input_data, DimsToShape(input_dims), output_data,
DimsToShape(output_dims));
}
+inline void LogSoftmax(const uint8* input_data, const RuntimeShape& input_shape,
+ int32 input_multiplier, int32 input_left_shift,
+ int32 reverse_scaling_divisor,
+ int32 reverse_scaling_right_shift, int diff_min,
+ uint8* output_data, const RuntimeShape& output_shape) {
+ SoftmaxParams params;
+ params.input_multiplier = input_multiplier;
+ params.input_left_shift = input_left_shift;
+ params.reverse_scaling_divisor = reverse_scaling_divisor;
+ params.reverse_scaling_right_shift = reverse_scaling_right_shift;
+ params.diff_min = diff_min;
+ LogSoftmax(params, input_shape, input_data, output_shape, output_data);
+}
+
inline void LogSoftmax(const uint8* input_data, const Dims<4>& input_dims,
int32 input_multiplier, int32 input_left_shift,
int32 reverse_scaling_divisor,
@@ -607,6 +1496,18 @@
DimsToShape(output_dims));
}
+inline void Logistic(const uint8* input_data, const RuntimeShape& input_shape,
+ int32 input_zero_point, int32 input_range_radius,
+ int32 input_multiplier, int input_left_shift,
+ uint8* output_data, const RuntimeShape& output_shape) {
+ LogisticParams params;
+ params.input_zero_point = input_zero_point;
+ params.input_range_radius = input_range_radius;
+ params.input_multiplier = input_multiplier;
+ params.input_left_shift = input_left_shift;
+ Logistic(params, input_shape, input_data, output_shape, output_data);
+}
+
inline void Logistic(const float* input_data, const Dims<4>& input_dims,
float* output_data, const Dims<4>& output_dims) {
Logistic(DimsToShape(input_dims), input_data, DimsToShape(output_dims),
@@ -622,6 +1523,20 @@
DimsToShape(output_dims));
}
+inline void Logistic(const RuntimeShape& input_shape, const int16* input_data,
+ const RuntimeShape& output_shape, int16* output_data) {
+ LogisticParams params;
+ // No params currently needed by int16 Logistic.
+ Logistic(params, input_shape, input_data, output_shape, output_data);
+}
+
+inline void Logistic(const int16* input_data, const RuntimeShape& input_shape,
+ int16* output_data, const RuntimeShape& output_shape) {
+ LogisticParams params;
+ // No params currently needed by int16 Logistic.
+ Logistic(params, input_shape, input_data, output_shape, output_data);
+}
+
inline void Logistic(const int16* input_data, const Dims<4>& input_dims,
int16* output_data, const Dims<4>& output_dims) {
Logistic(input_data, DimsToShape(input_dims), output_data,
@@ -634,6 +1549,18 @@
output_data);
}
+inline void Tanh(const uint8* input_data, const RuntimeShape& input_shape,
+ int32 input_zero_point, int32 input_range_radius,
+ int32 input_multiplier, int input_left_shift,
+ uint8* output_data, const RuntimeShape& output_shape) {
+ TanhParams params;
+ params.input_zero_point = input_zero_point;
+ params.input_range_radius = input_range_radius;
+ params.input_multiplier = input_multiplier;
+ params.input_left_shift = input_left_shift;
+ Tanh(params, input_shape, input_data, output_shape, output_data);
+}
+
inline void Tanh(const uint8* input_data, const Dims<4>& input_dims,
int32 input_zero_point, int32 input_range_radius,
int32 input_multiplier, int input_left_shift,
@@ -643,6 +1570,14 @@
DimsToShape(output_dims));
}
+inline void Tanh(const int16* input_data, const RuntimeShape& input_shape,
+ int input_left_shift, int16* output_data,
+ const RuntimeShape& output_shape) {
+ TanhParams params;
+ params.input_left_shift = input_left_shift;
+ Tanh(params, input_shape, input_data, output_shape, output_data);
+}
+
inline void Tanh(const int16* input_data, const Dims<4>& input_dims,
int input_left_shift, int16* output_data,
const Dims<4>& output_dims) {
@@ -777,7 +1712,6 @@
DimsToShape(output_dims), output_data);
}
-// Legacy Dims<4>.
inline void LocalResponseNormalization(const float* input_data,
const Dims<4>& input_dims, int range,
float bias, float alpha, float beta,
@@ -793,7 +1727,6 @@
DimsToShape(output_dims), output_data);
}
-// Legacy Dims<4> version.
template <typename SrcT, typename DstT>
void Cast(const SrcT* input_data, const Dims<4>& input_dims, DstT* output_data,
const Dims<4>& output_dims) {
@@ -801,14 +1734,12 @@
output_data);
}
-// Legacy Dims<4> version.
inline void Floor(const float* input_data, const Dims<4>& input_dims,
float* output_data, const Dims<4>& output_dims) {
Floor(DimsToShape(input_dims), input_data, DimsToShape(output_dims),
output_data);
}
-// Legacy Dims<4>
inline void ResizeBilinear(const float* input_data, const Dims<4>& input_dims,
const int32* output_size_data,
const Dims<4>& output_size_dims, float* output_data,
@@ -820,7 +1751,6 @@
DimsToShape(output_dims), output_data);
}
-// Legacy Dims<4>
inline void ResizeBilinear(const uint8* input_data, const Dims<4>& input_dims,
const int32* output_size_data,
const Dims<4>& output_size_dims, uint8* output_data,
@@ -850,7 +1780,6 @@
output_data, output_dims, /*align_corners=*/false);
}
-// Legacy Dims<4>.
template <typename T>
inline void BatchToSpaceND(const T* input_data, const Dims<4>& input_dims,
const int32* block_shape_data,
diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/multithreaded_conv.h b/tensorflow/contrib/lite/kernels/internal/optimized/multithreaded_conv.h
index b5d001c..4139cf4 100644
--- a/tensorflow/contrib/lite/kernels/internal/optimized/multithreaded_conv.h
+++ b/tensorflow/contrib/lite/kernels/internal/optimized/multithreaded_conv.h
@@ -69,13 +69,13 @@
template <class T>
class EigenTensorConvFunctor {
private:
- Eigen::PaddingType TfLitePadding2EigenPadding(TfLitePadding padding) {
+ Eigen::PaddingType RuntimePadding2EigenPadding(PaddingType padding) {
switch (padding) {
- case kTfLitePaddingValid:
+ case PaddingType::kValid:
return Eigen::PADDING_VALID;
- case kTfLitePaddingSame:
+ case PaddingType::kSame:
return Eigen::PADDING_SAME;
- case kTfLitePaddingUnknown:
+ case PaddingType::kNone:
assert(false); // should never get here.
return Eigen::PADDING_VALID;
}
@@ -89,7 +89,7 @@
int input_width, int input_depth, const T* filter_data,
int filter_height, int filter_width, int filter_count,
int stride_rows, int stride_cols, int pad_width,
- int pad_height, TfLitePadding padding, T* output_data,
+ int pad_height, PaddingType padding, T* output_data,
int output_height, int output_width) {
const bool is_1x1_kernel = (filter_height == 1 && filter_width == 1 &&
stride_rows == 1 && stride_cols == 1);
@@ -127,28 +127,38 @@
input_depth, filter_count);
output.device(device) =
Eigen::SpatialConvolution(input, filter, stride_cols, stride_rows,
- TfLitePadding2EigenPadding(padding));
+ RuntimePadding2EigenPadding(padding));
}
}
};
-inline void Conv(const Eigen::ThreadPoolDevice& device, const float* input_data,
- const Dims<4>& input_dims, const float* filter_data,
- const Dims<4>& filter_dims, const float* bias_data,
- const Dims<4>& bias_dims, int stride_width, int stride_height,
- int pad_width, int pad_height, TfLitePadding padding,
- float output_activation_min, float output_activation_max,
- float* output_data, const Dims<4>& output_dims,
- float* im2col_data, const Dims<4>& im2col_dims) {
- const int batches = MatchingArraySize(input_dims, 3, output_dims, 3);
- const int input_depth = MatchingArraySize(input_dims, 0, filter_dims, 0);
- const int output_depth = MatchingArraySize(filter_dims, 3, output_dims, 0);
- const int input_height = ArraySize(input_dims, 2);
- const int input_width = ArraySize(input_dims, 1);
- const int filter_height = ArraySize(filter_dims, 2);
- const int filter_width = ArraySize(filter_dims, 1);
- const int output_height = ArraySize(output_dims, 2);
- const int output_width = ArraySize(output_dims, 1);
+inline void Conv(const Eigen::ThreadPoolDevice& device,
+ const ConvParams& params, const RuntimeShape& input_shape,
+ const float* input_data, const RuntimeShape& filter_shape,
+ const float* filter_data, const RuntimeShape& bias_shape,
+ const float* bias_data, const RuntimeShape& output_shape,
+ float* output_data, const RuntimeShape& im2col_shape,
+ float* im2col_data) {
+ const int stride_width = params.stride_width;
+ const int stride_height = params.stride_height;
+ const PaddingType padding = params.padding_type;
+ const int pad_width = params.padding_values.width;
+ const int pad_height = params.padding_values.height;
+ const float output_activation_min = params.float_activation_min;
+ const float output_activation_max = params.float_activation_max;
+ TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
+ TFLITE_DCHECK_EQ(filter_shape.DimensionsCount(), 4);
+ TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
+
+ const int batches = MatchingDim(input_shape, 0, output_shape, 0);
+ const int input_depth = MatchingDim(input_shape, 3, filter_shape, 3);
+ const int output_depth = MatchingDim(filter_shape, 0, output_shape, 3);
+ const int input_height = input_shape.Dims(1);
+ const int input_width = input_shape.Dims(2);
+ const int filter_height = filter_shape.Dims(1);
+ const int filter_width = filter_shape.Dims(2);
+ const int output_height = output_shape.Dims(1);
+ const int output_width = output_shape.Dims(2);
EigenTensorConvFunctor<float> conv_functor;
conv_functor(device, input_data, im2col_data, batches, input_height,
input_width, input_depth, filter_data, filter_height,
@@ -157,8 +167,8 @@
output_width);
optimized_ops::AddBiasAndEvalActivationFunction(
- output_activation_min, output_activation_max, DimsToShape(bias_dims),
- bias_data, DimsToShape(output_dims), output_data);
+ output_activation_min, output_activation_max, bias_shape, bias_data,
+ output_shape, output_data);
}
} // namespace multithreaded_ops
diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h b/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h
index 0999738..77f84e0 100644
--- a/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h
+++ b/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h
@@ -52,13 +52,10 @@
using reference_ops::Broadcast4DSlowLessEqualWithScaling;
using reference_ops::Broadcast4DSlowLessWithScaling;
using reference_ops::BroadcastAdd4DSlow;
-using reference_ops::BroadcastGreater;
-using reference_ops::BroadcastGreaterEqual;
-using reference_ops::BroadcastLess;
-using reference_ops::BroadcastLessEqual;
using reference_ops::BroadcastMul4DSlow;
using reference_ops::BroadcastSub4DSlow;
using reference_ops::Concatenation;
+using reference_ops::ConcatenationWithScaling;
using reference_ops::DepthConcatenation;
using reference_ops::Dequantize;
using reference_ops::Div;
@@ -81,7 +78,6 @@
using reference_ops::SpaceToBatchND;
using reference_ops::Split;
using reference_ops::StridedSlice;
-using reference_ops::TensorFlowSplit;
using reference_ops::Transpose;
// TODO(b/80247582) Remove this constant.
@@ -111,12 +107,6 @@
return VectorMap<Scalar>(data, size, 1);
}
-template <typename Scalar, int N>
-VectorMap<Scalar> MapAsVector(Scalar* data, const Dims<N>& dims) {
- const int size = FlatSize(dims);
- return VectorMap<Scalar>(data, size, 1);
-}
-
// Make a local VectorMap typedef allowing to map a float array
// as a Eigen matrix expression. The same explanation as for VectorMap
// above also applies here.
@@ -144,28 +134,6 @@
return MatrixMap<Scalar>(data, rows, cols);
}
-template <typename Scalar, int N>
-MatrixMap<Scalar> MapAsMatrixWithFirstDimAsRows(Scalar* data,
- const Dims<N>& dims) {
- const int rows = dims.sizes[0];
- int cols = 1;
- for (int d = 1; d < N; d++) {
- cols *= dims.sizes[d];
- }
- return MatrixMap<Scalar>(data, rows, cols);
-}
-
-template <typename Scalar, int N>
-MatrixMap<Scalar> MapAsMatrixWithLastDimAsCols(Scalar* data,
- const Dims<N>& dims) {
- const int cols = dims.sizes[N - 1];
- int rows = 1;
- for (int d = 0; d < N - 1; d++) {
- rows *= dims.sizes[d];
- }
- return MatrixMap<Scalar>(data, rows, cols);
-}
-
template <typename Scalar>
using ArrayMap = typename std::conditional<
std::is_const<Scalar>::value,
@@ -173,17 +141,6 @@
Eigen::Dynamic, Eigen::Dynamic>>,
Eigen::Map<Eigen::Array<Scalar, Eigen::Dynamic, Eigen::Dynamic>>>::type;
-template <typename Scalar, int N>
-ArrayMap<Scalar> MapAsArrayWithFirstDimAsRows(Scalar* data,
- const Dims<N>& dims) {
- const int rows = dims.sizes[0];
- int cols = 1;
- for (int d = 1; d < N; d++) {
- cols *= dims.sizes[d];
- }
- return ArrayMap<Scalar>(data, rows, cols);
-}
-
template <typename Scalar>
ArrayMap<Scalar> MapAsArrayWithLastDimAsRows(Scalar* data,
const RuntimeShape& shape) {
@@ -205,20 +162,6 @@
UnalignedConstMatrix;
};
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-// TODO(b/62193649): this function is only needed as long
-// as we have the --variable_batch hack.
-template <typename Scalar, int N>
-MatrixMap<Scalar> MapAsMatrixWithGivenNumberOfRows(Scalar* data,
- const Dims<N>& dims,
- int rows) {
- const int flatsize = FlatSize(dims);
- TFLITE_DCHECK((flatsize % rows) == 0);
- const int cols = flatsize / rows;
- return MatrixMap<Scalar>(data, rows, cols);
-}
-
// TODO(b/62193649): this function is only needed as long
// as we have the --variable_batch hack.
template <typename Scalar>
@@ -270,15 +213,6 @@
SaturatingRoundingMultiplyByPOTParam(a.raw(), exponent));
}
-inline bool AreSameDims(const Dims<4>& dims1, const Dims<4>& dims2) {
- for (int i = 0; i < 4; i++) {
- if (dims1.sizes[i] != dims2.sizes[i]) {
- return false;
- }
- }
- return true;
-}
-
inline void AddBiasAndEvalActivationFunction(float output_activation_min,
float output_activation_max,
const RuntimeShape& bias_shape,
@@ -352,33 +286,6 @@
#endif
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void AddBiasAndEvalActivationFunction(const float* bias_data,
- const Dims<4>& bias_dims,
- float* array_data,
- const Dims<4>& array_dims,
- float output_activation_min,
- float output_activation_max) {
- AddBiasAndEvalActivationFunction(output_activation_min, output_activation_max,
- DimsToShape(bias_dims), bias_data,
- DimsToShape(array_dims), array_data);
-}
-
-// Note: This to be converted to RuntimeShapes along with Conv.
-// legacy, for compatibility with old checked-in code
-template <FusedActivationFunctionType Ac>
-void AddBiasAndEvalActivationFunction(const float* bias_data,
- const Dims<4>& bias_dims,
- float* array_data,
- const Dims<4>& array_dims) {
- float output_activation_min, output_activation_max;
- GetActivationMinMax(Ac, &output_activation_min, &output_activation_max);
- AddBiasAndEvalActivationFunction(bias_data, bias_dims, array_data, array_dims,
- output_activation_min,
- output_activation_max);
-}
-
template <typename Lhs, typename Rhs, typename Result>
void Gemm(const Eigen::MatrixBase<Lhs>& lhs, const Eigen::MatrixBase<Rhs>& rhs,
Eigen::MatrixBase<Result>* result) {
@@ -925,38 +832,6 @@
output_data);
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void FullyConnected(const float* input_data, const Dims<4>& input_dims,
- const float* weights_data,
- const Dims<4>& weights_dims, const float* bias_data,
- const Dims<4>& bias_dims,
- float output_activation_min,
- float output_activation_max, float* output_data,
- const Dims<4>& output_dims) {
- tflite::FullyConnectedParams op_params;
- op_params.float_activation_min = output_activation_min;
- op_params.float_activation_max = output_activation_max;
-
- FullyConnected(op_params, DimsToShape(input_dims), input_data,
- DimsToShape(weights_dims), weights_data,
- DimsToShape(bias_dims), bias_data, DimsToShape(output_dims),
- output_data);
-}
-
-// legacy, for compatibility with old checked-in code
-template <FusedActivationFunctionType Ac>
-void FullyConnected(const float* input_data, const Dims<4>& input_dims,
- const float* weights_data, const Dims<4>& weights_dims,
- const float* bias_data, const Dims<4>& bias_dims,
- float* output_data, const Dims<4>& output_dims) {
- float output_activation_min, output_activation_max;
- GetActivationMinMax(Ac, &output_activation_min, &output_activation_max);
- FullyConnected(input_data, input_dims, weights_data, weights_dims, bias_data,
- bias_dims, output_activation_min, output_activation_max,
- output_data, output_dims);
-}
-
#ifdef USE_NEON
inline void FullyConnectedAsGEMV(
const RuntimeShape& input_shape, const uint8* input_data,
@@ -1203,33 +1078,6 @@
input_offset, output_pipeline);
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void FullyConnected(const uint8* input_data, const Dims<4>& input_dims,
- int32 input_offset, const uint8* filter_data,
- const Dims<4>& filter_dims, int32 filter_offset,
- const int32* bias_data, const Dims<4>& bias_dims,
- int32 output_offset, int32 output_multiplier,
- int output_shift, int32 output_activation_min,
- int32 output_activation_max, uint8* output_data,
- const Dims<4>& output_dims,
- gemmlowp::GemmContext* gemm_context) {
- tflite::FullyConnectedParams op_params;
- op_params.input_offset = input_offset;
- op_params.weights_offset = filter_offset;
- op_params.output_offset = output_offset;
- op_params.output_multiplier = output_multiplier;
- // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
- op_params.output_shift = kReverseShift * output_shift;
- op_params.quantized_activation_min = output_activation_min;
- op_params.quantized_activation_max = output_activation_max;
-
- FullyConnected(op_params, DimsToShape(input_dims), input_data,
- DimsToShape(filter_dims), filter_data, DimsToShape(bias_dims),
- bias_data, DimsToShape(output_dims), output_data,
- gemm_context);
-}
-
inline void FullyConnected(
const FullyConnectedParams& params, const RuntimeShape& input_shape,
const uint8* input_data, const RuntimeShape& filter_shape,
@@ -1317,54 +1165,6 @@
input_offset, output_pipeline);
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void FullyConnected(
- const uint8* input_data, const Dims<4>& input_dims, int32 input_offset,
- const uint8* filter_data, const Dims<4>& filter_dims, int32 filter_offset,
- const int32* bias_data_int32, const Dims<4>& bias_dims, int32 output_offset,
- int32 output_multiplier, int output_shift, int32 output_activation_min,
- int32 output_activation_max, int16* output_data, const Dims<4>& output_dims,
- gemmlowp::GemmContext* gemm_context) {
- tflite::FullyConnectedParams op_params;
- op_params.input_offset = input_offset;
- op_params.weights_offset = filter_offset;
- op_params.output_offset = output_offset;
- op_params.output_multiplier = output_multiplier;
- // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
- op_params.output_shift = kReverseShift * output_shift;
- op_params.quantized_activation_min = output_activation_min;
- op_params.quantized_activation_max = output_activation_max;
-
- FullyConnected(op_params, DimsToShape(input_dims), input_data,
- DimsToShape(filter_dims), filter_data, DimsToShape(bias_dims),
- bias_data_int32, DimsToShape(output_dims), output_data,
- gemm_context);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// legacy, for compatibility with old checked-in code
-template <FusedActivationFunctionType Ac>
-void FullyConnected(const uint8* input_data, const Dims<4>& input_dims,
- int32 input_offset, const uint8* filter_data,
- const Dims<4>& filter_dims, int32 filter_offset,
- const int32* bias_data, const Dims<4>& bias_dims,
- int32 output_offset, int32 output_multiplier,
- int output_shift, int32 output_activation_min,
- int32 output_activation_max, uint8* output_data,
- const Dims<4>& output_dims,
- gemmlowp::GemmContext* gemm_context) {
- static_assert(Ac == FusedActivationFunctionType::kNone ||
- Ac == FusedActivationFunctionType::kRelu ||
- Ac == FusedActivationFunctionType::kRelu6 ||
- Ac == FusedActivationFunctionType::kRelu1,
- "");
- FullyConnected(input_data, input_dims, input_offset, filter_data, filter_dims,
- filter_offset, bias_data, bias_dims, output_offset,
- output_multiplier, output_shift, output_activation_min,
- output_activation_max, output_data, output_dims, gemm_context);
-}
-
// Internal function doing the actual arithmetic work for
// ShuffledFullyConnected.
// May be called either directly by it (single-threaded case) or may be used
@@ -1809,29 +1609,6 @@
gemm_context->workers_pool()->Execute(tasks);
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void ShuffledFullyConnected(
- const uint8* input_data, const Dims<4>& input_dims,
- const uint8* shuffled_weights_data, const Dims<4>& weights_dims,
- const int32* bias_data, const Dims<4>& bias_dims, int32 output_multiplier,
- int output_shift, int32 output_activation_min, int32 output_activation_max,
- int16* output_data, const Dims<4>& output_dims,
- uint8* shuffled_input_workspace_data, gemmlowp::GemmContext* gemm_context) {
- tflite::FullyConnectedParams op_params;
- op_params.output_multiplier = output_multiplier;
- // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
- op_params.output_shift = kReverseShift * output_shift;
- op_params.quantized_activation_min = output_activation_min;
- op_params.quantized_activation_max = output_activation_max;
-
- ShuffledFullyConnected(op_params, DimsToShape(input_dims), input_data,
- DimsToShape(weights_dims), shuffled_weights_data,
- DimsToShape(bias_dims), bias_data,
- DimsToShape(output_dims), output_data,
- shuffled_input_workspace_data, gemm_context);
-}
-
template <typename T>
inline void ExtractPatchIntoBufferColumn(const RuntimeShape& input_shape, int w,
int h, int b, int kheight, int kwidth,
@@ -1922,20 +1699,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-template <typename T>
-inline void ExtractPatchIntoBufferColumn(
- const Dims<4>& input_dims, int w, int h, int b, int kheight, int kwidth,
- int stride_width, int stride_height, int pad_width, int pad_height,
- int in_width, int in_height, int in_depth, int single_buffer_length,
- int buffer_id, const T* in_data, T* conv_buffer_data, uint8 zero_byte) {
- ExtractPatchIntoBufferColumn(
- DimsToShape(input_dims), w, h, b, kheight, kwidth, stride_width,
- stride_height, pad_width, pad_height, in_width, in_height, in_depth,
- single_buffer_length, buffer_id, in_data, conv_buffer_data, zero_byte);
-}
-
template <typename T>
void DilatedIm2col(const ConvParams& params, uint8 zero_byte,
const RuntimeShape& input_shape, const T* input_data,
@@ -2019,30 +1782,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-template <typename T>
-void DilatedIm2col(const T* input_data, const Dims<4>& input_dims,
- const Dims<4>& filter_dims, int stride_width,
- int stride_height, int dilation_width_factor,
- int dilation_height_factor, int pad_width, int pad_height,
- const Dims<4>& output_dims, uint8 zero_byte,
- T* im2col_data) {
- tflite::ConvParams op_params;
- // Padding type is ignored, but still set.
- op_params.padding_type = PaddingType::kSame;
- op_params.padding_values.width = pad_width;
- op_params.padding_values.height = pad_height;
- op_params.stride_width = stride_width;
- op_params.stride_height = stride_height;
- op_params.dilation_width_factor = dilation_width_factor;
- op_params.dilation_height_factor = dilation_height_factor;
-
- DilatedIm2col(op_params, zero_byte, DimsToShape(input_dims), input_data,
- DimsToShape(filter_dims), DimsToShape(output_dims),
- im2col_data);
-}
-
template <typename T>
void Im2col(const ConvParams& params, int kheight, int kwidth, uint8 zero_byte,
const RuntimeShape& input_shape, const T* input_data,
@@ -2078,36 +1817,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-template <typename T>
-void Im2col(const T* input_data, const Dims<4>& input_dims, int stride_width,
- int stride_height, int pad_width, int pad_height, int kheight,
- int kwidth, uint8 zero_byte, T* output_data,
- const Dims<4>& output_dims) {
- tflite::ConvParams op_params;
- // Padding type is ignored, but still set.
- op_params.padding_type = PaddingType::kSame;
- op_params.padding_values.width = pad_width;
- op_params.padding_values.height = pad_height;
- op_params.stride_width = stride_width;
- op_params.stride_height = stride_height;
- op_params.dilation_width_factor = 1;
- op_params.dilation_height_factor = 1;
-
- Im2col(op_params, kheight, kwidth, zero_byte, DimsToShape(input_dims),
- input_data, DimsToShape(output_dims), output_data);
-}
-
-// legacy, for compatibility with old checked-in code
-template <typename T>
-void Im2col(const T* input_data, const Dims<4>& input_dims, int stride,
- int pad_width, int pad_height, int kheight, int kwidth,
- uint8 zero_byte, T* output_data, const Dims<4>& output_dims) {
- Im2col(input_data, input_dims, stride, stride, pad_width, pad_height, kheight,
- kwidth, zero_byte, output_data, output_dims);
-}
-
inline void Conv(const ConvParams& params, const RuntimeShape& input_shape,
const float* input_data, const RuntimeShape& filter_shape,
const float* filter_data, const RuntimeShape& bias_shape,
@@ -2171,33 +1880,6 @@
output_data);
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void Conv(const float* input_data, const Dims<4>& input_dims,
- const float* filter_data, const Dims<4>& filter_dims,
- const float* bias_data, const Dims<4>& bias_dims,
- int stride_width, int stride_height, int dilation_width_factor,
- int dilation_height_factor, int pad_width, int pad_height,
- float output_activation_min, float output_activation_max,
- float* output_data, const Dims<4>& output_dims,
- float* im2col_data, const Dims<4>& im2col_dims) {
- tflite::ConvParams op_params;
- // Padding type is ignored, but still set.
- op_params.padding_type = PaddingType::kSame;
- op_params.padding_values.width = pad_width;
- op_params.padding_values.height = pad_height;
- op_params.stride_width = stride_width;
- op_params.stride_height = stride_height;
- op_params.dilation_width_factor = dilation_width_factor;
- op_params.dilation_height_factor = dilation_height_factor;
- op_params.float_activation_min = output_activation_min;
- op_params.float_activation_max = output_activation_max;
-
- Conv(op_params, DimsToShape(input_dims), input_data, DimsToShape(filter_dims),
- filter_data, DimsToShape(bias_dims), bias_data, DimsToShape(output_dims),
- output_data, DimsToShape(im2col_dims), im2col_data);
-}
-
inline void HybridConv(const ConvParams& params, float* scaling_factors_ptr,
const RuntimeShape& input_shape,
const int8_t* input_data,
@@ -2278,82 +1960,6 @@
output_data);
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void HybridConv(const int8_t* input_data, const Dims<4>& input_dims,
- const int8_t* filter_data, const Dims<4>& filter_dims,
- const float* bias_data, const Dims<4>& bias_dims,
- int stride_width, int stride_height, int pad_width,
- int pad_height, float* scaling_factors_ptr,
- float output_activation_min, float output_activation_max,
- float* output_data, const Dims<4>& output_dims,
- int8_t* im2col_data, const Dims<4>& im2col_dims) {
- tflite::ConvParams op_params;
- // Padding type is ignored, but still set.
- op_params.padding_type = PaddingType::kSame;
- op_params.padding_values.width = pad_width;
- op_params.padding_values.height = pad_height;
- op_params.stride_width = stride_width;
- op_params.stride_height = stride_height;
- op_params.float_activation_min = output_activation_min;
- op_params.float_activation_max = output_activation_max;
-
- HybridConv(op_params, scaling_factors_ptr, DimsToShape(input_dims),
- input_data, DimsToShape(filter_dims), filter_data,
- DimsToShape(bias_dims), bias_data, DimsToShape(output_dims),
- output_data, DimsToShape(im2col_dims), im2col_data);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-template <FusedActivationFunctionType Ac>
-void Conv(const float* input_data, const Dims<4>& input_dims,
- const float* filter_data, const Dims<4>& filter_dims,
- const float* bias_data, const Dims<4>& bias_dims, int stride_width,
- int stride_height, int dilation_width_factor,
- int dilation_height_factor, int pad_width, int pad_height,
- float* output_data, const Dims<4>& output_dims, float* im2col_data,
- const Dims<4>& im2col_dims) {
- float output_activation_min, output_activation_max;
- GetActivationMinMax(Ac, &output_activation_min, &output_activation_max);
- Conv(input_data, input_dims, filter_data, filter_dims, bias_data, bias_dims,
- stride_width, stride_height, dilation_width_factor,
- dilation_height_factor, pad_width, pad_height, output_activation_min,
- output_activation_max, output_data, output_dims, im2col_data,
- im2col_dims);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// legacy, for compatibility with old checked-in code
-template <FusedActivationFunctionType Ac>
-void Conv(const float* input_data, const Dims<4>& input_dims,
- const float* filter_data, const Dims<4>& filter_dims,
- const float* bias_data, const Dims<4>& bias_dims, int stride_width,
- int stride_height, int pad_width, int pad_height, float* output_data,
- const Dims<4>& output_dims, float* im2col_data,
- const Dims<4>& im2col_dims) {
- float output_activation_min, output_activation_max;
- GetActivationMinMax(Ac, &output_activation_min, &output_activation_max);
- Conv(input_data, input_dims, filter_data, filter_dims, bias_data, bias_dims,
- stride_width, stride_height, 1, 1, pad_width, pad_height,
- output_activation_min, output_activation_max, output_data, output_dims,
- im2col_data, im2col_dims);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// legacy, for compatibility with old checked-in code
-template <FusedActivationFunctionType Ac>
-void Conv(const float* input_data, const Dims<4>& input_dims,
- const float* filter_data, const Dims<4>& filter_dims,
- const float* bias_data, const Dims<4>& bias_dims, int stride,
- int pad_width, int pad_height, float* output_data,
- const Dims<4>& output_dims, float* im2col_data,
- const Dims<4>& im2col_dims) {
- Conv<Ac>(input_data, input_dims, filter_data, filter_dims, bias_data,
- bias_dims, stride, stride, 1, 1, pad_width, pad_height, output_data,
- output_dims, im2col_data, im2col_dims);
-}
-
inline void Conv(const ConvParams& params, const RuntimeShape& input_shape,
const uint8* input_data, const RuntimeShape& filter_shape,
const uint8* filter_data, const RuntimeShape& bias_shape,
@@ -2445,192 +2051,6 @@
input_offset, output_pipeline);
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void Conv(const uint8* input_data, const Dims<4>& input_dims,
- int32 input_offset, const uint8* filter_data,
- const Dims<4>& filter_dims, int32 filter_offset,
- const int32* bias_data, const Dims<4>& bias_dims,
- int stride_width, int stride_height, int dilation_width_factor,
- int dilation_height_factor, int pad_width, int pad_height,
- int32 output_offset, int32 output_multiplier, int output_shift,
- int32 output_activation_min, int32 output_activation_max,
- uint8* output_data, const Dims<4>& output_dims,
- uint8* im2col_data, const Dims<4>& im2col_dims,
- gemmlowp::GemmContext* gemm_context) {
- tflite::ConvParams op_params;
- // Padding type is ignored, but still set.
- op_params.padding_type = PaddingType::kSame;
- op_params.padding_values.width = pad_width;
- op_params.padding_values.height = pad_height;
- op_params.stride_width = stride_width;
- op_params.stride_height = stride_height;
- op_params.dilation_width_factor = dilation_width_factor;
- op_params.dilation_height_factor = dilation_height_factor;
- op_params.input_offset = input_offset;
- op_params.weights_offset = filter_offset;
- op_params.output_offset = output_offset;
- op_params.output_multiplier = output_multiplier;
- // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
- op_params.output_shift = kReverseShift * output_shift;
- op_params.quantized_activation_min = output_activation_min;
- op_params.quantized_activation_max = output_activation_max;
-
- Conv(op_params, DimsToShape(input_dims), input_data, DimsToShape(filter_dims),
- filter_data, DimsToShape(bias_dims), bias_data, DimsToShape(output_dims),
- output_data, DimsToShape(im2col_dims), im2col_data, gemm_context);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void Conv(const uint8* input_data, const Dims<4>& input_dims,
- int32 input_offset, const uint8* filter_data,
- const Dims<4>& filter_dims, int32 filter_offset,
- const int32* bias_data, const Dims<4>& bias_dims,
- int stride_width, int stride_height, int pad_width,
- int pad_height, int32 output_offset, int32 output_multiplier,
- int output_shift, int32 output_activation_min,
- int32 output_activation_max, uint8* output_data,
- const Dims<4>& output_dims, uint8* im2col_data,
- const Dims<4>& im2col_dims,
- gemmlowp::GemmContext* gemm_context) {
- Conv(input_data, input_dims, input_offset, filter_data, filter_dims,
- filter_offset, bias_data, bias_dims, stride_width, stride_height, 1, 1,
- pad_width, pad_height, output_offset, output_multiplier, output_shift,
- output_activation_min, output_activation_max, output_data, output_dims,
- im2col_data, im2col_dims, gemm_context);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// legacy, for compatibility with old checked-in code
-template <FusedActivationFunctionType Ac>
-inline void Conv(const uint8* input_data, const Dims<4>& input_dims,
- int32 input_offset, const uint8* filter_data,
- const Dims<4>& filter_dims, int32 filter_offset,
- const int32* bias_data, const Dims<4>& bias_dims,
- int stride_width, int stride_height, int pad_width,
- int pad_height, int32 output_offset, int32 output_multiplier,
- int output_shift, int32 output_activation_min,
- int32 output_activation_max, uint8* output_data,
- const Dims<4>& output_dims, uint8* im2col_data,
- const Dims<4>& im2col_dims,
- gemmlowp::GemmContext* gemm_context) {
- static_assert(Ac == FusedActivationFunctionType::kNone ||
- Ac == FusedActivationFunctionType::kRelu ||
- Ac == FusedActivationFunctionType::kRelu6 ||
- Ac == FusedActivationFunctionType::kRelu1,
- "");
- if (Ac == FusedActivationFunctionType::kNone) {
- TFLITE_DCHECK_EQ(output_activation_min, 0);
- TFLITE_DCHECK_EQ(output_activation_max, 255);
- }
- Conv(input_data, input_dims, input_offset, filter_data, filter_dims,
- filter_offset, bias_data, bias_dims, stride_width, stride_height,
- pad_width, pad_height, output_offset, output_multiplier, output_shift,
- output_activation_min, output_activation_max, output_data, output_dims,
- im2col_data, im2col_dims, gemm_context);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// legacy, for compatibility with old checked-in code
-template <FusedActivationFunctionType Ac>
-void Conv(const uint8* input_data, const Dims<4>& input_dims,
- int32 input_offset, const uint8* filter_data,
- const Dims<4>& filter_dims, int32 filter_offset,
- const int32* bias_data, const Dims<4>& bias_dims, int stride,
- int pad_width, int pad_height, int32 output_offset,
- int32 output_multiplier, int output_shift,
- int32 output_activation_min, int32 output_activation_max,
- uint8* output_data, const Dims<4>& output_dims, uint8* im2col_data,
- const Dims<4>& im2col_dims, gemmlowp::GemmContext* gemm_context) {
- static_assert(Ac == FusedActivationFunctionType::kNone ||
- Ac == FusedActivationFunctionType::kRelu ||
- Ac == FusedActivationFunctionType::kRelu6 ||
- Ac == FusedActivationFunctionType::kRelu1,
- "");
- Conv(input_data, input_dims, input_offset, filter_data, filter_dims,
- filter_offset, bias_data, bias_dims, stride, stride, pad_width,
- pad_height, output_offset, output_multiplier, output_shift,
- output_activation_min, output_activation_max, output_data, output_dims,
- im2col_data, im2col_dims, gemm_context);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// legacy, for compatibility with old checked-in code
-template <FusedActivationFunctionType Ac, typename T>
-void Im2col(const T* input_data, const Dims<4>& input_dims, int stride,
- int pad_width, int pad_height, int kheight, int kwidth,
- uint8 zero_byte, T* output_data, const Dims<4>& output_dims) {
- Im2col(input_data, input_dims, stride, stride, pad_width, pad_height, kheight,
- kwidth, zero_byte, output_data, output_dims);
-}
-
-// legacy, for compatibility with old checked-in code
-template <FusedActivationFunctionType Ac>
-void ConvAsGemm(const float* input_data, const Dims<4>& input_dims,
- const float* filter_data, const Dims<4>& filter_dims,
- const float* bias_data, const Dims<4>& bias_dims,
- float* output_data, const Dims<4>& output_dims) {
- gemmlowp::ScopedProfilingLabel label("ConvAsGemm");
-
- const auto input_matrix_map =
- MapAsMatrixWithFirstDimAsRows(input_data, input_dims);
- const auto filter_matrix_map =
- MapAsMatrixWithLastDimAsCols(filter_data, filter_dims);
- auto output_matrix_map =
- MapAsMatrixWithFirstDimAsRows(output_data, output_dims);
-
- Gemm(filter_matrix_map.transpose(), input_matrix_map, &output_matrix_map);
-
- AddBiasAndEvalActivationFunction<Ac>(bias_data, bias_dims, output_data,
- output_dims);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// legacy, for compatibility with old checked-in code
-template <FusedActivationFunctionType Ac>
-void ConvAsGemm(const uint8* input_data, const Dims<4>& input_dims,
- int32 input_offset, const uint8* filter_data,
- const Dims<4>& filter_dims, int32 filter_offset,
- const int32* bias_data, const Dims<4>& bias_dims,
- int32 output_offset, int32 output_multiplier, int output_shift,
- int32 output_activation_min, int32 output_activation_max,
- uint8* output_data, const Dims<4>& output_dims,
- gemmlowp::GemmContext* gemm_context) {
- gemmlowp::ScopedProfilingLabel label("ConvAsGemm/8bit");
- static_assert(Ac == FusedActivationFunctionType::kNone ||
- Ac == FusedActivationFunctionType::kRelu ||
- Ac == FusedActivationFunctionType::kRelu6 ||
- Ac == FusedActivationFunctionType::kRelu1,
- "");
- const int input_rows = input_dims.sizes[0];
- const int input_cols = FlatSizeSkipDim(input_dims, 0);
- const int filter_rows = filter_dims.sizes[3];
- const int filter_cols = FlatSizeSkipDim(filter_dims, 3);
- const int output_rows = output_dims.sizes[0];
- const int output_cols = FlatSizeSkipDim(output_dims, 0);
- TFLITE_DCHECK_EQ(output_rows, filter_rows);
- TFLITE_DCHECK_EQ(output_cols, input_cols);
- TFLITE_DCHECK_EQ(filter_cols, input_rows);
- TFLITE_DCHECK_EQ(bias_dims.sizes[0], output_rows);
- TFLITE_DCHECK_EQ(bias_dims.sizes[1], 1);
- TFLITE_DCHECK_EQ(bias_dims.sizes[2], 1);
- TFLITE_DCHECK_EQ(bias_dims.sizes[3], 1);
- gemmlowp::MatrixMap<const uint8, gemmlowp::MapOrder::RowMajor> filter_matrix(
- filter_data, output_rows, filter_cols, filter_cols);
- gemmlowp::MatrixMap<const uint8, gemmlowp::MapOrder::ColMajor> input_matrix(
- input_data, filter_cols, output_cols, filter_cols);
- gemmlowp::MatrixMap<uint8, gemmlowp::MapOrder::ColMajor> output_matrix(
- output_data, output_rows, output_cols, output_rows);
- const auto& output_pipeline = GemmlowpOutputPipeline::MakeExp(
- bias_data, output_rows, output_offset, output_multiplier, -output_shift,
- output_activation_min, output_activation_max);
- gemmlowp::GemmWithOutputPipeline<uint8, uint8,
- gemmlowp::L8R8WithLhsNonzeroBitDepthParams>(
- gemm_context, filter_matrix, input_matrix, &output_matrix, filter_offset,
- input_offset, output_pipeline);
-}
-
template <typename T>
inline void DepthToSpace(const tflite::DepthToSpaceParams& op_params,
const RuntimeShape& unextended_input_shape,
@@ -3547,21 +2967,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy Dims<4>.
-template <typename T>
-void BroadcastDiv(const T* input1_data, const Dims<4>& input1_dims,
- const T* input2_data, const Dims<4>& input2_dims,
- T output_activation_min, T output_activation_max,
- T* output_data, const Dims<4>& output_dims) {
- tflite::ArithmeticParams op_params;
- SetActivationParams(output_activation_min, output_activation_max, &op_params);
-
- BroadcastDiv4DSlow(op_params, DimsToShape(input1_dims), input1_data,
- DimsToShape(input2_dims), input2_data,
- DimsToShape(output_dims), output_data);
-}
-
// TODO(aselle): This is not actually optimized yet.
inline void SubNonBroadcast(const ArithmeticParams& params,
const RuntimeShape& input1_shape,
@@ -3755,31 +3160,6 @@
output_state_map.tanh();
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void LstmCell(const float* input_data, const Dims<4>& input_dims,
- const float* prev_activ_data,
- const Dims<4>& prev_activ_dims, const float* weights_data,
- const Dims<4>& weights_dims, const float* bias_data,
- const Dims<4>& bias_dims, const float* prev_state_data,
- const Dims<4>& prev_state_dims, float* output_state_data,
- const Dims<4>& output_state_dims, float* output_activ_data,
- const Dims<4>& output_activ_dims, float* concat_temp_data,
- const Dims<4>& concat_temp_dims, float* activ_temp_data,
- const Dims<4>& activ_temp_dims) {
- tflite::LstmCellParams op_params;
- // Float LSTM cell does not need parameters to be set: leave untouched.
-
- LstmCell(op_params, DimsToShape(input_dims), input_data,
- DimsToShape(prev_activ_dims), prev_activ_data,
- DimsToShape(weights_dims), weights_data, DimsToShape(bias_dims),
- bias_data, DimsToShape(prev_state_dims), prev_state_data,
- DimsToShape(output_state_dims), output_state_data,
- DimsToShape(output_activ_dims), output_activ_data,
- DimsToShape(concat_temp_dims), concat_temp_data,
- DimsToShape(activ_temp_dims), activ_temp_data);
-}
-
// Quantized LSTM cell. Currently just a copy of the reference impl in
// reference_ops.h. See the big function comment there, not replicating it
// here.
@@ -4070,37 +3450,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-template <int StateIntegerBits>
-void LstmCell(const uint8* input_data_uint8, const Dims<4>& input_dims,
- const uint8* prev_activ_data_uint8,
- const Dims<4>& prev_activ_dims, const uint8* weights_data_uint8,
- const Dims<4>& weights_dims, const int32* bias_data_int32,
- const Dims<4>& bias_dims, const int16* prev_state_data_int16,
- const Dims<4>& prev_state_dims, int16* output_state_data_int16,
- const Dims<4>& output_state_dims, uint8* output_activ_data_uint8,
- const Dims<4>& output_activ_dims, uint8* concat_temp_data_uint8,
- const Dims<4>& concat_temp_dims, int16* activ_temp_data_int16,
- const Dims<4>& activ_temp_dims, int32 weights_zero_point,
- int32 accum_multiplier, int accum_shift,
- gemmlowp::GemmContext* gemm_context) {
- tflite::LstmCellParams op_params;
- op_params.weights_zero_point = weights_zero_point;
- op_params.accum_multiplier = accum_multiplier;
- op_params.accum_shift = accum_shift;
-
- LstmCell<StateIntegerBits>(
- op_params, DimsToShape(input_dims), input_data_uint8,
- DimsToShape(prev_activ_dims), prev_activ_data_uint8,
- DimsToShape(weights_dims), weights_data_uint8, DimsToShape(bias_dims),
- bias_data_int32, DimsToShape(prev_state_dims), prev_state_data_int16,
- DimsToShape(output_state_dims), output_state_data_int16,
- DimsToShape(output_activ_dims), output_activ_data_uint8,
- DimsToShape(concat_temp_dims), concat_temp_data_uint8,
- DimsToShape(activ_temp_dims), activ_temp_data_int16, gemm_context);
-}
-
inline int NodeOffset(int b, int h, int w, int height, int width) {
return (b * height + h) * width + w;
}
@@ -4560,16 +3909,6 @@
out_mat.array().rowwise() *= scale;
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void Softmax(const float* input_data, const RuntimeShape& input_shape,
- float beta, float* output_data,
- const RuntimeShape& output_shape) {
- SoftmaxParams params;
- params.beta = beta;
- Softmax(params, input_shape, input_data, output_shape, output_data);
-}
-
inline void Softmax(const SoftmaxParams& params,
const RuntimeShape& input_shape, const uint8* input_data,
const RuntimeShape& output_shape, uint8* output_data) {
@@ -4781,19 +4120,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void Softmax(const uint8* input_data, const RuntimeShape& input_shape,
- int32 input_beta_multiplier, int32 input_beta_left_shift,
- int diff_min, uint8* output_data,
- const RuntimeShape& output_shape) {
- SoftmaxParams params;
- params.input_multiplier = input_beta_multiplier;
- params.input_left_shift = input_beta_left_shift;
- params.diff_min = diff_min;
- Softmax(params, input_shape, input_data, output_shape, output_data);
-}
-
// TODO(myenik): This is the same as the reference implementation, not actually
// optimized yet.
inline void LogSoftmax(const SoftmaxParams& params,
@@ -4831,15 +4157,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy
-inline void LogSoftmax(const float* input_data, const RuntimeShape& input_shape,
- float* output_data, const RuntimeShape& output_shape) {
- SoftmaxParams params;
- // No params currently used for float LogSoftmax.
- LogSoftmax(params, input_shape, input_data, output_shape, output_data);
-}
-
template <int OutputIntegerBits, int InputIntegerBits>
inline gemmlowp::FixedPoint<int32, OutputIntegerBits>
log_x_for_x_greater_than_or_equal_to_1_impl(
@@ -5044,22 +4361,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void LogSoftmax(const uint8* input_data, const RuntimeShape& input_shape,
- int32 input_multiplier, int32 input_left_shift,
- int32 reverse_scaling_divisor,
- int32 reverse_scaling_right_shift, int diff_min,
- uint8* output_data, const RuntimeShape& output_shape) {
- SoftmaxParams params;
- params.input_multiplier = input_multiplier;
- params.input_left_shift = input_left_shift;
- params.reverse_scaling_divisor = reverse_scaling_divisor;
- params.reverse_scaling_right_shift = reverse_scaling_right_shift;
- params.diff_min = diff_min;
- LogSoftmax(params, input_shape, input_data, output_shape, output_data);
-}
-
inline void Logistic(const RuntimeShape& input_shape, const float* input_data,
const RuntimeShape& output_shape, float* output_data) {
gemmlowp::ScopedProfilingLabel label("Logistic");
@@ -5218,20 +4519,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void Logistic(const uint8* input_data, const RuntimeShape& input_shape,
- int32 input_zero_point, int32 input_range_radius,
- int32 input_multiplier, int input_left_shift,
- uint8* output_data, const RuntimeShape& output_shape) {
- LogisticParams params;
- params.input_zero_point = input_zero_point;
- params.input_range_radius = input_range_radius;
- params.input_multiplier = input_multiplier;
- params.input_left_shift = input_left_shift;
- Logistic(params, input_shape, input_data, output_shape, output_data);
-}
-
inline void Logistic(const LogisticParams& params,
const RuntimeShape& input_shape, const int16* input_data,
const RuntimeShape& output_shape, int16* output_data) {
@@ -5293,24 +4580,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy version.
-inline void Logistic(const RuntimeShape& input_shape, const int16* input_data,
- const RuntimeShape& output_shape, int16* output_data) {
- LogisticParams params;
- // No params currently needed by int16 Logistic.
- Logistic(params, input_shape, input_data, output_shape, output_data);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy version.
-inline void Logistic(const int16* input_data, const RuntimeShape& input_shape,
- int16* output_data, const RuntimeShape& output_shape) {
- LogisticParams params;
- // No params currently needed by int16 Logistic.
- Logistic(params, input_shape, input_data, output_shape, output_data);
-}
-
inline void Tanh(const RuntimeShape& input_shape, const float* input_data,
const RuntimeShape& output_shape, float* output_data) {
gemmlowp::ScopedProfilingLabel label("Tanh");
@@ -5478,20 +4747,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void Tanh(const uint8* input_data, const RuntimeShape& input_shape,
- int32 input_zero_point, int32 input_range_radius,
- int32 input_multiplier, int input_left_shift,
- uint8* output_data, const RuntimeShape& output_shape) {
- TanhParams params;
- params.input_zero_point = input_zero_point;
- params.input_range_radius = input_range_radius;
- params.input_multiplier = input_multiplier;
- params.input_left_shift = input_left_shift;
- Tanh(params, input_shape, input_data, output_shape, output_data);
-}
-
inline void Tanh(const TanhParams& params, const RuntimeShape& input_shape,
const int16* input_data, const RuntimeShape& output_shape,
int16* output_data) {
@@ -5593,16 +4848,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void Tanh(const int16* input_data, const RuntimeShape& input_shape,
- int input_left_shift, int16* output_data,
- const RuntimeShape& output_shape) {
- TanhParams params;
- params.input_left_shift = input_left_shift;
- Tanh(params, input_shape, input_data, output_shape, output_data);
-}
-
template <typename SrcT, typename DstT>
inline void Cast(const RuntimeShape& input_shape, const SrcT* input_data,
const RuntimeShape& output_shape, DstT* output_data) {
@@ -6485,27 +5730,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-template <typename T>
-void TransposeIm2col(const T* input_data, const Dims<4>& input_dims,
- const Dims<4>& filter_dims, int stride_width,
- int stride_height, int pad_width, int pad_height,
- const Dims<4>& output_dims, uint8 zero_byte,
- T* im2col_data) {
- tflite::ConvParams op_params;
- // Padding type is ignored, but still set.
- op_params.padding_type = PaddingType::kSame;
- op_params.padding_values.width = pad_width;
- op_params.padding_values.height = pad_height;
- op_params.stride_width = stride_width;
- op_params.stride_height = stride_height;
-
- TransposeIm2col(op_params, zero_byte, DimsToShape(input_dims), input_data,
- DimsToShape(filter_dims), DimsToShape(output_dims),
- im2col_data);
-}
-
inline void TransposeConv(
const ConvParams& params, const RuntimeShape& input_shape,
const float* input_data, const RuntimeShape& filter_shape,
@@ -6529,27 +5753,6 @@
Gemm(filter_matrix_map.transpose(), im2col_matrix_map, &output_matrix_map);
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void TransposeConv(const float* input_data, const Dims<4>& input_dims,
- const float* filter_data, const Dims<4>& filter_dims,
- int stride_width, int stride_height, int pad_width,
- int pad_height, float* output_data,
- const Dims<4>& output_dims, float* im2col_data,
- const Dims<4>& im2col_dims) {
- tflite::ConvParams op_params;
- // Padding type is ignored, but still set.
- op_params.padding_type = PaddingType::kSame;
- op_params.padding_values.width = pad_width;
- op_params.padding_values.height = pad_height;
- op_params.stride_width = stride_width;
- op_params.stride_height = stride_height;
-
- TransposeConv(op_params, DimsToShape(input_dims), input_data,
- DimsToShape(filter_dims), filter_data, DimsToShape(output_dims),
- output_data, DimsToShape(im2col_dims), im2col_data);
-}
-
} // namespace optimized_ops
} // namespace tflite
diff --git a/tensorflow/contrib/lite/kernels/internal/reference/depthwiseconv_float.h b/tensorflow/contrib/lite/kernels/internal/reference/depthwiseconv_float.h
index a842852..1122427 100644
--- a/tensorflow/contrib/lite/kernels/internal/reference/depthwiseconv_float.h
+++ b/tensorflow/contrib/lite/kernels/internal/reference/depthwiseconv_float.h
@@ -94,81 +94,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void DepthwiseConv(const float* input_data, const Dims<4>& input_dims,
- const float* filter_data, const Dims<4>& filter_dims,
- const float* bias_data, const Dims<4>& bias_dims,
- int stride_width, int stride_height,
- int dilation_width_factor, int dilation_height_factor,
- int pad_width, int pad_height, int depth_multiplier,
- float output_activation_min,
- float output_activation_max, float* output_data,
- const Dims<4>& output_dims) {
- tflite::DepthwiseParams op_params;
- // Padding type is ignored, but still set.
- op_params.padding_type = PaddingType::kSame;
- op_params.padding_values.width = pad_width;
- op_params.padding_values.height = pad_height;
- op_params.stride_width = stride_width;
- op_params.stride_height = stride_height;
- op_params.dilation_width_factor = dilation_width_factor;
- op_params.dilation_height_factor = dilation_height_factor;
- op_params.depth_multiplier = depth_multiplier;
- op_params.float_activation_min = output_activation_min;
- op_params.float_activation_max = output_activation_max;
-
- DepthwiseConv(op_params, DimsToShape(input_dims), input_data,
- DimsToShape(filter_dims), filter_data, DimsToShape(bias_dims),
- bias_data, DimsToShape(output_dims), output_data);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void DepthwiseConv(const float* input_data, const Dims<4>& input_dims,
- const float* filter_data, const Dims<4>& filter_dims,
- const float* bias_data, const Dims<4>& bias_dims,
- int stride_width, int stride_height, int pad_width,
- int pad_height, int depth_multiplier,
- float output_activation_min,
- float output_activation_max, float* output_data,
- const Dims<4>& output_dims) {
- DepthwiseConv(input_data, input_dims, filter_data, filter_dims, bias_data,
- bias_dims, stride_width, stride_height, 1, 1, pad_width,
- pad_height, depth_multiplier, output_activation_min,
- output_activation_max, output_data, output_dims);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy, for compatibility with old checked-in code.
-template <FusedActivationFunctionType Ac>
-void DepthwiseConv(const float* input_data, const Dims<4>& input_dims,
- const float* filter_data, const Dims<4>& filter_dims,
- const float* bias_data, const Dims<4>& bias_dims,
- int stride_width, int stride_height, int pad_width,
- int pad_height, int depth_multiplier, float* output_data,
- const Dims<4>& output_dims) {
- float output_activation_min, output_activation_max;
- GetActivationMinMax(Ac, &output_activation_min, &output_activation_max);
- DepthwiseConv(input_data, input_dims, filter_data, filter_dims, bias_data,
- bias_dims, stride_width, stride_height, pad_width, pad_height,
- depth_multiplier, output_activation_min, output_activation_max,
- output_data, output_dims);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy, for compatibility with old checked-in code.
-template <FusedActivationFunctionType Ac>
-void DepthwiseConv(const float* input_data, const Dims<4>& input_dims,
- const float* filter_data, const Dims<4>& filter_dims,
- const float* bias_data, const Dims<4>& bias_dims, int stride,
- int pad_width, int pad_height, int depth_multiplier,
- float* output_data, const Dims<4>& output_dims) {
- DepthwiseConv<Ac>(input_data, input_dims, filter_data, filter_dims, bias_data,
- bias_dims, stride, stride, pad_width, pad_height,
- depth_multiplier, output_data, output_dims);
-}
-
} // end namespace reference_ops
} // end namespace tflite
diff --git a/tensorflow/contrib/lite/kernels/internal/reference/depthwiseconv_uint8.h b/tensorflow/contrib/lite/kernels/internal/reference/depthwiseconv_uint8.h
index e8fc566..eab28e6 100644
--- a/tensorflow/contrib/lite/kernels/internal/reference/depthwiseconv_uint8.h
+++ b/tensorflow/contrib/lite/kernels/internal/reference/depthwiseconv_uint8.h
@@ -25,9 +25,6 @@
namespace tflite {
namespace reference_ops {
-// TODO(b/80418076): Move to legacy ops file, along with invocations.
-static constexpr int kDepthwiseReverseShift = -1;
-
inline void DepthwiseConv(
const DepthwiseParams& params, const RuntimeShape& input_shape,
const uint8* input_data, const RuntimeShape& filter_shape,
@@ -109,106 +106,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void DepthwiseConv(const uint8* input_data, const Dims<4>& input_dims,
- int32 input_offset, const uint8* filter_data,
- const Dims<4>& filter_dims, int32 filter_offset,
- const int32* bias_data, const Dims<4>& bias_dims,
- int stride_width, int stride_height,
- int dilation_width_factor, int dilation_height_factor,
- int pad_width, int pad_height, int depth_multiplier,
- int32 output_offset, int32 output_multiplier,
- int output_shift, int32 output_activation_min,
- int32 output_activation_max, uint8* output_data,
- const Dims<4>& output_dims) {
- tflite::DepthwiseParams op_params;
- // Padding type is ignored, but still set.
- op_params.padding_type = PaddingType::kSame;
- op_params.padding_values.width = pad_width;
- op_params.padding_values.height = pad_height;
- op_params.stride_width = stride_width;
- op_params.stride_height = stride_height;
- op_params.dilation_width_factor = dilation_width_factor;
- op_params.dilation_height_factor = dilation_height_factor;
- op_params.depth_multiplier = depth_multiplier;
- op_params.quantized_activation_min = output_activation_min;
- op_params.quantized_activation_max = output_activation_max;
- op_params.input_offset = input_offset;
- op_params.weights_offset = filter_offset;
- op_params.output_offset = output_offset;
- op_params.output_multiplier = output_multiplier;
- // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
- op_params.output_shift = kDepthwiseReverseShift * output_shift;
-
- DepthwiseConv(op_params, DimsToShape(input_dims), input_data,
- DimsToShape(filter_dims), filter_data, DimsToShape(bias_dims),
- bias_data, DimsToShape(output_dims), output_data);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-inline void DepthwiseConv(const uint8* input_data, const Dims<4>& input_dims,
- int32 input_offset, const uint8* filter_data,
- const Dims<4>& filter_dims, int32 filter_offset,
- const int32* bias_data, const Dims<4>& bias_dims,
- int stride_width, int stride_height, int pad_width,
- int pad_height, int depth_multiplier,
- int32 output_offset, int32 output_multiplier,
- int output_shift, int32 output_activation_min,
- int32 output_activation_max, uint8* output_data,
- const Dims<4>& output_dims) {
- DepthwiseConv(input_data, input_dims, input_offset, filter_data, filter_dims,
- filter_offset, bias_data, bias_dims, stride_width,
- stride_height, 1, 1, pad_width, pad_height, depth_multiplier,
- output_offset, output_multiplier, output_shift,
- output_activation_min, output_activation_max, output_data,
- output_dims);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy, for compatibility with old checked-in code.
-template <FusedActivationFunctionType Ac>
-void DepthwiseConv(const uint8* input_data, const Dims<4>& input_dims,
- int32 input_offset, const uint8* filter_data,
- const Dims<4>& filter_dims, int32 filter_offset,
- const int32* bias_data, const Dims<4>& bias_dims,
- int stride_width, int stride_height, int pad_width,
- int pad_height, int depth_multiplier, int32 output_offset,
- int32 output_multiplier, int output_shift,
- int32 output_activation_min, int32 output_activation_max,
- uint8* output_data, const Dims<4>& output_dims) {
- if (Ac == FusedActivationFunctionType::kNone) {
- TFLITE_DCHECK_EQ(output_activation_min, 0);
- TFLITE_DCHECK_EQ(output_activation_max, 255);
- }
- DepthwiseConv(input_data, input_dims, input_offset, filter_data, filter_dims,
- filter_offset, bias_data, bias_dims, stride_width,
- stride_height, pad_width, pad_height, depth_multiplier,
- output_offset, output_multiplier, output_shift,
- output_activation_min, output_activation_max, output_data,
- output_dims);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy, for compatibility with old checked-in code.
-template <FusedActivationFunctionType Ac>
-void DepthwiseConv(const uint8* input_data, const Dims<4>& input_dims,
- int32 input_offset, const uint8* filter_data,
- const Dims<4>& filter_dims, int32 filter_offset,
- const int32* bias_data, const Dims<4>& bias_dims, int stride,
- int pad_width, int pad_height, int depth_multiplier,
- int32 output_offset, int32 output_multiplier,
- int output_shift, int32 output_activation_min,
- int32 output_activation_max, uint8* output_data,
- const Dims<4>& output_dims) {
- DepthwiseConv<Ac>(input_data, input_dims, input_offset, filter_data,
- filter_dims, filter_offset, bias_data, bias_dims, stride,
- stride, pad_width, pad_height, depth_multiplier,
- output_offset, output_multiplier, output_shift,
- output_activation_min, output_activation_max, output_data,
- output_dims);
-}
-
} // end namespace reference_ops
} // end namespace tflite
diff --git a/tensorflow/contrib/lite/kernels/internal/reference/fully_connected.h b/tensorflow/contrib/lite/kernels/internal/reference/fully_connected.h
index 23325e8..3c7fd29 100644
--- a/tensorflow/contrib/lite/kernels/internal/reference/fully_connected.h
+++ b/tensorflow/contrib/lite/kernels/internal/reference/fully_connected.h
@@ -62,39 +62,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void FullyConnected(const float* input_data, const Dims<4>& input_dims,
- const float* weights_data,
- const Dims<4>& weights_dims, const float* bias_data,
- const Dims<4>& bias_dims,
- float output_activation_min,
- float output_activation_max, float* output_data,
- const Dims<4>& output_dims) {
- tflite::FullyConnectedParams op_params;
- op_params.float_activation_min = output_activation_min;
- op_params.float_activation_max = output_activation_max;
-
- FullyConnected(op_params, DimsToShape(input_dims), input_data,
- DimsToShape(weights_dims), weights_data,
- DimsToShape(bias_dims), bias_data, DimsToShape(output_dims),
- output_data);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// legacy, for compatibility with old checked-in code
-template <FusedActivationFunctionType Ac>
-void FullyConnected(const float* input_data, const Dims<4>& input_dims,
- const float* weights_data, const Dims<4>& weights_dims,
- const float* bias_data, const Dims<4>& bias_dims,
- float* output_data, const Dims<4>& output_dims) {
- float output_activation_min, output_activation_max;
- GetActivationMinMax(Ac, &output_activation_min, &output_activation_max);
- FullyConnected(input_data, input_dims, weights_data, weights_dims, bias_data,
- bias_dims, output_activation_min, output_activation_max,
- output_data, output_dims);
-}
-
inline void FullyConnected(
const FullyConnectedParams& params, const RuntimeShape& input_shape,
const uint8* input_data, const RuntimeShape& filter_shape,
@@ -144,32 +111,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void FullyConnected(const uint8* input_data, const Dims<4>& input_dims,
- int32 input_offset, const uint8* filter_data,
- const Dims<4>& filter_dims, int32 filter_offset,
- const int32* bias_data, const Dims<4>& bias_dims,
- int32 output_offset, int32 output_multiplier,
- int output_shift, int32 output_activation_min,
- int32 output_activation_max, uint8* output_data,
- const Dims<4>& output_dims, void* gemm_context) {
- tflite::FullyConnectedParams op_params;
- op_params.input_offset = input_offset;
- op_params.weights_offset = filter_offset;
- op_params.output_offset = output_offset;
- op_params.output_multiplier = output_multiplier;
- // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
- op_params.output_shift = kReverseShift * output_shift;
- op_params.quantized_activation_min = output_activation_min;
- op_params.quantized_activation_max = output_activation_max;
-
- FullyConnected(op_params, DimsToShape(input_dims), input_data,
- DimsToShape(filter_dims), filter_data, DimsToShape(bias_dims),
- bias_data, DimsToShape(output_dims), output_data,
- gemm_context);
-}
-
inline void FullyConnected(
const FullyConnectedParams& params, const RuntimeShape& input_shape,
const uint8* input_data, const RuntimeShape& filter_shape,
@@ -224,32 +165,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void FullyConnected(const uint8* input_data, const Dims<4>& input_dims,
- int32 input_offset, const uint8* filter_data,
- const Dims<4>& filter_dims, int32 filter_offset,
- const int32* bias_data, const Dims<4>& bias_dims,
- int32 output_offset, int32 output_multiplier,
- int output_shift, int32 output_activation_min,
- int32 output_activation_max, int16* output_data,
- const Dims<4>& output_dims, void* gemm_context) {
- tflite::FullyConnectedParams op_params;
- op_params.input_offset = input_offset;
- op_params.weights_offset = filter_offset;
- op_params.output_offset = output_offset;
- op_params.output_multiplier = output_multiplier;
- // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
- op_params.output_shift = kReverseShift * output_shift;
- op_params.quantized_activation_min = output_activation_min;
- op_params.quantized_activation_max = output_activation_max;
-
- FullyConnected(op_params, DimsToShape(input_dims), input_data,
- DimsToShape(filter_dims), filter_data, DimsToShape(bias_dims),
- bias_data, DimsToShape(output_dims), output_data,
- gemm_context);
-}
-
inline void ShuffledFullyConnected(
const FullyConnectedParams& params, const RuntimeShape& input_shape,
const uint8* input_data, const RuntimeShape& weights_shape,
@@ -405,55 +320,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void ShuffledFullyConnected(
- const uint8* input_data, const Dims<4>& input_dims,
- const uint8* shuffled_weights_data, const Dims<4>& weights_dims,
- const int32* bias_data, const Dims<4>& bias_dims, int32 output_multiplier,
- int output_shift, int32 output_activation_min, int32 output_activation_max,
- int16* output_data, const Dims<4>& output_dims,
- uint8* shuffled_input_workspace_data, void* gemm_context) {
- tflite::FullyConnectedParams op_params;
- op_params.output_multiplier = output_multiplier;
- // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
- op_params.output_shift = kReverseShift * output_shift;
- op_params.quantized_activation_min = output_activation_min;
- op_params.quantized_activation_max = output_activation_max;
-
- ShuffledFullyConnected(op_params, DimsToShape(input_dims), input_data,
- DimsToShape(weights_dims), shuffled_weights_data,
- DimsToShape(bias_dims), bias_data,
- DimsToShape(output_dims), output_data,
- shuffled_input_workspace_data, gemm_context);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// legacy, for compatibility with old checked-in code
-template <FusedActivationFunctionType Ac>
-void FullyConnected(const uint8* input_data, const Dims<4>& input_dims,
- int32 input_offset, const uint8* filter_data,
- const Dims<4>& filter_dims, int32 filter_offset,
- const int32* bias_data, const Dims<4>& bias_dims,
- int32 output_offset, int32 output_multiplier,
- int output_shift, int32 output_activation_min,
- int32 output_activation_max, uint8* output_data,
- const Dims<4>& output_dims, void* gemm_context) {
- static_assert(Ac == FusedActivationFunctionType::kNone ||
- Ac == FusedActivationFunctionType::kRelu ||
- Ac == FusedActivationFunctionType::kRelu6 ||
- Ac == FusedActivationFunctionType::kRelu1,
- "");
- if (Ac == FusedActivationFunctionType::kNone) {
- TFLITE_DCHECK_EQ(output_activation_min, 0);
- TFLITE_DCHECK_EQ(output_activation_max, 255);
- }
- FullyConnected(input_data, input_dims, input_offset, filter_data, filter_dims,
- filter_offset, bias_data, bias_dims, output_offset,
- output_multiplier, output_shift, output_activation_min,
- output_activation_max, output_data, output_dims, gemm_context);
-}
-
} // namespace reference_ops
} // namespace tflite
diff --git a/tensorflow/contrib/lite/kernels/internal/reference/legacy_reference_ops.h b/tensorflow/contrib/lite/kernels/internal/reference/legacy_reference_ops.h
index 683ccdc..be99240 100644
--- a/tensorflow/contrib/lite/kernels/internal/reference/legacy_reference_ops.h
+++ b/tensorflow/contrib/lite/kernels/internal/reference/legacy_reference_ops.h
@@ -19,6 +19,8 @@
#include <sys/types.h>
#include "tensorflow/contrib/lite/kernels/internal/common.h"
+#include "tensorflow/contrib/lite/kernels/internal/reference/depthwiseconv_float.h"
+#include "tensorflow/contrib/lite/kernels/internal/reference/depthwiseconv_uint8.h"
#include "tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h"
#include "tensorflow/contrib/lite/kernels/internal/types.h"
@@ -26,6 +28,1070 @@
namespace reference_ops {
+static constexpr int kDepthwiseReverseShift = -1;
+
+inline void DepthwiseConv(const float* input_data, const Dims<4>& input_dims,
+ const float* filter_data, const Dims<4>& filter_dims,
+ const float* bias_data, const Dims<4>& bias_dims,
+ int stride_width, int stride_height,
+ int dilation_width_factor, int dilation_height_factor,
+ int pad_width, int pad_height, int depth_multiplier,
+ float output_activation_min,
+ float output_activation_max, float* output_data,
+ const Dims<4>& output_dims) {
+ tflite::DepthwiseParams op_params;
+ // Padding type is ignored, but still set.
+ op_params.padding_type = PaddingType::kSame;
+ op_params.padding_values.width = pad_width;
+ op_params.padding_values.height = pad_height;
+ op_params.stride_width = stride_width;
+ op_params.stride_height = stride_height;
+ op_params.dilation_width_factor = dilation_width_factor;
+ op_params.dilation_height_factor = dilation_height_factor;
+ op_params.depth_multiplier = depth_multiplier;
+ op_params.float_activation_min = output_activation_min;
+ op_params.float_activation_max = output_activation_max;
+
+ DepthwiseConv(op_params, DimsToShape(input_dims), input_data,
+ DimsToShape(filter_dims), filter_data, DimsToShape(bias_dims),
+ bias_data, DimsToShape(output_dims), output_data);
+}
+
+inline void DepthwiseConv(const float* input_data, const Dims<4>& input_dims,
+ const float* filter_data, const Dims<4>& filter_dims,
+ const float* bias_data, const Dims<4>& bias_dims,
+ int stride_width, int stride_height, int pad_width,
+ int pad_height, int depth_multiplier,
+ float output_activation_min,
+ float output_activation_max, float* output_data,
+ const Dims<4>& output_dims) {
+ DepthwiseConv(input_data, input_dims, filter_data, filter_dims, bias_data,
+ bias_dims, stride_width, stride_height, 1, 1, pad_width,
+ pad_height, depth_multiplier, output_activation_min,
+ output_activation_max, output_data, output_dims);
+}
+
+// Legacy, for compatibility with old checked-in code.
+template <FusedActivationFunctionType Ac>
+void DepthwiseConv(const float* input_data, const Dims<4>& input_dims,
+ const float* filter_data, const Dims<4>& filter_dims,
+ const float* bias_data, const Dims<4>& bias_dims,
+ int stride_width, int stride_height, int pad_width,
+ int pad_height, int depth_multiplier, float* output_data,
+ const Dims<4>& output_dims) {
+ float output_activation_min, output_activation_max;
+ GetActivationMinMax(Ac, &output_activation_min, &output_activation_max);
+ DepthwiseConv(input_data, input_dims, filter_data, filter_dims, bias_data,
+ bias_dims, stride_width, stride_height, pad_width, pad_height,
+ depth_multiplier, output_activation_min, output_activation_max,
+ output_data, output_dims);
+}
+
+// Legacy, for compatibility with old checked-in code.
+template <FusedActivationFunctionType Ac>
+void DepthwiseConv(const float* input_data, const Dims<4>& input_dims,
+ const float* filter_data, const Dims<4>& filter_dims,
+ const float* bias_data, const Dims<4>& bias_dims, int stride,
+ int pad_width, int pad_height, int depth_multiplier,
+ float* output_data, const Dims<4>& output_dims) {
+ DepthwiseConv<Ac>(input_data, input_dims, filter_data, filter_dims, bias_data,
+ bias_dims, stride, stride, pad_width, pad_height,
+ depth_multiplier, output_data, output_dims);
+}
+
+inline void DepthwiseConv(const uint8* input_data, const Dims<4>& input_dims,
+ int32 input_offset, const uint8* filter_data,
+ const Dims<4>& filter_dims, int32 filter_offset,
+ const int32* bias_data, const Dims<4>& bias_dims,
+ int stride_width, int stride_height,
+ int dilation_width_factor, int dilation_height_factor,
+ int pad_width, int pad_height, int depth_multiplier,
+ int32 output_offset, int32 output_multiplier,
+ int output_shift, int32 output_activation_min,
+ int32 output_activation_max, uint8* output_data,
+ const Dims<4>& output_dims) {
+ tflite::DepthwiseParams op_params;
+ // Padding type is ignored, but still set.
+ op_params.padding_type = PaddingType::kSame;
+ op_params.padding_values.width = pad_width;
+ op_params.padding_values.height = pad_height;
+ op_params.stride_width = stride_width;
+ op_params.stride_height = stride_height;
+ op_params.dilation_width_factor = dilation_width_factor;
+ op_params.dilation_height_factor = dilation_height_factor;
+ op_params.depth_multiplier = depth_multiplier;
+ op_params.quantized_activation_min = output_activation_min;
+ op_params.quantized_activation_max = output_activation_max;
+ op_params.input_offset = input_offset;
+ op_params.weights_offset = filter_offset;
+ op_params.output_offset = output_offset;
+ op_params.output_multiplier = output_multiplier;
+ // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
+ op_params.output_shift = kDepthwiseReverseShift * output_shift;
+
+ DepthwiseConv(op_params, DimsToShape(input_dims), input_data,
+ DimsToShape(filter_dims), filter_data, DimsToShape(bias_dims),
+ bias_data, DimsToShape(output_dims), output_data);
+}
+
+inline void DepthwiseConv(const uint8* input_data, const Dims<4>& input_dims,
+ int32 input_offset, const uint8* filter_data,
+ const Dims<4>& filter_dims, int32 filter_offset,
+ const int32* bias_data, const Dims<4>& bias_dims,
+ int stride_width, int stride_height, int pad_width,
+ int pad_height, int depth_multiplier,
+ int32 output_offset, int32 output_multiplier,
+ int output_shift, int32 output_activation_min,
+ int32 output_activation_max, uint8* output_data,
+ const Dims<4>& output_dims) {
+ DepthwiseConv(input_data, input_dims, input_offset, filter_data, filter_dims,
+ filter_offset, bias_data, bias_dims, stride_width,
+ stride_height, 1, 1, pad_width, pad_height, depth_multiplier,
+ output_offset, output_multiplier, output_shift,
+ output_activation_min, output_activation_max, output_data,
+ output_dims);
+}
+
+// Legacy, for compatibility with old checked-in code.
+template <FusedActivationFunctionType Ac>
+void DepthwiseConv(const uint8* input_data, const Dims<4>& input_dims,
+ int32 input_offset, const uint8* filter_data,
+ const Dims<4>& filter_dims, int32 filter_offset,
+ const int32* bias_data, const Dims<4>& bias_dims,
+ int stride_width, int stride_height, int pad_width,
+ int pad_height, int depth_multiplier, int32 output_offset,
+ int32 output_multiplier, int output_shift,
+ int32 output_activation_min, int32 output_activation_max,
+ uint8* output_data, const Dims<4>& output_dims) {
+ if (Ac == FusedActivationFunctionType::kNone) {
+ TFLITE_DCHECK_EQ(output_activation_min, 0);
+ TFLITE_DCHECK_EQ(output_activation_max, 255);
+ }
+ DepthwiseConv(input_data, input_dims, input_offset, filter_data, filter_dims,
+ filter_offset, bias_data, bias_dims, stride_width,
+ stride_height, pad_width, pad_height, depth_multiplier,
+ output_offset, output_multiplier, output_shift,
+ output_activation_min, output_activation_max, output_data,
+ output_dims);
+}
+
+// Legacy, for compatibility with old checked-in code.
+template <FusedActivationFunctionType Ac>
+void DepthwiseConv(const uint8* input_data, const Dims<4>& input_dims,
+ int32 input_offset, const uint8* filter_data,
+ const Dims<4>& filter_dims, int32 filter_offset,
+ const int32* bias_data, const Dims<4>& bias_dims, int stride,
+ int pad_width, int pad_height, int depth_multiplier,
+ int32 output_offset, int32 output_multiplier,
+ int output_shift, int32 output_activation_min,
+ int32 output_activation_max, uint8* output_data,
+ const Dims<4>& output_dims) {
+ DepthwiseConv<Ac>(input_data, input_dims, input_offset, filter_data,
+ filter_dims, filter_offset, bias_data, bias_dims, stride,
+ stride, pad_width, pad_height, depth_multiplier,
+ output_offset, output_multiplier, output_shift,
+ output_activation_min, output_activation_max, output_data,
+ output_dims);
+}
+
+inline void Conv(const float* input_data, const Dims<4>& input_dims,
+ const float* filter_data, const Dims<4>& filter_dims,
+ const float* bias_data, const Dims<4>& bias_dims,
+ int stride_width, int stride_height, int dilation_width_factor,
+ int dilation_height_factor, int pad_width, int pad_height,
+ float output_activation_min, float output_activation_max,
+ float* output_data, const Dims<4>& output_dims,
+ float* im2col_data, const Dims<4>& im2col_dims) {
+ tflite::ConvParams op_params;
+ // Padding type is ignored, but still set.
+ op_params.padding_type = PaddingType::kSame;
+ op_params.padding_values.width = pad_width;
+ op_params.padding_values.height = pad_height;
+ op_params.stride_width = stride_width;
+ op_params.stride_height = stride_height;
+ op_params.dilation_width_factor = dilation_width_factor;
+ op_params.dilation_height_factor = dilation_height_factor;
+ op_params.float_activation_min = output_activation_min;
+ op_params.float_activation_max = output_activation_max;
+
+ Conv(op_params, DimsToShape(input_dims), input_data, DimsToShape(filter_dims),
+ filter_data, DimsToShape(bias_dims), bias_data, DimsToShape(output_dims),
+ output_data, DimsToShape(im2col_dims), im2col_data);
+}
+
+template <FusedActivationFunctionType Ac>
+void Conv(const float* input_data, const Dims<4>& input_dims,
+ const float* filter_data, const Dims<4>& filter_dims,
+ const float* bias_data, const Dims<4>& bias_dims, int stride_width,
+ int stride_height, int dilation_width_factor,
+ int dilation_height_factor, int pad_width, int pad_height,
+ float* output_data, const Dims<4>& output_dims, float* im2col_data,
+ const Dims<4>& im2col_dims) {
+ float output_activation_min, output_activation_max;
+ GetActivationMinMax(Ac, &output_activation_min, &output_activation_max);
+ Conv(input_data, input_dims, filter_data, filter_dims, bias_data, bias_dims,
+ stride_width, stride_height, dilation_width_factor,
+ dilation_height_factor, pad_width, pad_height, output_activation_min,
+ output_activation_max, output_data, output_dims, im2col_data,
+ im2col_dims);
+}
+
+// legacy, for compatibility with old checked-in code
+template <FusedActivationFunctionType Ac>
+void Conv(const float* input_data, const Dims<4>& input_dims,
+ const float* filter_data, const Dims<4>& filter_dims,
+ const float* bias_data, const Dims<4>& bias_dims, int stride_width,
+ int stride_height, int pad_width, int pad_height, float* output_data,
+ const Dims<4>& output_dims, float* im2col_data,
+ const Dims<4>& im2col_dims) {
+ float output_activation_min, output_activation_max;
+ GetActivationMinMax(Ac, &output_activation_min, &output_activation_max);
+ Conv(input_data, input_dims, filter_data, filter_dims, bias_data, bias_dims,
+ stride_width, stride_height, 1, 1, pad_width, pad_height,
+ output_activation_min, output_activation_max, output_data, output_dims,
+ im2col_data, im2col_dims);
+}
+
+// legacy, for compatibility with old checked-in code
+template <FusedActivationFunctionType Ac>
+void Conv(const float* input_data, const Dims<4>& input_dims,
+ const float* filter_data, const Dims<4>& filter_dims,
+ const float* bias_data, const Dims<4>& bias_dims, int stride,
+ int pad_width, int pad_height, float* output_data,
+ const Dims<4>& output_dims, float* im2col_data,
+ const Dims<4>& im2col_dims) {
+ Conv<Ac>(input_data, input_dims, filter_data, filter_dims, bias_data,
+ bias_dims, stride, stride, 1, 1, pad_width, pad_height, output_data,
+ output_dims, im2col_data, im2col_dims);
+}
+
+inline void Conv(const uint8* input_data, const Dims<4>& input_dims,
+ int32 input_offset, const uint8* filter_data,
+ const Dims<4>& filter_dims, int32 filter_offset,
+ const int32* bias_data, const Dims<4>& bias_dims,
+ int stride_width, int stride_height, int dilation_width_factor,
+ int dilation_height_factor, int pad_width, int pad_height,
+ int32 output_offset, int32 output_multiplier, int output_shift,
+ int32 output_activation_min, int32 output_activation_max,
+ uint8* output_data, const Dims<4>& output_dims,
+ uint8* im2col_data, const Dims<4>& im2col_dims,
+ gemmlowp::GemmContext* gemm_context) {
+ tflite::ConvParams op_params;
+ // Padding type is ignored, but still set.
+ op_params.padding_type = PaddingType::kSame;
+ op_params.padding_values.width = pad_width;
+ op_params.padding_values.height = pad_height;
+ op_params.stride_width = stride_width;
+ op_params.stride_height = stride_height;
+ op_params.dilation_width_factor = dilation_width_factor;
+ op_params.dilation_height_factor = dilation_height_factor;
+ op_params.input_offset = input_offset;
+ op_params.weights_offset = filter_offset;
+ op_params.output_offset = output_offset;
+ op_params.output_multiplier = output_multiplier;
+ // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
+ op_params.output_shift = kReverseShift * output_shift;
+ op_params.quantized_activation_min = output_activation_min;
+ op_params.quantized_activation_max = output_activation_max;
+
+ Conv(op_params, DimsToShape(input_dims), input_data, DimsToShape(filter_dims),
+ filter_data, DimsToShape(bias_dims), bias_data, DimsToShape(output_dims),
+ output_data, DimsToShape(im2col_dims), im2col_data, gemm_context);
+}
+
+inline void Conv(const uint8* input_data, const Dims<4>& input_dims,
+ int32 input_offset, const uint8* filter_data,
+ const Dims<4>& filter_dims, int32 filter_offset,
+ const int32* bias_data, const Dims<4>& bias_dims,
+ int stride_width, int stride_height, int pad_width,
+ int pad_height, int32 output_offset, int32 output_multiplier,
+ int output_shift, int32 output_activation_min,
+ int32 output_activation_max, uint8* output_data,
+ const Dims<4>& output_dims, uint8* im2col_data,
+ const Dims<4>& im2col_dims,
+ gemmlowp::GemmContext* gemm_context) {
+ Conv(input_data, input_dims, input_offset, filter_data, filter_dims,
+ filter_offset, bias_data, bias_dims, stride_width, stride_height, 1, 1,
+ pad_width, pad_height, output_offset, output_multiplier, output_shift,
+ output_activation_min, output_activation_max, output_data, output_dims,
+ im2col_data, im2col_dims, gemm_context);
+}
+
+// legacy, for compatibility with old checked-in code
+template <FusedActivationFunctionType Ac>
+inline void Conv(const uint8* input_data, const Dims<4>& input_dims,
+ int32 input_offset, const uint8* filter_data,
+ const Dims<4>& filter_dims, int32 filter_offset,
+ const int32* bias_data, const Dims<4>& bias_dims,
+ int stride_width, int stride_height, int pad_width,
+ int pad_height, int32 output_offset, int32 output_multiplier,
+ int output_shift, int32 output_activation_min,
+ int32 output_activation_max, uint8* output_data,
+ const Dims<4>& output_dims, uint8* im2col_data,
+ const Dims<4>& im2col_dims,
+ gemmlowp::GemmContext* gemm_context) {
+ static_assert(Ac == FusedActivationFunctionType::kNone ||
+ Ac == FusedActivationFunctionType::kRelu ||
+ Ac == FusedActivationFunctionType::kRelu6 ||
+ Ac == FusedActivationFunctionType::kRelu1,
+ "");
+ if (Ac == FusedActivationFunctionType::kNone) {
+ TFLITE_DCHECK_EQ(output_activation_min, 0);
+ TFLITE_DCHECK_EQ(output_activation_max, 255);
+ }
+ Conv(input_data, input_dims, input_offset, filter_data, filter_dims,
+ filter_offset, bias_data, bias_dims, stride_width, stride_height,
+ pad_width, pad_height, output_offset, output_multiplier, output_shift,
+ output_activation_min, output_activation_max, output_data, output_dims,
+ im2col_data, im2col_dims, gemm_context);
+}
+
+// legacy, for compatibility with old checked-in code
+template <FusedActivationFunctionType Ac>
+void Conv(const uint8* input_data, const Dims<4>& input_dims,
+ int32 input_offset, const uint8* filter_data,
+ const Dims<4>& filter_dims, int32 filter_offset,
+ const int32* bias_data, const Dims<4>& bias_dims, int stride,
+ int pad_width, int pad_height, int32 output_offset,
+ int32 output_multiplier, int output_shift,
+ int32 output_activation_min, int32 output_activation_max,
+ uint8* output_data, const Dims<4>& output_dims, uint8* im2col_data,
+ const Dims<4>& im2col_dims, gemmlowp::GemmContext* gemm_context) {
+ Conv<Ac>(input_data, input_dims, input_offset, filter_data, filter_dims,
+ filter_offset, bias_data, bias_dims, stride, stride, pad_width,
+ pad_height, output_offset, output_multiplier, output_shift,
+ output_activation_min, output_activation_max, output_data,
+ output_dims, im2col_data, im2col_dims, gemm_context);
+}
+
+inline void TransposeConv(const float* input_data, const Dims<4>& input_dims,
+ const float* filter_data, const Dims<4>& filter_dims,
+ int stride_width, int stride_height, int pad_width,
+ int pad_height, float* output_data,
+ const Dims<4>& output_dims, float* im2col_data,
+ const Dims<4>& im2col_dims) {
+ tflite::ConvParams op_params;
+ // Padding type is ignored, but still set.
+ op_params.padding_type = PaddingType::kSame;
+ op_params.padding_values.width = pad_width;
+ op_params.padding_values.height = pad_height;
+ op_params.stride_width = stride_width;
+ op_params.stride_height = stride_height;
+
+ TransposeConv(op_params, DimsToShape(input_dims), input_data,
+ DimsToShape(filter_dims), filter_data, DimsToShape(output_dims),
+ output_data, DimsToShape(im2col_dims), im2col_data);
+}
+
+inline void FullyConnected(const float* input_data, const Dims<4>& input_dims,
+ const float* weights_data,
+ const Dims<4>& weights_dims, const float* bias_data,
+ const Dims<4>& bias_dims,
+ float output_activation_min,
+ float output_activation_max, float* output_data,
+ const Dims<4>& output_dims) {
+ tflite::FullyConnectedParams op_params;
+ op_params.float_activation_min = output_activation_min;
+ op_params.float_activation_max = output_activation_max;
+
+ FullyConnected(op_params, DimsToShape(input_dims), input_data,
+ DimsToShape(weights_dims), weights_data,
+ DimsToShape(bias_dims), bias_data, DimsToShape(output_dims),
+ output_data);
+}
+
+// legacy, for compatibility with old checked-in code
+template <FusedActivationFunctionType Ac>
+void FullyConnected(const float* input_data, const Dims<4>& input_dims,
+ const float* weights_data, const Dims<4>& weights_dims,
+ const float* bias_data, const Dims<4>& bias_dims,
+ float* output_data, const Dims<4>& output_dims) {
+ float output_activation_min, output_activation_max;
+ GetActivationMinMax(Ac, &output_activation_min, &output_activation_max);
+ FullyConnected(input_data, input_dims, weights_data, weights_dims, bias_data,
+ bias_dims, output_activation_min, output_activation_max,
+ output_data, output_dims);
+}
+
+inline void FullyConnected(const uint8* input_data, const Dims<4>& input_dims,
+ int32 input_offset, const uint8* filter_data,
+ const Dims<4>& filter_dims, int32 filter_offset,
+ const int32* bias_data, const Dims<4>& bias_dims,
+ int32 output_offset, int32 output_multiplier,
+ int output_shift, int32 output_activation_min,
+ int32 output_activation_max, uint8* output_data,
+ const Dims<4>& output_dims,
+ gemmlowp::GemmContext* gemm_context) {
+ tflite::FullyConnectedParams op_params;
+ op_params.input_offset = input_offset;
+ op_params.weights_offset = filter_offset;
+ op_params.output_offset = output_offset;
+ op_params.output_multiplier = output_multiplier;
+ // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
+ op_params.output_shift = kReverseShift * output_shift;
+ op_params.quantized_activation_min = output_activation_min;
+ op_params.quantized_activation_max = output_activation_max;
+
+ FullyConnected(op_params, DimsToShape(input_dims), input_data,
+ DimsToShape(filter_dims), filter_data, DimsToShape(bias_dims),
+ bias_data, DimsToShape(output_dims), output_data,
+ gemm_context);
+}
+
+inline void FullyConnected(const uint8* input_data, const Dims<4>& input_dims,
+ int32 input_offset, const uint8* filter_data,
+ const Dims<4>& filter_dims, int32 filter_offset,
+ const int32* bias_data, const Dims<4>& bias_dims,
+ int32 output_offset, int32 output_multiplier,
+ int output_shift, int32 output_activation_min,
+ int32 output_activation_max, int16* output_data,
+ const Dims<4>& output_dims,
+ gemmlowp::GemmContext* gemm_context) {
+ tflite::FullyConnectedParams op_params;
+ op_params.input_offset = input_offset;
+ op_params.weights_offset = filter_offset;
+ op_params.output_offset = output_offset;
+ op_params.output_multiplier = output_multiplier;
+ // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
+ op_params.output_shift = kReverseShift * output_shift;
+ op_params.quantized_activation_min = output_activation_min;
+ op_params.quantized_activation_max = output_activation_max;
+
+ FullyConnected(op_params, DimsToShape(input_dims), input_data,
+ DimsToShape(filter_dims), filter_data, DimsToShape(bias_dims),
+ bias_data, DimsToShape(output_dims), output_data,
+ gemm_context);
+}
+
+inline void ShuffledFullyConnected(
+ const uint8* input_data, const Dims<4>& input_dims,
+ const uint8* shuffled_weights_data, const Dims<4>& weights_dims,
+ const int32* bias_data, const Dims<4>& bias_dims, int32 output_multiplier,
+ int output_shift, int32 output_activation_min, int32 output_activation_max,
+ int16* output_data, const Dims<4>& output_dims,
+ uint8* shuffled_input_workspace_data, gemmlowp::GemmContext* gemm_context) {
+ tflite::FullyConnectedParams op_params;
+ op_params.output_multiplier = output_multiplier;
+ // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
+ op_params.output_shift = kReverseShift * output_shift;
+ op_params.quantized_activation_min = output_activation_min;
+ op_params.quantized_activation_max = output_activation_max;
+
+ ShuffledFullyConnected(op_params, DimsToShape(input_dims), input_data,
+ DimsToShape(weights_dims), shuffled_weights_data,
+ DimsToShape(bias_dims), bias_data,
+ DimsToShape(output_dims), output_data,
+ shuffled_input_workspace_data, gemm_context);
+}
+
+// legacy, for compatibility with old checked-in code
+template <FusedActivationFunctionType Ac>
+void FullyConnected(const uint8* input_data, const Dims<4>& input_dims,
+ int32 input_offset, const uint8* filter_data,
+ const Dims<4>& filter_dims, int32 filter_offset,
+ const int32* bias_data, const Dims<4>& bias_dims,
+ int32 output_offset, int32 output_multiplier,
+ int output_shift, int32 output_activation_min,
+ int32 output_activation_max, uint8* output_data,
+ const Dims<4>& output_dims,
+ gemmlowp::GemmContext* gemm_context) {
+ static_assert(Ac == FusedActivationFunctionType::kNone ||
+ Ac == FusedActivationFunctionType::kRelu ||
+ Ac == FusedActivationFunctionType::kRelu6 ||
+ Ac == FusedActivationFunctionType::kRelu1,
+ "");
+ if (Ac == FusedActivationFunctionType::kNone) {
+ TFLITE_DCHECK_EQ(output_activation_min, 0);
+ TFLITE_DCHECK_EQ(output_activation_max, 255);
+ }
+ FullyConnected(input_data, input_dims, input_offset, filter_data, filter_dims,
+ filter_offset, bias_data, bias_dims, output_offset,
+ output_multiplier, output_shift, output_activation_min,
+ output_activation_max, output_data, output_dims, gemm_context);
+}
+
+inline void LstmCell(const float* input_data, const Dims<4>& input_dims,
+ const float* prev_activ_data,
+ const Dims<4>& prev_activ_dims, const float* weights_data,
+ const Dims<4>& weights_dims, const float* bias_data,
+ const Dims<4>& bias_dims, const float* prev_state_data,
+ const Dims<4>& prev_state_dims, float* output_state_data,
+ const Dims<4>& output_state_dims, float* output_activ_data,
+ const Dims<4>& output_activ_dims, float* concat_temp_data,
+ const Dims<4>& concat_temp_dims, float* activ_temp_data,
+ const Dims<4>& activ_temp_dims) {
+ tflite::LstmCellParams op_params;
+ // Float LSTM cell does not need parameters to be set: leave untouched.
+
+ LstmCell(op_params, DimsToShape(input_dims), input_data,
+ DimsToShape(prev_activ_dims), prev_activ_data,
+ DimsToShape(weights_dims), weights_data, DimsToShape(bias_dims),
+ bias_data, DimsToShape(prev_state_dims), prev_state_data,
+ DimsToShape(output_state_dims), output_state_data,
+ DimsToShape(output_activ_dims), output_activ_data,
+ DimsToShape(concat_temp_dims), concat_temp_data,
+ DimsToShape(activ_temp_dims), activ_temp_data);
+}
+
+template <int StateIntegerBits>
+void LstmCell(const uint8* input_data_uint8, const Dims<4>& input_dims,
+ const uint8* prev_activ_data_uint8,
+ const Dims<4>& prev_activ_dims, const uint8* weights_data_uint8,
+ const Dims<4>& weights_dims, const int32* bias_data_int32,
+ const Dims<4>& bias_dims, const int16* prev_state_data_int16,
+ const Dims<4>& prev_state_dims, int16* output_state_data_int16,
+ const Dims<4>& output_state_dims, uint8* output_activ_data_uint8,
+ const Dims<4>& output_activ_dims, uint8* concat_temp_data_uint8,
+ const Dims<4>& concat_temp_dims, int16* activ_temp_data_int16,
+ const Dims<4>& activ_temp_dims, int32 weights_zero_point,
+ int32 accum_multiplier, int accum_shift,
+ gemmlowp::GemmContext* gemm_context) {
+ tflite::LstmCellParams op_params;
+ op_params.weights_zero_point = weights_zero_point;
+ op_params.accum_multiplier = accum_multiplier;
+ op_params.accum_shift = accum_shift;
+
+ LstmCell<StateIntegerBits>(
+ op_params, DimsToShape(input_dims), input_data_uint8,
+ DimsToShape(prev_activ_dims), prev_activ_data_uint8,
+ DimsToShape(weights_dims), weights_data_uint8, DimsToShape(bias_dims),
+ bias_data_int32, DimsToShape(prev_state_dims), prev_state_data_int16,
+ DimsToShape(output_state_dims), output_state_data_int16,
+ DimsToShape(output_activ_dims), output_activ_data_uint8,
+ DimsToShape(concat_temp_dims), concat_temp_data_uint8,
+ DimsToShape(activ_temp_dims), activ_temp_data_int16, gemm_context);
+}
+
+template <typename T>
+void BroadcastDiv(const T* input1_data, const Dims<4>& input1_dims,
+ const T* input2_data, const Dims<4>& input2_dims,
+ T output_activation_min, T output_activation_max,
+ T* output_data, const Dims<4>& output_dims) {
+ tflite::ArithmeticParams op_params;
+ SetActivationParams(output_activation_min, output_activation_max, &op_params);
+
+ BroadcastDiv4DSlow(op_params, DimsToShape(input1_dims), input1_data,
+ DimsToShape(input2_dims), input2_data,
+ DimsToShape(output_dims), output_data);
+}
+
+template <typename T>
+inline void Div(const T* input1_data, const Dims<4>& input1_dims,
+ const T* input2_data, const Dims<4>& input2_dims,
+ T output_activation_min, T output_activation_max,
+ T* output_data, const Dims<4>& output_dims) {
+ tflite::ArithmeticParams op_params;
+ SetActivationParams(output_activation_min, output_activation_max, &op_params);
+
+ Div(op_params, DimsToShape(input1_dims), input1_data,
+ DimsToShape(input2_dims), input2_data, DimsToShape(output_dims),
+ output_data);
+}
+
+template <FusedActivationFunctionType Ac, typename Scalar>
+inline void Concatenation(int concat_dim, const Scalar* const* input_data,
+ const Dims<4>* const* input_dims, int inputs_count,
+ Scalar* output_data, const Dims<4>& output_dims) {
+ // For now we don't have a model with a Concatenation with fused activation.
+ TFLITE_DCHECK_EQ(Ac, FusedActivationFunctionType::kNone);
+
+ std::vector<RuntimeShape> input_shapes(inputs_count);
+ std::vector<const RuntimeShape*> input_shapes_indirect(inputs_count);
+ for (int i = 0; i < inputs_count; ++i) {
+ ShapeFromDims(*input_dims[i], &input_shapes[i]);
+ input_shapes_indirect[i] = &input_shapes[i];
+ }
+ tflite::ConcatenationParams op_params;
+ op_params.axis = 3 - concat_dim;
+ op_params.inputs_count = inputs_count;
+
+ Concatenation(op_params, input_shapes_indirect.data(), input_data,
+ DimsToShape(output_dims), output_data);
+}
+
+inline void Concatenation(int concat_dim, const uint8* const* input_data,
+ const Dims<4>* const* input_dims,
+ const int32* input_zeropoint,
+ const float* input_scale, int inputs_count,
+ uint8* output_data, const Dims<4>& output_dims,
+ const int32 output_zeropoint,
+ const float output_scale) {
+ std::vector<RuntimeShape> input_shapes(inputs_count);
+ std::vector<const RuntimeShape*> input_shapes_indirect(inputs_count);
+ for (int i = 0; i < inputs_count; ++i) {
+ ShapeFromDims(*input_dims[i], &input_shapes[i]);
+ input_shapes_indirect[i] = &input_shapes[i];
+ }
+ tflite::ConcatenationParams op_params;
+ op_params.axis = 3 - concat_dim;
+ op_params.input_zeropoint = input_zeropoint;
+ op_params.input_scale = input_scale;
+ op_params.inputs_count = inputs_count;
+ op_params.output_zeropoint = output_zeropoint;
+ op_params.output_scale = output_scale;
+
+ ConcatenationWithScaling(op_params, input_shapes_indirect.data(), input_data,
+ DimsToShape(output_dims), output_data);
+}
+
+template <FusedActivationFunctionType Ac, typename Scalar>
+void DepthConcatenation(const Scalar* const* input_data,
+ const Dims<4>* const* input_dims, int inputs_count,
+ Scalar* output_data, const Dims<4>& output_dims) {
+ // For now we don't have a model with a Concatenation with fused activation.
+ TFLITE_DCHECK_EQ(Ac, FusedActivationFunctionType::kNone);
+
+ std::vector<RuntimeShape> input_shapes(inputs_count);
+ std::vector<const RuntimeShape*> input_shapes_indirect(inputs_count);
+ for (int i = 0; i < inputs_count; ++i) {
+ ShapeFromDims(*input_dims[i], &input_shapes[i]);
+ input_shapes_indirect[i] = &input_shapes[i];
+ }
+ tflite::ConcatenationParams op_params;
+ op_params.inputs_count = inputs_count;
+
+ DepthConcatenation(op_params, input_shapes_indirect.data(), input_data,
+ DimsToShape(output_dims), output_data);
+}
+
+template <typename Scalar>
+void TensorFlowSplit(const Scalar* input_data, const Dims<4>& input_dims,
+ int axis, int outputs_count, Scalar* const* output_data,
+ const Dims<4>* const* output_dims) {
+ std::vector<RuntimeShape> output_shapes(outputs_count);
+ std::vector<const RuntimeShape*> output_shapes_indirect(outputs_count);
+ for (int i = 0; i < outputs_count; ++i) {
+ ShapeFromDims(*output_dims[i], &output_shapes[i]);
+ output_shapes_indirect[i] = &output_shapes[i];
+ }
+ tflite::SplitParams op_params;
+ op_params.axis = 3 - axis;
+ op_params.num_split = outputs_count;
+
+ Split(op_params, DimsToShape(input_dims), input_data,
+ output_shapes_indirect.data(), output_data);
+}
+
+template <FusedActivationFunctionType Ac, typename Scalar>
+void TensorFlowSplit(const Scalar* input_data, const Dims<4>& input_dims,
+ int outputs_count, Scalar* const* output_data,
+ const Dims<4>* const* output_dims) {
+ TFLITE_DCHECK_GE(outputs_count, 1);
+ for (int i = 0; i < outputs_count; i++) {
+ /* batches = */ MatchingArraySize(*output_dims[i], 3, input_dims, 3);
+ /* height = */ MatchingArraySize(*output_dims[i], 2, input_dims, 2);
+ /* width = */ MatchingArraySize(*output_dims[i], 1, input_dims, 1);
+ }
+ // For now we don't have a model with a Split with fused activation.
+ TFLITE_DCHECK_EQ(Ac, FusedActivationFunctionType::kNone);
+
+ TensorFlowSplit(input_data, input_dims, /*axis=*/0, outputs_count,
+ output_data, output_dims);
+}
+
+inline void Softmax(const float* input_data, const RuntimeShape& input_shape,
+ float beta, float* output_data,
+ const RuntimeShape& output_shape) {
+ SoftmaxParams params;
+ params.beta = beta;
+ Softmax(params, input_shape, input_data, output_shape, output_data);
+}
+
+inline void Softmax(const uint8* input_data, const RuntimeShape& input_shape,
+ int32 input_beta_multiplier, int32 input_beta_left_shift,
+ int diff_min, uint8* output_data,
+ const RuntimeShape& output_shape) {
+ SoftmaxParams params;
+ params.input_multiplier = input_beta_multiplier;
+ params.input_left_shift = input_beta_left_shift;
+ params.diff_min = diff_min;
+ Softmax(params, input_shape, input_data, output_shape, output_data);
+}
+
+inline void LogSoftmax(const float* input_data, const RuntimeShape& input_shape,
+ float* output_data, const RuntimeShape& output_shape) {
+ SoftmaxParams params;
+ // No params currently used for float LogSoftmax.
+ LogSoftmax(params, input_shape, input_data, output_shape, output_data);
+}
+
+inline void LogSoftmax(const uint8* input_data, const RuntimeShape& input_shape,
+ int32 input_multiplier, int32 input_left_shift,
+ int32 reverse_scaling_divisor,
+ int32 reverse_scaling_right_shift, int diff_min,
+ uint8* output_data, const RuntimeShape& output_shape) {
+ SoftmaxParams params;
+ params.input_multiplier = input_multiplier;
+ params.input_left_shift = input_left_shift;
+ params.reverse_scaling_divisor = reverse_scaling_divisor;
+ params.reverse_scaling_right_shift = reverse_scaling_right_shift;
+ params.diff_min = diff_min;
+ LogSoftmax(params, input_shape, input_data, output_shape, output_data);
+}
+
+inline void Logistic(const uint8* input_data, const RuntimeShape& input_shape,
+ int32 input_zero_point, int32 input_range_radius,
+ int32 input_multiplier, int input_left_shift,
+ uint8* output_data, const RuntimeShape& output_shape) {
+ LogisticParams params;
+ params.input_zero_point = input_zero_point;
+ params.input_range_radius = input_range_radius;
+ params.input_multiplier = input_multiplier;
+ params.input_left_shift = input_left_shift;
+ Logistic(params, input_shape, input_data, output_shape, output_data);
+}
+
+inline void Logistic(const RuntimeShape& input_shape, const int16* input_data,
+ const RuntimeShape& output_shape, int16* output_data) {
+ LogisticParams params;
+ // No params currently needed by int16 Logistic.
+ Logistic(params, input_shape, input_data, output_shape, output_data);
+}
+
+inline void Tanh(const uint8* input_data, const RuntimeShape& input_shape,
+ int32 input_zero_point, int32 input_range_radius,
+ int32 input_multiplier, int input_left_shift,
+ uint8* output_data, const RuntimeShape& output_shape) {
+ TanhParams params;
+ params.input_zero_point = input_zero_point;
+ params.input_range_radius = input_range_radius;
+ params.input_multiplier = input_multiplier;
+ params.input_left_shift = input_left_shift;
+ Tanh(params, input_shape, input_data, output_shape, output_data);
+}
+
+inline void Tanh(const int16* input_data, const RuntimeShape& input_shape,
+ int input_left_shift, int16* output_data,
+ const RuntimeShape& output_shape) {
+ TanhParams params;
+ params.input_left_shift = input_left_shift;
+ Tanh(params, input_shape, input_data, output_shape, output_data);
+}
+
+inline void Dequantize(const uint8* input_data, const Dims<4>& input_dims,
+ int32 zero_point, double scale, float* output_data,
+ const Dims<4>& output_dims) {
+ tflite::DequantizationParams op_params;
+ op_params.zero_point = zero_point;
+ op_params.scale = scale;
+
+ Dequantize(op_params, DimsToShape(input_dims), input_data,
+ DimsToShape(output_dims), output_data);
+}
+
+inline void FakeQuant(const float* input_data, const Dims<4>& input_dims,
+ float rmin, float rmax, int num_bits, float* output_data,
+ const Dims<4>& output_dims) {
+ tflite::FakeQuantParams op_params;
+ op_params.num_bits = num_bits;
+ op_params.minmax.min = rmin;
+ op_params.minmax.max = rmax;
+
+ FakeQuant(op_params, DimsToShape(input_dims), input_data,
+ DimsToShape(output_dims), output_data);
+}
+
+template <typename T>
+inline void Gather(const T* input_data, const Dims<4>& input_dims,
+ int input_rank, const int32* coords_data,
+ const Dims<4>& coords_dims, T* output_data,
+ const Dims<4>& output_dims) {
+ tflite::GatherParams op_params;
+ op_params.input_rank = input_rank;
+
+ Gather(op_params, DimsToShape(input_dims), input_data,
+ DimsToShape(coords_dims), coords_data, DimsToShape(output_dims),
+ output_data);
+}
+
+inline uint32 LegacyReverseBits32(uint32 n) {
+ n = ((n >> 1) & 0x55555555) | ((n & 0x55555555) << 1);
+ n = ((n >> 2) & 0x33333333) | ((n & 0x33333333) << 2);
+ n = ((n >> 4) & 0x0F0F0F0F) | ((n & 0x0F0F0F0F) << 4);
+ return (((n & 0xFF) << 24) | ((n & 0xFF00) << 8) | ((n & 0xFF0000) >> 8) |
+ ((n & 0xFF000000) >> 24));
+}
+
+inline void StridedSliceReverseIndices(tflite::StridedSliceParams* p) {
+ TFLITE_CHECK_EQ(p->start_indices_count, p->stop_indices_count);
+ TFLITE_CHECK_EQ(p->stop_indices_count, p->strides_count);
+
+ std::reverse(p->start_indices, p->start_indices + p->start_indices_count);
+ std::reverse(p->stop_indices, p->stop_indices + p->stop_indices_count);
+ std::reverse(p->strides, p->strides + p->strides_count);
+
+ p->begin_mask = LegacyReverseBits32(static_cast<uint32>(p->begin_mask)) >>
+ (32 - p->start_indices_count);
+ p->ellipsis_mask =
+ LegacyReverseBits32(static_cast<uint32>(p->ellipsis_mask)) >>
+ (32 - p->start_indices_count);
+ p->end_mask = LegacyReverseBits32(static_cast<uint32>(p->end_mask)) >>
+ (32 - p->start_indices_count);
+ p->new_axis_mask =
+ LegacyReverseBits32(static_cast<uint32>(p->new_axis_mask)) >>
+ (32 - p->start_indices_count);
+ p->shrink_axis_mask =
+ LegacyReverseBits32(static_cast<uint32>(p->shrink_axis_mask)) >>
+ (32 - p->start_indices_count);
+}
+
+template <typename T>
+inline void StridedSlice(const T* input_data, const Dims<4>& input_dims,
+ int begin_mask, int end_mask, int shrink_axis_mask,
+ const std::vector<int>& start_indices,
+ const std::vector<int>& stop_indices,
+ const std::vector<int>& strides, T* output_data,
+ const Dims<4>& output_dims) {
+ TFLITE_DCHECK_EQ(start_indices.size(), 4);
+ auto op_params = strided_slice::BuildStridedSliceParams(
+ begin_mask, end_mask, shrink_axis_mask, start_indices, stop_indices,
+ strides);
+ StridedSliceReverseIndices(&op_params);
+
+ StridedSlice(op_params, DimsToShape(input_dims), input_data,
+ DimsToShape(output_dims), output_data);
+}
+
+template <typename T>
+inline void Mean(const T* input_data, const Dims<4>& input_dims,
+ const std::vector<int>& reduction_indices, T* output_data,
+ const Dims<4>& output_dims) {
+ tflite::MeanParams op_params;
+ op_params.axis_count = reduction_indices.size();
+ for (int i = 0; i < op_params.axis_count; ++i) {
+ op_params.axis[i] = reduction_indices[op_params.axis_count - 1 - i];
+ }
+
+ Mean(op_params, DimsToShape(input_dims), input_data, DimsToShape(output_dims),
+ output_data);
+}
+
+template <typename T>
+void Transpose(const T* input, const Dims<4>& input_dims, T* output,
+ const Dims<4>& output_dims, const int* permuted_axes) {
+ TransposeParams params;
+ params.perm_count = 4;
+ for (int i = 0; i < 4; ++i) {
+ params.perm[i] = 3 - permuted_axes[3 - i];
+ }
+ Transpose(params, DimsToShape(input_dims), input, DimsToShape(output_dims),
+ output);
+}
+
+template <typename T, ComparisonFn<T> F>
+inline void Comparison(const T* input1_data, const Dims<4>& input1_dims,
+ const T* input2_data, const Dims<4>& input2_dims,
+ bool* output_data, const Dims<4>& output_dims) {
+ ComparisonParams op_params;
+ // No parameters needed.
+ ComparisonImpl<T, F>(op_params, DimsToShape(input1_dims), input1_data,
+ DimsToShape(input2_dims), input2_data,
+ DimsToShape(output_dims), output_data);
+}
+
+template <typename T, ComparisonFn<int32> F>
+inline void Comparison(int left_shift, const T* input1_data,
+ const Dims<4>& input1_dims, int32 input1_offset,
+ int32 input1_multiplier, int input1_shift,
+ const T* input2_data, const Dims<4>& input2_dims,
+ int32 input2_offset, int32 input2_multiplier,
+ int input2_shift, bool* output_data,
+ const Dims<4>& output_dims) {
+ tflite::ComparisonParams op_params;
+ op_params.left_shift = left_shift;
+ op_params.input1_offset = input1_offset;
+ op_params.input1_multiplier = input1_multiplier;
+ // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
+ op_params.input1_shift = kReverseShift * input1_shift;
+ op_params.input2_offset = input2_offset;
+ op_params.input2_multiplier = input2_multiplier;
+ // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
+ op_params.input2_shift = kReverseShift * input2_shift;
+
+ ComparisonWithScaling<T, F>(op_params, DimsToShape(input1_dims), input1_data,
+ DimsToShape(input2_dims), input2_data,
+ DimsToShape(output_dims), output_data);
+}
+
+template <typename T, ComparisonFn<T> F>
+inline void BroadcastComparison(const T* input1_data,
+ const Dims<4>& input1_dims,
+ const T* input2_data,
+ const Dims<4>& input2_dims, bool* output_data,
+ const Dims<4>& output_dims) {
+ ComparisonParams op_params;
+ // No parameters needed.
+ BroadcastComparison4DSlowImpl<T, F>(op_params, DimsToShape(input1_dims),
+ input1_data, DimsToShape(input2_dims),
+ input2_data, DimsToShape(output_dims),
+ output_data);
+}
+
+template <typename T, ComparisonFn<int32> F>
+inline void BroadcastComparison(int left_shift, const T* input1_data,
+ const Dims<4>& input1_dims, int32 input1_offset,
+ int32 input1_multiplier, int input1_shift,
+ const T* input2_data,
+ const Dims<4>& input2_dims, int32 input2_offset,
+ int32 input2_multiplier, int input2_shift,
+ bool* output_data, const Dims<4>& output_dims) {
+ ComparisonParams op_params;
+
+ op_params.left_shift = left_shift;
+ op_params.input1_offset = input1_offset;
+ op_params.input1_multiplier = input1_multiplier;
+ // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
+ op_params.input1_shift = kReverseShift * input1_shift;
+ op_params.input2_offset = input2_offset;
+ op_params.input2_multiplier = input2_multiplier;
+ // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
+ op_params.input2_shift = kReverseShift * input2_shift;
+
+ BroadcastComparison4DSlowWithScaling<T, F>(
+ op_params, DimsToShape(input1_dims), input1_data,
+ DimsToShape(input2_dims), input2_data, DimsToShape(output_dims),
+ output_data);
+}
+
+#define TFLITE_LEGACY_COMPARISON_OP(name) \
+ template <typename T> \
+ inline void name(const T* input1_data, const Dims<4>& input1_dims, \
+ const T* input2_data, const Dims<4>& input2_dims, \
+ bool* output_data, const Dims<4>& output_dims) { \
+ gemmlowp::ScopedProfilingLabel label(#name); \
+ Comparison<T, name##Fn>(input1_data, input1_dims, input2_data, \
+ input2_dims, output_data, output_dims); \
+ } \
+ template <typename T> \
+ inline void name( \
+ int left_shift, const T* input1_data, const Dims<4>& input1_dims, \
+ int32 input1_offset, int32 input1_multiplier, int input1_shift, \
+ const T* input2_data, const Dims<4>& input2_dims, int32 input2_offset, \
+ int32 input2_multiplier, int input2_shift, bool* output_data, \
+ const Dims<4>& output_dims) { \
+ gemmlowp::ScopedProfilingLabel label(#name "/8bit"); \
+ Comparison<T, name##Fn>(left_shift, input1_data, input1_dims, \
+ input1_offset, input1_multiplier, input1_shift, \
+ input2_data, input2_dims, input2_offset, \
+ input2_multiplier, input2_shift, output_data, \
+ output_dims); \
+ } \
+ template <typename T> \
+ inline void Broadcast##name( \
+ const T* input1_data, const Dims<4>& input1_dims, const T* input2_data, \
+ const Dims<4>& input2_dims, bool* output_data, \
+ const Dims<4>& output_dims) { \
+ gemmlowp::ScopedProfilingLabel label("Broadcast" #name); \
+ BroadcastComparison<T, name##Fn>(input1_data, input1_dims, input2_data, \
+ input2_dims, output_data, output_dims); \
+ } \
+ template <typename T> \
+ inline void Broadcast##name( \
+ int left_shift, const T* input1_data, const Dims<4>& input1_dims, \
+ int32 input1_offset, int32 input1_multiplier, int input1_shift, \
+ const T* input2_data, const Dims<4>& input2_dims, int32 input2_offset, \
+ int32 input2_multiplier, int input2_shift, bool* output_data, \
+ const Dims<4>& output_dims) { \
+ gemmlowp::ScopedProfilingLabel label("Broadcast" #name "/8bit"); \
+ BroadcastComparison<T, name##Fn>(left_shift, input1_data, input1_dims, \
+ input1_offset, input1_multiplier, \
+ input1_shift, input2_data, input2_dims, \
+ input2_offset, input2_multiplier, \
+ input2_shift, output_data, output_dims); \
+ }
+TFLITE_LEGACY_COMPARISON_OP(Equal);
+TFLITE_LEGACY_COMPARISON_OP(NotEqual);
+TFLITE_LEGACY_COMPARISON_OP(Greater);
+TFLITE_LEGACY_COMPARISON_OP(GreaterEqual);
+TFLITE_LEGACY_COMPARISON_OP(Less);
+TFLITE_LEGACY_COMPARISON_OP(LessEqual);
+#undef TFLITE_LEGACY_COMPARISON_OP
+
+template <typename D, typename T>
+inline void Select(const D* input_condition_data,
+ const Dims<4>& input_condition_dims, const T* input_x_data,
+ const Dims<4>& input_x_dims, const T* input_y_data,
+ const Dims<4>& input_y_dims, T* output_data,
+ const Dims<4>& output_dims) {
+ Select(DimsToShape(input_condition_dims), input_condition_data,
+ DimsToShape(input_x_dims), input_x_data, DimsToShape(input_y_dims),
+ input_y_data, DimsToShape(output_dims), output_data);
+}
+
+template <typename D, typename T>
+inline void RankOneSelect(const D* input_condition_data,
+ const Dims<4>& input_condition_dims,
+ const T* input_x_data, const Dims<4>& input_x_dims,
+ const T* input_y_data, const Dims<4>& input_y_dims,
+ T* output_data, const Dims<4>& output_dims) {
+ RankOneSelect(DimsToShape(input_condition_dims), input_condition_data,
+ DimsToShape(input_x_dims), input_x_data,
+ DimsToShape(input_y_dims), input_y_data,
+ DimsToShape(output_dims), output_data);
+}
+
+template <typename T, typename TI>
+inline void SparseToDense(const std::vector<std::vector<TI>>& indices,
+ const T* values, T default_value, T* output_data,
+ const Dims<4>& output_dims, bool value_is_scalar) {
+ SparseToDense(indices, values, default_value, value_is_scalar,
+ DimsToShape(output_dims), output_data);
+}
+
+template <typename Scalar>
+void Pack(int dim, const Scalar* const* input_data,
+ const Dims<4>* const* input_dims, int inputs_count,
+ Scalar* output_data, const Dims<4>& output_dims) {
+ std::vector<RuntimeShape> input_shapes(inputs_count);
+ std::vector<const RuntimeShape*> input_shapes_indirect(inputs_count);
+ for (int i = 0; i < inputs_count; ++i) {
+ ShapeFromDims(*input_dims[i], &input_shapes[i]);
+ input_shapes_indirect[i] = &input_shapes[i];
+ }
+ tflite::PackParams op_params;
+ op_params.axis = 3 - dim;
+ op_params.inputs_count = inputs_count;
+
+ Pack(op_params, input_shapes_indirect.data(), input_data,
+ DimsToShape(output_dims), output_data);
+}
+
+template <typename Scalar>
+void Unpack(int axis, const Scalar* input_data, const Dims<4>& input_dims,
+ int dimensions, int outputs_count, Scalar* const* output_datas,
+ const Dims<4>& output_dims) {
+ tflite::UnpackParams op_params;
+ op_params.axis = 3 - axis;
+ op_params.num_split = outputs_count;
+
+ Unpack(op_params, DimsToShape(input_dims), input_data,
+ DimsToShape(output_dims), output_datas);
+}
+
+template <typename Scalar>
+void Pack(int dim, const Scalar* const* input_data,
+ const Dims<4>* const* input_dims, const int32* input_zeropoint,
+ const float* input_scale, int inputs_count, Scalar* output_data,
+ const Dims<4>& output_dims, const int32 output_zeropoint,
+ const float output_scale) {
+ std::vector<RuntimeShape> input_shapes(inputs_count);
+ std::vector<const RuntimeShape*> input_shapes_indirect(inputs_count);
+ for (int i = 0; i < inputs_count; ++i) {
+ ShapeFromDims(*input_dims[i], &input_shapes[i]);
+ input_shapes_indirect[i] = &input_shapes[i];
+ }
+ tflite::PackParams op_params;
+ op_params.axis = 3 - dim;
+ op_params.input_zeropoint = input_zeropoint;
+ op_params.input_scale = input_scale;
+ op_params.inputs_count = inputs_count;
+ op_params.output_zeropoint = output_zeropoint;
+ op_params.output_scale = output_scale;
+
+ PackWithScaling(op_params, input_shapes_indirect.data(), input_data,
+ DimsToShape(output_dims), output_data);
+}
+
template <FusedActivationFunctionType Ac>
void L2Normalization(const float* input_data, const RuntimeShape& input_shape,
float* output_data, const RuntimeShape& output_shape) {
@@ -342,7 +1408,6 @@
DimsToShape(output_dims), output_data);
}
-// Legacy.
// Transitional version that will be moved shortly to legacy_reference_ops, as
// part of RuntimeShape revisions.
inline void BroadcastMul4DSlow(const uint8* input1_data,
diff --git a/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h b/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h
index 7a55354..59f17ae 100644
--- a/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h
+++ b/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h
@@ -231,83 +231,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void Conv(const float* input_data, const Dims<4>& input_dims,
- const float* filter_data, const Dims<4>& filter_dims,
- const float* bias_data, const Dims<4>& bias_dims,
- int stride_width, int stride_height, int dilation_width_factor,
- int dilation_height_factor, int pad_width, int pad_height,
- float output_activation_min, float output_activation_max,
- float* output_data, const Dims<4>& output_dims,
- float* im2col_data, const Dims<4>& im2col_dims) {
- tflite::ConvParams op_params;
- // Padding type is ignored, but still set.
- op_params.padding_type = PaddingType::kSame;
- op_params.padding_values.width = pad_width;
- op_params.padding_values.height = pad_height;
- op_params.stride_width = stride_width;
- op_params.stride_height = stride_height;
- op_params.dilation_width_factor = dilation_width_factor;
- op_params.dilation_height_factor = dilation_height_factor;
- op_params.float_activation_min = output_activation_min;
- op_params.float_activation_max = output_activation_max;
-
- Conv(op_params, DimsToShape(input_dims), input_data, DimsToShape(filter_dims),
- filter_data, DimsToShape(bias_dims), bias_data, DimsToShape(output_dims),
- output_data, DimsToShape(im2col_dims), im2col_data);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-template <FusedActivationFunctionType Ac>
-void Conv(const float* input_data, const Dims<4>& input_dims,
- const float* filter_data, const Dims<4>& filter_dims,
- const float* bias_data, const Dims<4>& bias_dims, int stride_width,
- int stride_height, int dilation_width_factor,
- int dilation_height_factor, int pad_width, int pad_height,
- float* output_data, const Dims<4>& output_dims, float* im2col_data,
- const Dims<4>& im2col_dims) {
- float output_activation_min, output_activation_max;
- GetActivationMinMax(Ac, &output_activation_min, &output_activation_max);
- Conv(input_data, input_dims, filter_data, filter_dims, bias_data, bias_dims,
- stride_width, stride_height, dilation_width_factor,
- dilation_height_factor, pad_width, pad_height, output_activation_min,
- output_activation_max, output_data, output_dims, im2col_data,
- im2col_dims);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// legacy, for compatibility with old checked-in code
-template <FusedActivationFunctionType Ac>
-void Conv(const float* input_data, const Dims<4>& input_dims,
- const float* filter_data, const Dims<4>& filter_dims,
- const float* bias_data, const Dims<4>& bias_dims, int stride_width,
- int stride_height, int pad_width, int pad_height, float* output_data,
- const Dims<4>& output_dims, float* im2col_data,
- const Dims<4>& im2col_dims) {
- float output_activation_min, output_activation_max;
- GetActivationMinMax(Ac, &output_activation_min, &output_activation_max);
- Conv(input_data, input_dims, filter_data, filter_dims, bias_data, bias_dims,
- stride_width, stride_height, 1, 1, pad_width, pad_height,
- output_activation_min, output_activation_max, output_data, output_dims,
- im2col_data, im2col_dims);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// legacy, for compatibility with old checked-in code
-template <FusedActivationFunctionType Ac>
-void Conv(const float* input_data, const Dims<4>& input_dims,
- const float* filter_data, const Dims<4>& filter_dims,
- const float* bias_data, const Dims<4>& bias_dims, int stride,
- int pad_width, int pad_height, float* output_data,
- const Dims<4>& output_dims, float* im2col_data,
- const Dims<4>& im2col_dims) {
- Conv<Ac>(input_data, input_dims, filter_data, filter_dims, bias_data,
- bias_dims, stride, stride, 1, 1, pad_width, pad_height, output_data,
- output_dims, im2col_data, im2col_dims);
-}
-
inline void Conv(const ConvParams& params, const RuntimeShape& input_shape,
const uint8* input_data, const RuntimeShape& filter_shape,
const uint8* filter_data, const RuntimeShape& bias_shape,
@@ -391,111 +314,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void Conv(const uint8* input_data, const Dims<4>& input_dims,
- int32 input_offset, const uint8* filter_data,
- const Dims<4>& filter_dims, int32 filter_offset,
- const int32* bias_data, const Dims<4>& bias_dims,
- int stride_width, int stride_height, int dilation_width_factor,
- int dilation_height_factor, int pad_width, int pad_height,
- int32 output_offset, int32 output_multiplier, int output_shift,
- int32 output_activation_min, int32 output_activation_max,
- uint8* output_data, const Dims<4>& output_dims,
- uint8* im2col_data, const Dims<4>& im2col_dims,
- gemmlowp::GemmContext* gemm_context) {
- tflite::ConvParams op_params;
- // Padding type is ignored, but still set.
- op_params.padding_type = PaddingType::kSame;
- op_params.padding_values.width = pad_width;
- op_params.padding_values.height = pad_height;
- op_params.stride_width = stride_width;
- op_params.stride_height = stride_height;
- op_params.dilation_width_factor = dilation_width_factor;
- op_params.dilation_height_factor = dilation_height_factor;
- op_params.input_offset = input_offset;
- op_params.weights_offset = filter_offset;
- op_params.output_offset = output_offset;
- op_params.output_multiplier = output_multiplier;
- // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
- op_params.output_shift = kReverseShift * output_shift;
- op_params.quantized_activation_min = output_activation_min;
- op_params.quantized_activation_max = output_activation_max;
-
- Conv(op_params, DimsToShape(input_dims), input_data, DimsToShape(filter_dims),
- filter_data, DimsToShape(bias_dims), bias_data, DimsToShape(output_dims),
- output_data, DimsToShape(im2col_dims), im2col_data, gemm_context);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void Conv(const uint8* input_data, const Dims<4>& input_dims,
- int32 input_offset, const uint8* filter_data,
- const Dims<4>& filter_dims, int32 filter_offset,
- const int32* bias_data, const Dims<4>& bias_dims,
- int stride_width, int stride_height, int pad_width,
- int pad_height, int32 output_offset, int32 output_multiplier,
- int output_shift, int32 output_activation_min,
- int32 output_activation_max, uint8* output_data,
- const Dims<4>& output_dims, uint8* im2col_data,
- const Dims<4>& im2col_dims,
- gemmlowp::GemmContext* gemm_context) {
- Conv(input_data, input_dims, input_offset, filter_data, filter_dims,
- filter_offset, bias_data, bias_dims, stride_width, stride_height, 1, 1,
- pad_width, pad_height, output_offset, output_multiplier, output_shift,
- output_activation_min, output_activation_max, output_data, output_dims,
- im2col_data, im2col_dims, gemm_context);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// legacy, for compatibility with old checked-in code
-template <FusedActivationFunctionType Ac>
-inline void Conv(const uint8* input_data, const Dims<4>& input_dims,
- int32 input_offset, const uint8* filter_data,
- const Dims<4>& filter_dims, int32 filter_offset,
- const int32* bias_data, const Dims<4>& bias_dims,
- int stride_width, int stride_height, int pad_width,
- int pad_height, int32 output_offset, int32 output_multiplier,
- int output_shift, int32 output_activation_min,
- int32 output_activation_max, uint8* output_data,
- const Dims<4>& output_dims, uint8* im2col_data,
- const Dims<4>& im2col_dims,
- gemmlowp::GemmContext* gemm_context) {
- static_assert(Ac == FusedActivationFunctionType::kNone ||
- Ac == FusedActivationFunctionType::kRelu ||
- Ac == FusedActivationFunctionType::kRelu6 ||
- Ac == FusedActivationFunctionType::kRelu1,
- "");
- if (Ac == FusedActivationFunctionType::kNone) {
- TFLITE_DCHECK_EQ(output_activation_min, 0);
- TFLITE_DCHECK_EQ(output_activation_max, 255);
- }
- Conv(input_data, input_dims, input_offset, filter_data, filter_dims,
- filter_offset, bias_data, bias_dims, stride_width, stride_height,
- pad_width, pad_height, output_offset, output_multiplier, output_shift,
- output_activation_min, output_activation_max, output_data, output_dims,
- im2col_data, im2col_dims, gemm_context);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// legacy, for compatibility with old checked-in code
-template <FusedActivationFunctionType Ac>
-void Conv(const uint8* input_data, const Dims<4>& input_dims,
- int32 input_offset, const uint8* filter_data,
- const Dims<4>& filter_dims, int32 filter_offset,
- const int32* bias_data, const Dims<4>& bias_dims, int stride,
- int pad_width, int pad_height, int32 output_offset,
- int32 output_multiplier, int output_shift,
- int32 output_activation_min, int32 output_activation_max,
- uint8* output_data, const Dims<4>& output_dims, uint8* im2col_data,
- const Dims<4>& im2col_dims, gemmlowp::GemmContext* gemm_context) {
- Conv<Ac>(input_data, input_dims, input_offset, filter_data, filter_dims,
- filter_offset, bias_data, bias_dims, stride, stride, pad_width,
- pad_height, output_offset, output_multiplier, output_shift,
- output_activation_min, output_activation_max, output_data,
- output_dims, im2col_data, im2col_dims, gemm_context);
-}
-
template <typename T>
inline void DepthToSpace(const tflite::DepthToSpaceParams& op_params,
const RuntimeShape& unextended_input_shape,
@@ -1385,21 +1203,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy Dims<4>.
-template <typename T>
-void BroadcastDiv(const T* input1_data, const Dims<4>& input1_dims,
- const T* input2_data, const Dims<4>& input2_dims,
- T output_activation_min, T output_activation_max,
- T* output_data, const Dims<4>& output_dims) {
- tflite::ArithmeticParams op_params;
- SetActivationParams(output_activation_min, output_activation_max, &op_params);
-
- BroadcastDiv4DSlow(op_params, DimsToShape(input1_dims), input1_data,
- DimsToShape(input2_dims), input2_data,
- DimsToShape(output_dims), output_data);
-}
-
template <typename T>
inline void Div(const ArithmeticParams& params,
const RuntimeShape& input1_shape, const T* input1_data,
@@ -1418,21 +1221,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy Dims<4>.
-template <typename T>
-inline void Div(const T* input1_data, const Dims<4>& input1_dims,
- const T* input2_data, const Dims<4>& input2_dims,
- T output_activation_min, T output_activation_max,
- T* output_data, const Dims<4>& output_dims) {
- tflite::ArithmeticParams op_params;
- SetActivationParams(output_activation_min, output_activation_max, &op_params);
-
- Div(op_params, DimsToShape(input1_dims), input1_data,
- DimsToShape(input2_dims), input2_data, DimsToShape(output_dims),
- output_data);
-}
-
inline void SubNonBroadcast(const ArithmeticParams& params,
const RuntimeShape& input1_shape,
const float* input1_data,
@@ -1772,35 +1560,10 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy Dims<4>.
-template <FusedActivationFunctionType Ac, typename Scalar>
-inline void Concatenation(int concat_dim, const Scalar* const* input_data,
- const Dims<4>* const* input_dims, int inputs_count,
- Scalar* output_data, const Dims<4>& output_dims) {
- // For now we don't have a model with a Concatenation with fused activation.
- TFLITE_DCHECK_EQ(Ac, FusedActivationFunctionType::kNone);
-
- std::vector<RuntimeShape> input_shapes(inputs_count);
- std::vector<const RuntimeShape*> input_shapes_indirect(inputs_count);
- for (int i = 0; i < inputs_count; ++i) {
- ShapeFromDims(*input_dims[i], &input_shapes[i]);
- input_shapes_indirect[i] = &input_shapes[i];
- }
- tflite::ConcatenationParams op_params;
- op_params.axis = 3 - concat_dim;
- op_params.inputs_count = inputs_count;
-
- Concatenation(op_params, input_shapes_indirect.data(), input_data,
- DimsToShape(output_dims), output_data);
-}
-
// TODO(prabhumk): This is the same as the optimized implementation.
// TODO(prabhumk): The quantized implementation of concatentation isn't fully
// quantized as it takes scale as a floating point value. This should be fixed
// when optimizng this routine further.
-
-// template <>
inline void ConcatenationWithScaling(const ConcatenationParams& params,
const RuntimeShape* const* input_shapes,
const uint8* const* input_data,
@@ -1813,15 +1576,13 @@
const int32 output_zeropoint = params.output_zeropoint;
const float output_scale = params.output_scale;
- // The arguments input_zeropoint and input_scale are expected to be an array
- // that have the quantization parameters for all the inputs to the concat
- // operator.
- TFLITE_DCHECK_GT(inputs_count, 1);
- TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
+ const int concat_dimensions = output_shape.DimensionsCount();
+ TFLITE_DCHECK_LT(axis, concat_dimensions);
+
int64_t concat_size = 0;
for (int i = 0; i < inputs_count; i++) {
- TFLITE_DCHECK_EQ(input_shapes[i]->DimensionsCount(), 4);
- for (int j = 0; j < 4; j++) {
+ TFLITE_DCHECK_EQ(input_shapes[i]->DimensionsCount(), concat_dimensions);
+ for (int j = 0; j < concat_dimensions; j++) {
if (j != axis) {
MatchingDim(*input_shapes[i], j, output_shape, j);
}
@@ -1836,9 +1597,10 @@
// For all input arrays,
// FlatSize() = outer_size * Dims(axis) * base_inner_size;
int64_t base_inner_size = 1;
- for (int i = axis + 1; i < 4; ++i) {
+ for (int i = axis + 1; i < concat_dimensions; ++i) {
base_inner_size *= output_shape.Dims(i);
}
+
const float inverse_output_scale = 1.f / output_scale;
uint8* output_ptr = output_data;
for (int k = 0; k < outer_size; k++) {
@@ -1864,65 +1626,52 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy Dims<4>.
-inline void Concatenation(int concat_dim, const uint8* const* input_data,
- const Dims<4>* const* input_dims,
- const int32* input_zeropoint,
- const float* input_scale, int inputs_count,
- uint8* output_data, const Dims<4>& output_dims,
- const int32 output_zeropoint,
- const float output_scale) {
- std::vector<RuntimeShape> input_shapes(inputs_count);
- std::vector<const RuntimeShape*> input_shapes_indirect(inputs_count);
- for (int i = 0; i < inputs_count; ++i) {
- ShapeFromDims(*input_dims[i], &input_shapes[i]);
- input_shapes_indirect[i] = &input_shapes[i];
- }
- tflite::ConcatenationParams op_params;
- op_params.axis = 3 - concat_dim;
- op_params.input_zeropoint = input_zeropoint;
- op_params.input_scale = input_scale;
- op_params.inputs_count = inputs_count;
- op_params.output_zeropoint = output_zeropoint;
- op_params.output_scale = output_scale;
-
- ConcatenationWithScaling(op_params, input_shapes_indirect.data(), input_data,
- DimsToShape(output_dims), output_data);
-}
-
template <typename Scalar>
-void Pack(int dim, const Scalar* const* input_data,
- const Dims<4>* const* input_dims, int inputs_count,
- Scalar* output_data, const Dims<4>& output_dims) {
- TFLITE_DCHECK(IsPackedWithoutStrides(output_dims));
+void Pack(const PackParams& params, const RuntimeShape* const* input_shapes,
+ const Scalar* const* input_data, const RuntimeShape& output_shape,
+ Scalar* output_data) {
+ const int dimensions = output_shape.DimensionsCount();
+ int axis = params.axis;
+ int inputs_count = params.inputs_count;
+
int outer_size = 1;
- for (int i = dim + 1; i < 4; i++) {
- outer_size *= output_dims.sizes[i];
+ for (int i = 0; i < axis; i++) {
+ outer_size *= output_shape.Dims(i);
}
- Scalar* output_ptr = output_data;
- const int copy_size = FlatSize(**input_dims) / outer_size;
- for (int k = 0; k < outer_size; k++) {
- for (int i = 0; i < inputs_count; ++i) {
- memcpy(output_ptr, input_data[i] + k * copy_size,
- copy_size * sizeof(Scalar));
- output_ptr += copy_size;
+ int copy_size = 1;
+ for (int i = params.axis + 1; i < dimensions; i++) {
+ copy_size *= output_shape.Dims(i);
+ }
+ TFLITE_DCHECK_EQ((**input_shapes).FlatSize(), copy_size * outer_size);
+
+ for (int i = 0; i < inputs_count; ++i) {
+ for (int k = 0; k < outer_size; k++) {
+ const Scalar* input_ptr = input_data[i] + copy_size * k;
+ int loc = k * inputs_count * copy_size + i * copy_size;
+ memcpy(output_data + loc, input_ptr, copy_size * sizeof(Scalar));
}
}
}
template <typename Scalar>
-void Unpack(int axis, const Scalar* input_data, const Dims<4>& input_dims,
- int dimensions, int outputs_count, Scalar* const* output_datas,
- const Dims<4>& output_dims) {
- int outer_size = 1;
- for (int i = dimensions - axis; i < 4; i++) {
- outer_size *= input_dims.sizes[i];
- }
+void Unpack(const UnpackParams& params, const RuntimeShape& input_shape,
+ const Scalar* input_data, const RuntimeShape& output_shape,
+ Scalar* const* output_datas) {
+ const int dimensions = input_shape.DimensionsCount();
+ const int outputs_count = params.num_split;
- const int copy_size = FlatSize(input_dims) / outer_size / outputs_count;
- for (int k = 0; k < outer_size; k++) {
- for (int i = 0; i < outputs_count; ++i) {
+ int outer_size = 1;
+ for (int i = 0; i < params.axis; i++) {
+ outer_size *= input_shape.Dims(i);
+ }
+ int copy_size = 1;
+ for (int i = params.axis + 1; i < dimensions; i++) {
+ copy_size *= input_shape.Dims(i);
+ }
+ TFLITE_DCHECK_EQ(output_shape.FlatSize(), copy_size * outer_size);
+
+ for (int i = 0; i < outputs_count; ++i) {
+ for (int k = 0; k < outer_size; k++) {
Scalar* output_ptr = output_datas[i] + copy_size * k;
int loc = k * outputs_count * copy_size + i * copy_size;
memcpy(output_ptr, input_data + loc, copy_size * sizeof(Scalar));
@@ -1931,18 +1680,29 @@
}
template <typename Scalar>
-void Pack(int dim, const Scalar* const* input_data,
- const Dims<4>* const* input_dims, const int32* input_zeropoint,
- const float* input_scale, int inputs_count, Scalar* output_data,
- const Dims<4>& output_dims, const int32 output_zeropoint,
- const float output_scale) {
- TFLITE_DCHECK(IsPackedWithoutStrides(output_dims));
+void PackWithScaling(const PackParams& params,
+ const RuntimeShape* const* input_shapes,
+ const uint8* const* input_data,
+ const RuntimeShape& output_shape, uint8* output_data) {
+ const int dimensions = output_shape.DimensionsCount();
+ int axis = params.axis;
+ const int32* input_zeropoint = params.input_zeropoint;
+ const float* input_scale = params.input_scale;
+ int inputs_count = params.inputs_count;
+ const int32 output_zeropoint = params.output_zeropoint;
+ const float output_scale = params.output_scale;
+
int outer_size = 1;
- for (int i = dim + 1; i < 4; i++) {
- outer_size *= output_dims.sizes[i];
+ for (int i = 0; i < axis; i++) {
+ outer_size *= output_shape.Dims(i);
}
+ int copy_size = 1;
+ for (int i = axis + 1; i < dimensions; i++) {
+ copy_size *= output_shape.Dims(i);
+ }
+ TFLITE_DCHECK_EQ((**input_shapes).FlatSize(), copy_size * outer_size);
+
Scalar* output_ptr = output_data;
- const int copy_size = FlatSize(**input_dims) / outer_size;
const float inverse_output_scale = 1.f / output_scale;
for (int k = 0; k < outer_size; k++) {
for (int i = 0; i < inputs_count; ++i) {
@@ -1968,12 +1728,15 @@
}
}
-template <FusedActivationFunctionType Ac, typename Scalar>
-void DepthConcatenation(const Scalar* const* input_data,
- const Dims<4>* const* input_dims, int inputs_count,
- Scalar* output_data, const Dims<4>& output_dims) {
- Concatenation<Ac, Scalar>(0, input_data, input_dims, inputs_count,
- output_data, output_dims);
+template <typename Scalar>
+void DepthConcatenation(const ConcatenationParams& params,
+ const RuntimeShape* const* input_shapes,
+ const Scalar* const* input_data,
+ const RuntimeShape& output_shape, Scalar* output_data) {
+ auto params_copy = params;
+ params_copy.axis = 3;
+ Concatenation(params_copy, input_shapes, input_data, output_shape,
+ output_data);
}
inline void LstmCell(
@@ -2093,31 +1856,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void LstmCell(const float* input_data, const Dims<4>& input_dims,
- const float* prev_activ_data,
- const Dims<4>& prev_activ_dims, const float* weights_data,
- const Dims<4>& weights_dims, const float* bias_data,
- const Dims<4>& bias_dims, const float* prev_state_data,
- const Dims<4>& prev_state_dims, float* output_state_data,
- const Dims<4>& output_state_dims, float* output_activ_data,
- const Dims<4>& output_activ_dims, float* concat_temp_data,
- const Dims<4>& concat_temp_dims, float* activ_temp_data,
- const Dims<4>& activ_temp_dims) {
- tflite::LstmCellParams op_params;
- // Float LSTM cell does not need parameters to be set: leave untouched.
-
- LstmCell(op_params, DimsToShape(input_dims), input_data,
- DimsToShape(prev_activ_dims), prev_activ_data,
- DimsToShape(weights_dims), weights_data, DimsToShape(bias_dims),
- bias_data, DimsToShape(prev_state_dims), prev_state_data,
- DimsToShape(output_state_dims), output_state_data,
- DimsToShape(output_activ_dims), output_activ_data,
- DimsToShape(concat_temp_dims), concat_temp_data,
- DimsToShape(activ_temp_dims), activ_temp_data);
-}
-
// Quantized LSTM cell implementation.
// The quantization of the input, output arrays is as follows:
// - The input activations are quantized as uint8 on the interval
@@ -2392,37 +2130,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-template <int StateIntegerBits>
-void LstmCell(const uint8* input_data_uint8, const Dims<4>& input_dims,
- const uint8* prev_activ_data_uint8,
- const Dims<4>& prev_activ_dims, const uint8* weights_data_uint8,
- const Dims<4>& weights_dims, const int32* bias_data_int32,
- const Dims<4>& bias_dims, const int16* prev_state_data_int16,
- const Dims<4>& prev_state_dims, int16* output_state_data_int16,
- const Dims<4>& output_state_dims, uint8* output_activ_data_uint8,
- const Dims<4>& output_activ_dims, uint8* concat_temp_data_uint8,
- const Dims<4>& concat_temp_dims, int16* activ_temp_data_int16,
- const Dims<4>& activ_temp_dims, int32 weights_zero_point,
- int32 accum_multiplier, int accum_shift,
- gemmlowp::GemmContext* gemm_context) {
- tflite::LstmCellParams op_params;
- op_params.weights_zero_point = weights_zero_point;
- op_params.accum_multiplier = accum_multiplier;
- op_params.accum_shift = accum_shift;
-
- LstmCell<StateIntegerBits>(
- op_params, DimsToShape(input_dims), input_data_uint8,
- DimsToShape(prev_activ_dims), prev_activ_data_uint8,
- DimsToShape(weights_dims), weights_data_uint8, DimsToShape(bias_dims),
- bias_data_int32, DimsToShape(prev_state_dims), prev_state_data_int16,
- DimsToShape(output_state_dims), output_state_data_int16,
- DimsToShape(output_activ_dims), output_activ_data_uint8,
- DimsToShape(concat_temp_dims), concat_temp_data_uint8,
- DimsToShape(activ_temp_dims), activ_temp_data_int16, gemm_context);
-}
-
template <typename Scalar>
void Split(const SplitParams& params, const RuntimeShape& input_shape,
const Scalar* input_data, const RuntimeShape* const* output_shapes,
@@ -2465,45 +2172,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy Dims<4>.
-template <typename Scalar>
-void TensorFlowSplit(const Scalar* input_data, const Dims<4>& input_dims,
- int axis, int outputs_count, Scalar* const* output_data,
- const Dims<4>* const* output_dims) {
- std::vector<RuntimeShape> output_shapes(outputs_count);
- std::vector<const RuntimeShape*> output_shapes_indirect(outputs_count);
- for (int i = 0; i < outputs_count; ++i) {
- ShapeFromDims(*output_dims[i], &output_shapes[i]);
- output_shapes_indirect[i] = &output_shapes[i];
- }
- tflite::SplitParams op_params;
- op_params.axis = 3 - axis;
- op_params.num_split = outputs_count;
-
- Split(op_params, DimsToShape(input_dims), input_data,
- output_shapes_indirect.data(), output_data);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy Dims<4>.
-template <FusedActivationFunctionType Ac, typename Scalar>
-void TensorFlowSplit(const Scalar* input_data, const Dims<4>& input_dims,
- int outputs_count, Scalar* const* output_data,
- const Dims<4>* const* output_dims) {
- TFLITE_DCHECK_GE(outputs_count, 1);
- for (int i = 0; i < outputs_count; i++) {
- /* batches = */ MatchingArraySize(*output_dims[i], 3, input_dims, 3);
- /* height = */ MatchingArraySize(*output_dims[i], 2, input_dims, 2);
- /* width = */ MatchingArraySize(*output_dims[i], 1, input_dims, 1);
- }
- // For now we don't have a model with a Split with fused activation.
- TFLITE_DCHECK_EQ(Ac, FusedActivationFunctionType::kNone);
-
- TensorFlowSplit(input_data, input_dims, /*axis=*/0, outputs_count,
- output_data, output_dims);
-}
-
inline int NodeOffset(int b, int h, int w, int height, int width) {
return (b * height + h) * width + w;
}
@@ -2834,15 +2502,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy
-inline void LogSoftmax(const float* input_data, const RuntimeShape& input_shape,
- float* output_data, const RuntimeShape& output_shape) {
- SoftmaxParams params;
- // No params currently used for float LogSoftmax.
- LogSoftmax(params, input_shape, input_data, output_shape, output_data);
-}
-
// Although currently the name of this function says that it cannot handle
// values less than 1, in practice it can handle as low as 1/x_max, where
// x_max is the largest representable input. In other words, the output range
@@ -3047,22 +2706,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void LogSoftmax(const uint8* input_data, const RuntimeShape& input_shape,
- int32 input_multiplier, int32 input_left_shift,
- int32 reverse_scaling_divisor,
- int32 reverse_scaling_right_shift, int diff_min,
- uint8* output_data, const RuntimeShape& output_shape) {
- SoftmaxParams params;
- params.input_multiplier = input_multiplier;
- params.input_left_shift = input_left_shift;
- params.reverse_scaling_divisor = reverse_scaling_divisor;
- params.reverse_scaling_right_shift = reverse_scaling_right_shift;
- params.diff_min = diff_min;
- LogSoftmax(params, input_shape, input_data, output_shape, output_data);
-}
-
inline void Logistic(const RuntimeShape& input_shape, const float* input_data,
const RuntimeShape& output_shape, float* output_data) {
const int flat_size = MatchingFlatSize(input_shape, output_shape);
@@ -3124,20 +2767,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void Logistic(const uint8* input_data, const RuntimeShape& input_shape,
- int32 input_zero_point, int32 input_range_radius,
- int32 input_multiplier, int input_left_shift,
- uint8* output_data, const RuntimeShape& output_shape) {
- LogisticParams params;
- params.input_zero_point = input_zero_point;
- params.input_range_radius = input_range_radius;
- params.input_multiplier = input_multiplier;
- params.input_left_shift = input_left_shift;
- Logistic(params, input_shape, input_data, output_shape, output_data);
-}
-
inline void Logistic(const LogisticParams& params,
const RuntimeShape& input_shape, const int16* input_data,
const RuntimeShape& output_shape, int16* output_data) {
@@ -3157,15 +2786,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void Logistic(const RuntimeShape& input_shape, const int16* input_data,
- const RuntimeShape& output_shape, int16* output_data) {
- LogisticParams params;
- // No params currently needed by int16 Logistic.
- Logistic(params, input_shape, input_data, output_shape, output_data);
-}
-
inline void Tanh(const RuntimeShape& input_shape, const float* input_data,
const RuntimeShape& output_shape, float* output_data) {
const int flat_size = MatchingFlatSize(input_shape, output_shape);
@@ -3229,20 +2849,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void Tanh(const uint8* input_data, const RuntimeShape& input_shape,
- int32 input_zero_point, int32 input_range_radius,
- int32 input_multiplier, int input_left_shift,
- uint8* output_data, const RuntimeShape& output_shape) {
- TanhParams params;
- params.input_zero_point = input_zero_point;
- params.input_range_radius = input_range_radius;
- params.input_multiplier = input_multiplier;
- params.input_left_shift = input_left_shift;
- Tanh(params, input_shape, input_data, output_shape, output_data);
-}
-
inline void Tanh(const TanhParams& params, const RuntimeShape& input_shape,
const int16* input_data, const RuntimeShape& output_shape,
int16* output_data) {
@@ -3277,16 +2883,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void Tanh(const int16* input_data, const RuntimeShape& input_shape,
- int input_left_shift, int16* output_data,
- const RuntimeShape& output_shape) {
- TanhParams params;
- params.input_left_shift = input_left_shift;
- Tanh(params, input_shape, input_data, output_shape, output_data);
-}
-
inline void Dequantize(const tflite::DequantizationParams& op_params,
const RuntimeShape& input_shape, const uint8* input_data,
const RuntimeShape& output_shape, float* output_data) {
@@ -3301,19 +2897,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy Dims<4>.
-inline void Dequantize(const uint8* input_data, const Dims<4>& input_dims,
- int32 zero_point, double scale, float* output_data,
- const Dims<4>& output_dims) {
- tflite::DequantizationParams op_params;
- op_params.zero_point = zero_point;
- op_params.scale = scale;
-
- Dequantize(op_params, DimsToShape(input_dims), input_data,
- DimsToShape(output_dims), output_data);
-}
-
inline void FakeQuant(const tflite::FakeQuantParams& op_params,
const RuntimeShape& input_shape, const float* input_data,
const RuntimeShape& output_shape, float* output_data) {
@@ -3337,20 +2920,6 @@
output_data, flat_size);
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy Dims<4>.
-inline void FakeQuant(const float* input_data, const Dims<4>& input_dims,
- float rmin, float rmax, int num_bits, float* output_data,
- const Dims<4>& output_dims) {
- tflite::FakeQuantParams op_params;
- op_params.num_bits = num_bits;
- op_params.minmax.min = rmin;
- op_params.minmax.max = rmax;
-
- FakeQuant(op_params, DimsToShape(input_dims), input_data,
- DimsToShape(output_dims), output_data);
-}
-
template <typename SrcT, typename DstT>
inline void Cast(const RuntimeShape& input_shape, const SrcT* input_data,
const RuntimeShape& output_shape, DstT* output_data) {
@@ -3374,15 +2943,21 @@
template <typename T>
inline void Gather(const tflite::GatherParams& op_params,
- const RuntimeShape& input_shape, const T* input_data,
- const RuntimeShape& coords_shape, const int32* coords_data,
- const RuntimeShape& output_shape, T* output_data) {
- // Enable these checks when moving legacy ops to legacy_reference_ops.
- //
- // TFLITE_DCHECK_EQ(coords_shape.DimensionsCount(), 1);
+ const RuntimeShape& unextended_input_shape,
+ const T* input_data, const RuntimeShape& coords_shape,
+ const int32* coords_data,
+ const RuntimeShape& unextended_output_shape,
+ T* output_data) {
+ TFLITE_DCHECK_LE(unextended_input_shape.DimensionsCount(), 4);
+ TFLITE_DCHECK_LE(unextended_output_shape.DimensionsCount(), 4);
+ const RuntimeShape input_shape =
+ RuntimeShape::ExtendedShape(4, unextended_input_shape);
+ const RuntimeShape output_shape =
+ RuntimeShape::ExtendedShape(4, unextended_output_shape);
+
const int input_rank = op_params.input_rank;
const int gather_dimensions = output_shape.DimensionsCount();
- TFLITE_DCHECK_LE(input_shape.DimensionsCount(), gather_dimensions);
+ TFLITE_DCHECK_GE(input_shape.DimensionsCount(), gather_dimensions);
const int axis = gather_dimensions - input_rank;
TFLITE_DCHECK_LT(axis, gather_dimensions);
TFLITE_DCHECK_GE(axis, 0);
@@ -3404,23 +2979,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy Dims<4> version.
-// When moving legacy ops to legacy_reference_ops, replace content with looser
-// implementation.
-template <typename T>
-inline void Gather(const T* input_data, const Dims<4>& input_dims,
- int input_rank, const int32* coords_data,
- const Dims<4>& coords_dims, T* output_data,
- const Dims<4>& output_dims) {
- tflite::GatherParams op_params;
- op_params.input_rank = input_rank;
-
- Gather(op_params, DimsToShape(input_dims), input_data,
- DimsToShape(coords_dims), coords_data, DimsToShape(output_dims),
- output_data);
-}
-
template <typename T>
inline void ResizeBilinear(const tflite::ResizeBilinearParams& op_params,
const RuntimeShape& unextended_input_shape,
@@ -3750,58 +3308,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline uint32 LegacyReverseBits32(uint32 n) {
- n = ((n >> 1) & 0x55555555) | ((n & 0x55555555) << 1);
- n = ((n >> 2) & 0x33333333) | ((n & 0x33333333) << 2);
- n = ((n >> 4) & 0x0F0F0F0F) | ((n & 0x0F0F0F0F) << 4);
- return (((n & 0xFF) << 24) | ((n & 0xFF00) << 8) | ((n & 0xFF0000) >> 8) |
- ((n & 0xFF000000) >> 24));
-}
-
-inline void StridedSliceReverseIndices(tflite::StridedSliceParams* p) {
- TFLITE_CHECK_EQ(p->start_indices_count, p->stop_indices_count);
- TFLITE_CHECK_EQ(p->stop_indices_count, p->strides_count);
-
- std::reverse(p->start_indices, p->start_indices + p->start_indices_count);
- std::reverse(p->stop_indices, p->stop_indices + p->stop_indices_count);
- std::reverse(p->strides, p->strides + p->strides_count);
-
- p->begin_mask = LegacyReverseBits32(static_cast<uint32>(p->begin_mask)) >>
- (32 - p->start_indices_count);
- p->ellipsis_mask =
- LegacyReverseBits32(static_cast<uint32>(p->ellipsis_mask)) >>
- (32 - p->start_indices_count);
- p->end_mask = LegacyReverseBits32(static_cast<uint32>(p->end_mask)) >>
- (32 - p->start_indices_count);
- p->new_axis_mask =
- LegacyReverseBits32(static_cast<uint32>(p->new_axis_mask)) >>
- (32 - p->start_indices_count);
- p->shrink_axis_mask =
- LegacyReverseBits32(static_cast<uint32>(p->shrink_axis_mask)) >>
- (32 - p->start_indices_count);
-}
-
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-template <typename T>
-inline void StridedSlice(const T* input_data, const Dims<4>& input_dims,
- int begin_mask, int end_mask, int shrink_axis_mask,
- const std::vector<int>& start_indices,
- const std::vector<int>& stop_indices,
- const std::vector<int>& strides, T* output_data,
- const Dims<4>& output_dims) {
- TFLITE_DCHECK_EQ(start_indices.size(), 4);
- auto op_params = strided_slice::BuildStridedSliceParams(
- begin_mask, end_mask, shrink_axis_mask, start_indices, stop_indices,
- strides);
- StridedSliceReverseIndices(&op_params);
-
- StridedSlice(op_params, DimsToShape(input_dims), input_data,
- DimsToShape(output_dims), output_data);
-}
-
template <typename T>
inline void Slice(const tflite::SliceParams& op_params,
const RuntimeShape& input_shape, const T* input_data,
@@ -4067,22 +3573,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy Dims<4>.
-template <typename T>
-inline void Mean(const T* input_data, const Dims<4>& input_dims,
- const std::vector<int>& reduction_indices, T* output_data,
- const Dims<4>& output_dims) {
- tflite::MeanParams op_params;
- op_params.axis_count = reduction_indices.size();
- for (int i = 0; i < op_params.axis_count; ++i) {
- op_params.axis[i] = reduction_indices[op_params.axis_count - 1 - i];
- }
-
- Mean(op_params, DimsToShape(input_dims), input_data, DimsToShape(output_dims),
- output_data);
-}
-
// Computes the mean of elements across dimensions given in axis.
// It does so in two stages, first calculates the sum of elements along the axis
// then divides it by the number of element in axis for quantized values.
@@ -4340,20 +3830,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-template <typename T>
-void Transpose(const T* input, const Dims<4>& input_dims, T* output,
- const Dims<4>& output_dims, const int* permuted_axes) {
- TransposeParams params;
- params.perm_count = 4;
- for (int i = 0; i < 4; ++i) {
- params.perm[i] = 3 - permuted_axes[3 - i];
- }
- Transpose(params, DimsToShape(input_dims), input, DimsToShape(output_dims),
- output);
-}
-
inline void TransposeConv(
const ConvParams& params, const RuntimeShape& input_shape,
const float* input_data, const RuntimeShape& filter_shape,
@@ -4427,27 +3903,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void TransposeConv(const float* input_data, const Dims<4>& input_dims,
- const float* filter_data, const Dims<4>& filter_dims,
- int stride_width, int stride_height, int pad_width,
- int pad_height, float* output_data,
- const Dims<4>& output_dims, float* im2col_data,
- const Dims<4>& im2col_dims) {
- tflite::ConvParams op_params;
- // Padding type is ignored, but still set.
- op_params.padding_type = PaddingType::kSame;
- op_params.padding_values.width = pad_width;
- op_params.padding_values.height = pad_height;
- op_params.stride_width = stride_width;
- op_params.stride_height = stride_height;
-
- TransposeConv(op_params, DimsToShape(input_dims), input_data,
- DimsToShape(filter_dims), filter_data, DimsToShape(output_dims),
- output_data, DimsToShape(im2col_dims), im2col_data);
-}
-
template <typename T>
inline bool EqualFn(T lhs, T rhs) {
return lhs == rhs;
@@ -4501,19 +3956,6 @@
input2_data, output_shape, output_data);
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-template <typename T, ComparisonFn<T> F>
-inline void Comparison(const T* input1_data, const Dims<4>& input1_dims,
- const T* input2_data, const Dims<4>& input2_dims,
- bool* output_data, const Dims<4>& output_dims) {
- ComparisonParams op_params;
- // No parameters needed.
- ComparisonImpl<T, F>(op_params, DimsToShape(input1_dims), input1_data,
- DimsToShape(input2_dims), input2_data,
- DimsToShape(output_dims), output_data);
-}
-
template <typename T, ComparisonFn<int32> F>
inline void ComparisonWithScaling(
const ComparisonParams& op_params, const RuntimeShape& input1_shape,
@@ -4544,32 +3986,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-template <typename T, ComparisonFn<int32> F>
-inline void Comparison(int left_shift, const T* input1_data,
- const Dims<4>& input1_dims, int32 input1_offset,
- int32 input1_multiplier, int input1_shift,
- const T* input2_data, const Dims<4>& input2_dims,
- int32 input2_offset, int32 input2_multiplier,
- int input2_shift, bool* output_data,
- const Dims<4>& output_dims) {
- tflite::ComparisonParams op_params;
- op_params.left_shift = left_shift;
- op_params.input1_offset = input1_offset;
- op_params.input1_multiplier = input1_multiplier;
- // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
- op_params.input1_shift = kReverseShift * input1_shift;
- op_params.input2_offset = input2_offset;
- op_params.input2_multiplier = input2_multiplier;
- // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
- op_params.input2_shift = kReverseShift * input2_shift;
-
- ComparisonWithScaling<T, F>(op_params, DimsToShape(input1_dims), input1_data,
- DimsToShape(input2_dims), input2_data,
- DimsToShape(output_dims), output_data);
-}
-
template <typename T, ComparisonFn<T> F>
inline void BroadcastComparison4DSlowImpl(
const ComparisonParams& op_params,
@@ -4613,22 +4029,6 @@
output_shape, output_data);
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-template <typename T, ComparisonFn<T> F>
-inline void BroadcastComparison(const T* input1_data,
- const Dims<4>& input1_dims,
- const T* input2_data,
- const Dims<4>& input2_dims, bool* output_data,
- const Dims<4>& output_dims) {
- ComparisonParams op_params;
- // No parameters needed.
- BroadcastComparison4DSlowImpl<T, F>(op_params, DimsToShape(input1_dims),
- input1_data, DimsToShape(input2_dims),
- input2_data, DimsToShape(output_dims),
- output_data);
-}
-
template <typename T, ComparisonFn<int32> F>
inline void BroadcastComparison4DSlowWithScaling(
const ComparisonParams& op_params,
@@ -4679,80 +4079,7 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-template <typename T, ComparisonFn<int32> F>
-inline void BroadcastComparison(int left_shift, const T* input1_data,
- const Dims<4>& input1_dims, int32 input1_offset,
- int32 input1_multiplier, int input1_shift,
- const T* input2_data,
- const Dims<4>& input2_dims, int32 input2_offset,
- int32 input2_multiplier, int input2_shift,
- bool* output_data, const Dims<4>& output_dims) {
- ComparisonParams op_params;
-
- op_params.left_shift = left_shift;
- op_params.input1_offset = input1_offset;
- op_params.input1_multiplier = input1_multiplier;
- // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
- op_params.input1_shift = kReverseShift * input1_shift;
- op_params.input2_offset = input2_offset;
- op_params.input2_multiplier = input2_multiplier;
- // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
- op_params.input2_shift = kReverseShift * input2_shift;
-
- BroadcastComparison4DSlowWithScaling<T, F>(
- op_params, DimsToShape(input1_dims), input1_data,
- DimsToShape(input2_dims), input2_data, DimsToShape(output_dims),
- output_data);
-}
-
#define TFLITE_COMPARISON_OP(name) \
- template <typename T> \
- inline void name(const T* input1_data, const Dims<4>& input1_dims, \
- const T* input2_data, const Dims<4>& input2_dims, \
- bool* output_data, const Dims<4>& output_dims) { \
- gemmlowp::ScopedProfilingLabel label(#name); \
- Comparison<T, name##Fn>(input1_data, input1_dims, input2_data, \
- input2_dims, output_data, output_dims); \
- } \
- template <typename T> \
- inline void name( \
- int left_shift, const T* input1_data, const Dims<4>& input1_dims, \
- int32 input1_offset, int32 input1_multiplier, int input1_shift, \
- const T* input2_data, const Dims<4>& input2_dims, int32 input2_offset, \
- int32 input2_multiplier, int input2_shift, bool* output_data, \
- const Dims<4>& output_dims) { \
- gemmlowp::ScopedProfilingLabel label(#name "/8bit"); \
- Comparison<T, name##Fn>(left_shift, input1_data, input1_dims, \
- input1_offset, input1_multiplier, input1_shift, \
- input2_data, input2_dims, input2_offset, \
- input2_multiplier, input2_shift, output_data, \
- output_dims); \
- } \
- template <typename T> \
- inline void Broadcast##name( \
- const T* input1_data, const Dims<4>& input1_dims, const T* input2_data, \
- const Dims<4>& input2_dims, bool* output_data, \
- const Dims<4>& output_dims) { \
- gemmlowp::ScopedProfilingLabel label("Broadcast" #name); \
- BroadcastComparison<T, name##Fn>(input1_data, input1_dims, input2_data, \
- input2_dims, output_data, output_dims); \
- } \
- template <typename T> \
- inline void Broadcast##name( \
- int left_shift, const T* input1_data, const Dims<4>& input1_dims, \
- int32 input1_offset, int32 input1_multiplier, int input1_shift, \
- const T* input2_data, const Dims<4>& input2_dims, int32 input2_offset, \
- int32 input2_multiplier, int input2_shift, bool* output_data, \
- const Dims<4>& output_dims) { \
- gemmlowp::ScopedProfilingLabel label("Broadcast" #name "/8bit"); \
- BroadcastComparison<T, name##Fn>(left_shift, input1_data, input1_dims, \
- input1_offset, input1_multiplier, \
- input1_shift, input2_data, input2_dims, \
- input2_offset, input2_multiplier, \
- input2_shift, output_data, output_dims); \
- } \
inline void name(const ComparisonParams& op_params, \
const RuntimeShape& input1_shape, const float* input1_data, \
const RuntimeShape& input2_shape, const float* input2_data, \
@@ -4762,22 +4089,44 @@
input2_data, output_shape, output_data); \
} \
template <typename T> \
+ inline void name##NoScaling( \
+ const ComparisonParams& op_params, const RuntimeShape& input1_shape, \
+ const T* input1_data, const RuntimeShape& input2_shape, \
+ const T* input2_data, const RuntimeShape& output_shape, \
+ bool* output_data) { \
+ gemmlowp::ScopedProfilingLabel label(#name "NoScaling"); \
+ ComparisonImpl<T, name##Fn>(op_params, input1_shape, input1_data, \
+ input2_shape, input2_data, output_shape, \
+ output_data); \
+ } \
+ template <typename T> \
inline void name##WithScaling( \
const ComparisonParams& op_params, const RuntimeShape& input1_shape, \
const T* input1_data, const RuntimeShape& input2_shape, \
const T* input2_data, const RuntimeShape& output_shape, \
bool* output_data) { \
- gemmlowp::ScopedProfilingLabel label(#name "/8bit"); \
+ gemmlowp::ScopedProfilingLabel label(#name "WithScaling/8bit"); \
ComparisonWithScaling<T, name##Fn>(op_params, input1_shape, input1_data, \
input2_shape, input2_data, \
output_shape, output_data); \
} \
+ template <typename T> \
+ inline void Broadcast4DSlow##name##NoScaling( \
+ const ComparisonParams& op_params, const RuntimeShape& input1_shape, \
+ const T* input1_data, const RuntimeShape& input2_shape, \
+ const T* input2_data, const RuntimeShape& output_shape, \
+ bool* output_data) { \
+ gemmlowp::ScopedProfilingLabel label("Broadcast4DSlow" #name "NoScaling"); \
+ BroadcastComparison4DSlowImpl<T, name##Fn>( \
+ op_params, input1_shape, input1_data, input2_shape, input2_data, \
+ output_shape, output_data); \
+ } \
inline void Broadcast4DSlow##name( \
const ComparisonParams& op_params, const RuntimeShape& input1_shape, \
const float* input1_data, const RuntimeShape& input2_shape, \
const float* input2_data, const RuntimeShape& output_shape, \
bool* output_data) { \
- gemmlowp::ScopedProfilingLabel label("Broadcast" #name); \
+ gemmlowp::ScopedProfilingLabel label("Broadcast4DSlow" #name); \
BroadcastComparison4DSlow<name##Fn>(op_params, input1_shape, input1_data, \
input2_shape, input2_data, \
output_shape, output_data); \
@@ -4788,7 +4137,7 @@
const T* input1_data, const RuntimeShape& input2_shape, \
const T* input2_data, const RuntimeShape& output_shape, \
bool* output_data) { \
- gemmlowp::ScopedProfilingLabel label("Broadcast" #name "/8bit"); \
+ gemmlowp::ScopedProfilingLabel label("Broadcast4DSlow" #name "/8bit"); \
BroadcastComparison4DSlowWithScaling<T, name##Fn>( \
op_params, input1_shape, input1_data, input2_shape, input2_data, \
output_shape, output_data); \
@@ -4815,19 +4164,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-template <typename D, typename T>
-inline void Select(const D* input_condition_data,
- const Dims<4>& input_condition_dims, const T* input_x_data,
- const Dims<4>& input_x_dims, const T* input_y_data,
- const Dims<4>& input_y_dims, T* output_data,
- const Dims<4>& output_dims) {
- Select(DimsToShape(input_condition_dims), input_condition_data,
- DimsToShape(input_x_dims), input_x_data, DimsToShape(input_y_dims),
- input_y_data, DimsToShape(output_dims), output_data);
-}
-
template <typename D, typename T>
void RankOneSelect(const RuntimeShape& input_condition_shape,
const D* input_condition_data,
@@ -4849,20 +4185,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-template <typename D, typename T>
-inline void RankOneSelect(const D* input_condition_data,
- const Dims<4>& input_condition_dims,
- const T* input_x_data, const Dims<4>& input_x_dims,
- const T* input_y_data, const Dims<4>& input_y_dims,
- T* output_data, const Dims<4>& output_dims) {
- RankOneSelect(DimsToShape(input_condition_dims), input_condition_data,
- DimsToShape(input_x_dims), input_x_data,
- DimsToShape(input_y_dims), input_y_data,
- DimsToShape(output_dims), output_data);
-}
-
// For easy implementation, the indices is always a vector of size-4 vectors.
template <typename T, typename TI>
inline void SparseToDense(const std::vector<std::vector<TI>>& indices,
@@ -4904,16 +4226,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-template <typename T, typename TI>
-inline void SparseToDense(const std::vector<std::vector<TI>>& indices,
- const T* values, T default_value, T* output_data,
- const Dims<4>& output_dims, bool value_is_scalar) {
- SparseToDense(indices, values, default_value, value_is_scalar,
- DimsToShape(output_dims), output_data);
-}
-
template <typename T>
inline void Pow(const RuntimeShape& input1_shape, const T* input1_data,
const RuntimeShape& input2_shape, const T* input2_data,
diff --git a/tensorflow/contrib/lite/kernels/internal/reference/softmax.h b/tensorflow/contrib/lite/kernels/internal/reference/softmax.h
index 006174e..7d44296 100644
--- a/tensorflow/contrib/lite/kernels/internal/reference/softmax.h
+++ b/tensorflow/contrib/lite/kernels/internal/reference/softmax.h
@@ -57,16 +57,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy.
-inline void Softmax(const float* input_data, const RuntimeShape& input_shape,
- float beta, float* output_data,
- const RuntimeShape& output_shape) {
- SoftmaxParams params;
- params.beta = beta;
- Softmax(params, input_shape, input_data, output_shape, output_data);
-}
-
inline void Softmax(const SoftmaxParams& params,
const RuntimeShape& input_shape, const uint8* input_data,
const RuntimeShape& output_shape, uint8* output_data) {
@@ -151,19 +141,6 @@
}
}
-// TODO(b/80418076): Move to legacy ops file, update invocations.
-// Legacy
-inline void Softmax(const uint8* input_data, const RuntimeShape& input_shape,
- int32 input_beta_multiplier, int32 input_beta_left_shift,
- int diff_min, uint8* output_data,
- const RuntimeShape& output_shape) {
- SoftmaxParams params;
- params.input_multiplier = input_beta_multiplier;
- params.input_left_shift = input_beta_left_shift;
- params.diff_min = diff_min;
- Softmax(params, input_shape, input_data, output_shape, output_data);
-}
-
// Performs softmax along the input of size (input_size * batch_size).
inline void Softmax(const float* in, const int input_size, const int batch_size,
const float beta, float* out) {
diff --git a/tensorflow/contrib/lite/kernels/internal/softmax_quantized_test.cc b/tensorflow/contrib/lite/kernels/internal/softmax_quantized_test.cc
index ca94e77..831fb3c 100644
--- a/tensorflow/contrib/lite/kernels/internal/softmax_quantized_test.cc
+++ b/tensorflow/contrib/lite/kernels/internal/softmax_quantized_test.cc
@@ -43,11 +43,15 @@
// Reference data generated via Dequant of input into float, and then applying
// float Softmax.
- reference_ops::Dequantize(
- input_data, ToRuntimeDims(shape_common), input_offset, input_scale,
- reference_dequant_data.data(), ToRuntimeDims(shape_common));
- optimized_ops::Softmax(reference_dequant_data.data(), shape_common, beta,
- reference_output_float_data.data(), shape_common);
+ DequantizationParams dq_params;
+ dq_params.zero_point = input_offset;
+ dq_params.scale = input_scale;
+ reference_ops::Dequantize(dq_params, shape_common, input_data, shape_common,
+ reference_dequant_data.data());
+ SoftmaxParams sm_params;
+ sm_params.beta = beta;
+ optimized_ops::Softmax(sm_params, shape_common, reference_dequant_data.data(),
+ shape_common, reference_output_float_data.data());
// Work with quantized scaling for Softmax, under which 256 represents 1, but
// we limit this to 255.
for (int i = 0; i < ref_buffer_size; i++) {
@@ -116,12 +120,14 @@
const int diff_min = -tflite::CalculateInputRadius(kScaledDiffIntegerBits,
input_beta_left_shift);
- optimized_ops::Softmax(input_data, shape_common, input_beta_multiplier,
- input_beta_left_shift, diff_min,
- optimized_softmax_output.data(), shape_common);
- reference_ops::Softmax(input_data, shape_common, input_beta_multiplier,
- input_beta_left_shift, diff_min,
- reference_quant_softmax_output.data(), shape_common);
+ SoftmaxParams params;
+ params.input_multiplier = input_beta_multiplier;
+ params.input_left_shift = input_beta_left_shift;
+ params.diff_min = diff_min;
+ optimized_ops::Softmax(params, shape_common, input_data, shape_common,
+ optimized_softmax_output.data());
+ reference_ops::Softmax(params, shape_common, input_data, shape_common,
+ reference_quant_softmax_output.data());
CheckOutputData(optimized_softmax_output.data(),
reference_float_softmax_output.data(), shape_common,
diff --git a/tensorflow/contrib/lite/kernels/internal/tensor.h b/tensorflow/contrib/lite/kernels/internal/tensor.h
index 1310645..689cea0 100644
--- a/tensorflow/contrib/lite/kernels/internal/tensor.h
+++ b/tensorflow/contrib/lite/kernels/internal/tensor.h
@@ -37,10 +37,6 @@
: nullptr;
}
-inline Dims<4> GetTensorDims(std::vector<int32_t> data) {
- return GetTensorDims(data.data(), data.size());
-}
-
inline RuntimeShape GetTensorShape(std::vector<int32_t> data) {
return RuntimeShape(data.size(), data.data());
}
@@ -56,20 +52,20 @@
int num_tensors = tensor_list.size;
all_data_.reserve(num_tensors);
- all_dims_.reserve(num_tensors);
- all_dims_ptr_.reserve(num_tensors);
+ all_shape_.reserve(num_tensors);
+ all_shape_ptr_.reserve(num_tensors);
for (int i = 0; i < num_tensors; ++i) {
TfLiteTensor* t = &context.tensors[tensor_list.data[i]];
all_data_.push_back(GetTensorData<T>(t));
- all_dims_.push_back(GetTensorDims(t));
+ all_shape_.push_back(GetTensorShape(t));
}
// Taking the pointer from inside a std::vector is only OK if the vector is
- // never modified, so we populate all_dims in the previous loop and then we
+ // never modified, so we populate all_shape in the previous loop and then we
// are free to grab iterators here.
for (int i = 0; i < num_tensors; ++i) {
- all_dims_ptr_.push_back(&all_dims_[i]);
+ all_shape_ptr_.push_back(&all_shape_[i]);
}
}
// Return a pointer to the data pointers of all tensors in the list. For
@@ -78,16 +74,16 @@
// f[0][1] is the second element of the first tensor.
T* const* data() const { return all_data_.data(); }
- // Return a pointer the dim pointers of all tensors in the list. For
+ // Return a pointer the shape pointers of all tensors in the list. For
// example:
- // const Dims<4>* const* d = v.dims();
+ // const RuntimeShape* const* d = v.dims();
// dims[1] are the dimensions of the second tensor in the list.
- const Dims<4>* const* dims() const { return all_dims_ptr_.data(); }
+ const RuntimeShape* const* shapes() const { return all_shape_ptr_.data(); }
private:
std::vector<T*> all_data_;
- std::vector<Dims<4>> all_dims_;
- std::vector<Dims<4>*> all_dims_ptr_;
+ std::vector<RuntimeShape> all_shape_;
+ std::vector<RuntimeShape*> all_shape_ptr_;
};
// A list of quantized tensors in a format that can be used by kernels like
diff --git a/tensorflow/contrib/lite/kernels/internal/tensor_ctypes.h b/tensorflow/contrib/lite/kernels/internal/tensor_ctypes.h
index 77e22a0..9f5b33d 100644
--- a/tensorflow/contrib/lite/kernels/internal/tensor_ctypes.h
+++ b/tensorflow/contrib/lite/kernels/internal/tensor_ctypes.h
@@ -86,39 +86,6 @@
return tensor != nullptr ? tensor->data.b : nullptr;
}
-inline int RemapDim(int max_dimensions, int d) {
- return max_dimensions - d - 1;
-}
-
-// TODO(ahentz): the implementations in kernels/internal/ take a Dims<4> object
-// even if the original tensors were not 4D. We should consider rewriting them
-// to take a more generic 'shape' object.
-inline Dims<4> GetTensorDims(const int data[], const int size) {
- Dims<4> d;
- for (int i = 0; i < 4; ++i) {
- int src = size - i - 1;
- if (src >= 0) {
- d.sizes[i] = data[src];
- } else {
- d.sizes[i] = 1;
- }
- }
- d.strides[0] = 1;
- for (int i = 1; i < 4; i++) {
- d.strides[i] = d.strides[i - 1] * d.sizes[i - 1];
- }
- return d;
-}
-
-inline Dims<4> GetTensorDims(const TfLiteTensor* tensor) {
- if (tensor == nullptr) {
- return Dims<4>();
- }
-
- auto* dims = tensor->dims;
- return GetTensorDims(dims->data, dims->size);
-}
-
inline RuntimeShape GetTensorShape(const TfLiteTensor* tensor) {
if (tensor == nullptr) {
return RuntimeShape();
diff --git a/tensorflow/contrib/lite/kernels/internal/tensor_test.cc b/tensorflow/contrib/lite/kernels/internal/tensor_test.cc
index bf2068d..2ed73ba 100644
--- a/tensorflow/contrib/lite/kernels/internal/tensor_test.cc
+++ b/tensorflow/contrib/lite/kernels/internal/tensor_test.cc
@@ -21,28 +21,32 @@
using ::testing::ElementsAre;
-TEST(TensorTest, GetTensorDims4D) {
- Dims<4> d = GetTensorDims({2, 3, 4, 5});
- EXPECT_THAT(d.sizes, ElementsAre(5, 4, 3, 2));
- EXPECT_THAT(d.strides, ElementsAre(1, 5, 20, 60));
+TEST(TensorTest, GetTensorShape4D) {
+ RuntimeShape d = GetTensorShape({2, 3, 4, 5});
+ EXPECT_THAT(
+ std::vector<int32>(d.DimsData(), d.DimsData() + d.DimensionsCount()),
+ ElementsAre(2, 3, 4, 5));
}
-TEST(TensorTest, GetTensorDims3D) {
- Dims<4> d = GetTensorDims({3, 4, 5});
- EXPECT_THAT(d.sizes, ElementsAre(5, 4, 3, 1));
- EXPECT_THAT(d.strides, ElementsAre(1, 5, 20, 60));
+TEST(TensorTest, GetTensorShape3D) {
+ RuntimeShape d = GetTensorShape({3, 4, 5});
+ EXPECT_THAT(
+ std::vector<int32>(d.DimsData(), d.DimsData() + d.DimensionsCount()),
+ ElementsAre(3, 4, 5));
}
-TEST(TensorTest, GetTensorDims2D) {
- Dims<4> d = GetTensorDims({4, 5});
- EXPECT_THAT(d.sizes, ElementsAre(5, 4, 1, 1));
- EXPECT_THAT(d.strides, ElementsAre(1, 5, 20, 20));
+TEST(TensorTest, GetTensorShape2D) {
+ RuntimeShape d = GetTensorShape({4, 5});
+ EXPECT_THAT(
+ std::vector<int32>(d.DimsData(), d.DimsData() + d.DimensionsCount()),
+ ElementsAre(4, 5));
}
-TEST(TensorTest, GetTensorDims1D) {
- Dims<4> d = GetTensorDims({5});
- EXPECT_THAT(d.sizes, ElementsAre(5, 1, 1, 1));
- EXPECT_THAT(d.strides, ElementsAre(1, 5, 5, 5));
+TEST(TensorTest, GetTensorShape1D) {
+ RuntimeShape d = GetTensorShape({5});
+ EXPECT_THAT(
+ std::vector<int32>(d.DimsData(), d.DimsData() + d.DimensionsCount()),
+ ElementsAre(5));
}
} // namespace
diff --git a/tensorflow/contrib/lite/kernels/internal/types.h b/tensorflow/contrib/lite/kernels/internal/types.h
index a3a5994..b393477 100644
--- a/tensorflow/contrib/lite/kernels/internal/types.h
+++ b/tensorflow/contrib/lite/kernels/internal/types.h
@@ -875,6 +875,15 @@
int16 axis[4];
};
+struct PackParams {
+ int8 axis;
+ const int32* input_zeropoint;
+ const float* input_scale;
+ uint16 inputs_count;
+ int32 output_zeropoint;
+ float output_scale;
+};
+
struct PadParams {
int8 left_padding_count;
int32 left_padding[4];
@@ -975,6 +984,11 @@
int32 perm[4];
};
+struct UnpackParams {
+ uint16 num_split;
+ int16 axis;
+};
+
template <typename P>
inline void SetActivationParams(float min, float max, P* params) {
params->float_activation_min = min;
diff --git a/tensorflow/contrib/lite/kernels/log_softmax_test.cc b/tensorflow/contrib/lite/kernels/log_softmax_test.cc
index 9a8d35e..1acc966 100644
--- a/tensorflow/contrib/lite/kernels/log_softmax_test.cc
+++ b/tensorflow/contrib/lite/kernels/log_softmax_test.cc
@@ -91,8 +91,9 @@
std::unique_ptr<float[]> output_buffer(new float[input_size * batch_size]);
auto input_shape = RuntimeShape({batch_size, 1, 1, input_size});
- tflite::reference_ops::LogSoftmax(input_buffer, input_shape,
- output_buffer.get(), input_shape);
+ SoftmaxParams params;
+ tflite::reference_ops::LogSoftmax(params, input_shape, input_buffer,
+ input_shape, output_buffer.get());
std::vector<float> expected;
expected.insert(expected.end(), output_buffer.get(),
diff --git a/tensorflow/contrib/lite/kernels/pack.cc b/tensorflow/contrib/lite/kernels/pack.cc
index 4cb98fd..c368582 100644
--- a/tensorflow/contrib/lite/kernels/pack.cc
+++ b/tensorflow/contrib/lite/kernels/pack.cc
@@ -85,9 +85,12 @@
void PackImpl(TfLiteContext* context, TfLiteNode* node, TfLiteTensor* output,
int values_count, int axis) {
VectorOfTensors<T> all_inputs(*context, *node->inputs);
- reference_ops::Pack<T>(RemapDim(NumDimensions(output), axis),
- all_inputs.data(), all_inputs.dims(), values_count,
- GetTensorData<T>(output), GetTensorDims(output));
+ tflite::PackParams op_params;
+ op_params.axis = axis;
+ op_params.inputs_count = values_count;
+
+ reference_ops::Pack<T>(op_params, all_inputs.shapes(), all_inputs.data(),
+ GetTensorShape(output), GetTensorData<T>(output));
}
TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
diff --git a/tensorflow/contrib/lite/kernels/register.cc b/tensorflow/contrib/lite/kernels/register.cc
index 2f4b663..9402105 100644
--- a/tensorflow/contrib/lite/kernels/register.cc
+++ b/tensorflow/contrib/lite/kernels/register.cc
@@ -125,7 +125,7 @@
context->ReportError(
context,
"Regular TensorFlow ops are not supported by this interpreter. Make sure "
- "you invoke the Eager delegate before inference.");
+ "you invoke the Flex delegate before inference.");
return kTfLiteError;
}
@@ -136,13 +136,13 @@
const TfLiteRegistration* BuiltinOpResolver::FindOp(const char* op,
int version) const {
- // Return the NULL Op for all ops whose name start with "Eager", allowing
+ // Return the NULL Op for all ops whose name start with "Flex", allowing
// the interpreter to delegate their execution.
- if (IsEagerOp(op)) {
+ if (IsFlexOp(op)) {
static TfLiteRegistration null_op{
nullptr, nullptr, &UnsupportedTensorFlowOp,
nullptr, nullptr, BuiltinOperator_CUSTOM,
- "Eager", 1};
+ "Flex", 1};
return &null_op;
}
return MutableOpResolver::FindOp(op, version);
diff --git a/tensorflow/contrib/lite/kernels/select.cc b/tensorflow/contrib/lite/kernels/select.cc
index 3959502..4780a86 100644
--- a/tensorflow/contrib/lite/kernels/select.cc
+++ b/tensorflow/contrib/lite/kernels/select.cc
@@ -70,12 +70,12 @@
bool is_rank_one = !HaveSameShapes(input_condition, input_x);
-#define TF_LITE_SELECT(type, op) \
- reference_ops::op(GetTensorData<bool>(input_condition), \
- GetTensorDims(input_condition), \
- GetTensorData<type>(input_x), GetTensorDims(input_x), \
- GetTensorData<type>(input_y), GetTensorDims(input_y), \
- GetTensorData<type>(output), GetTensorDims(output));
+#define TF_LITE_SELECT(type, op) \
+ reference_ops::op(GetTensorShape(input_condition), \
+ GetTensorData<bool>(input_condition), \
+ GetTensorShape(input_x), GetTensorData<type>(input_x), \
+ GetTensorShape(input_y), GetTensorData<type>(input_y), \
+ GetTensorShape(output), GetTensorData<type>(output));
#define TF_LITE_SWITCH(type, op) \
switch (type) { \
diff --git a/tensorflow/contrib/lite/kernels/softmax_test.cc b/tensorflow/contrib/lite/kernels/softmax_test.cc
index 727822f..bd66980 100644
--- a/tensorflow/contrib/lite/kernels/softmax_test.cc
+++ b/tensorflow/contrib/lite/kernels/softmax_test.cc
@@ -93,8 +93,10 @@
std::unique_ptr<float[]> output_buffer(new float[input_size * batch_size]);
auto input_shape = RuntimeShape({batch_size, 1, 1, input_size});
- tflite::reference_ops::Softmax(input_buffer, input_shape, beta,
- output_buffer.get(), input_shape);
+ SoftmaxParams params;
+ params.beta = beta;
+ tflite::reference_ops::Softmax(params, input_shape, input_buffer, input_shape,
+ output_buffer.get());
std::vector<float> expected;
expected.insert(expected.end(), output_buffer.get(),
@@ -120,8 +122,10 @@
std::unique_ptr<float[]> output_buffer(new float[input_size * batch_size]);
auto input_shape = RuntimeShape({batch_size, 1, 1, input_size});
- tflite::reference_ops::Softmax(input_buffer, input_shape, beta,
- output_buffer.get(), input_shape);
+ SoftmaxParams params;
+ params.beta = beta;
+ tflite::reference_ops::Softmax(params, input_shape, input_buffer, input_shape,
+ output_buffer.get());
std::vector<float> expected;
expected.insert(expected.end(), output_buffer.get(),
diff --git a/tensorflow/contrib/lite/kernels/sparse_to_dense.cc b/tensorflow/contrib/lite/kernels/sparse_to_dense.cc
index 178568e..349fa0b 100644
--- a/tensorflow/contrib/lite/kernels/sparse_to_dense.cc
+++ b/tensorflow/contrib/lite/kernels/sparse_to_dense.cc
@@ -210,8 +210,9 @@
&indices_vector));
reference_ops::SparseToDense(indices_vector, GetTensorData<T>(values),
*GetTensorData<T>(default_value),
- GetTensorData<T>(output), GetTensorDims(output),
- value_is_scalar);
+ value_is_scalar, GetTensorShape(output),
+ GetTensorData<T>(output));
+
return kTfLiteOk;
}
diff --git a/tensorflow/contrib/lite/kernels/split.cc b/tensorflow/contrib/lite/kernels/split.cc
index 719e2dc..dab887b 100644
--- a/tensorflow/contrib/lite/kernels/split.cc
+++ b/tensorflow/contrib/lite/kernels/split.cc
@@ -109,25 +109,24 @@
if (axis_value < 0) {
axis_value += NumDimensions(op_context.input);
}
- axis_value = RemapDim(NumDimensions(op_context.input), axis_value);
// TODO(ahentz): Our usage of VectorOfTensors could be optimized by
// calculating it in Prepare, unless we defer shape calculation.
// TODO(ahentz): We can improve the optimized_ops version to handle other
// cases too.
-#define TF_LITE_SPLIT(scalar) \
- VectorOfTensors<scalar> all_outputs(*context, *node->outputs); \
- if (axis_value == NumDimensions(op_context.input)) { \
- optimized_ops::TensorFlowSplit<FusedActivationFunctionType::kNone, \
- scalar>( \
- GetTensorData<scalar>(op_context.input), \
- GetTensorDims(op_context.input), NumOutputs(node), all_outputs.data(), \
- all_outputs.dims()); \
- } else { \
- reference_ops::TensorFlowSplit<scalar>( \
- GetTensorData<scalar>(op_context.input), \
- GetTensorDims(op_context.input), axis_value, NumOutputs(node), \
- all_outputs.data(), all_outputs.dims()); \
+#define TF_LITE_SPLIT(scalar) \
+ VectorOfTensors<scalar> all_outputs(*context, *node->outputs); \
+ tflite::SplitParams op_params; \
+ op_params.num_split = NumOutputs(node); \
+ op_params.axis = axis_value; \
+ if (axis_value == 0) { \
+ optimized_ops::Split(op_params, GetTensorShape(op_context.input), \
+ GetTensorData<scalar>(op_context.input), \
+ all_outputs.shapes(), all_outputs.data()); \
+ } else { \
+ reference_ops::Split(op_params, GetTensorShape(op_context.input), \
+ GetTensorData<scalar>(op_context.input), \
+ all_outputs.shapes(), all_outputs.data()); \
}
switch (op_context.input->type) {
case kTfLiteFloat32: {
diff --git a/tensorflow/contrib/lite/kernels/strided_slice.cc b/tensorflow/contrib/lite/kernels/strided_slice.cc
index 87ffcc4..06b36dd 100644
--- a/tensorflow/contrib/lite/kernels/strided_slice.cc
+++ b/tensorflow/contrib/lite/kernels/strided_slice.cc
@@ -57,17 +57,6 @@
int dims;
};
-// Reverse order of bits in the mask to match the expected order in kernel
-inline int ReverseMaskBits(int mask, int num_dimensions) {
- int out = 0;
- for (int dim = 0; dim < num_dimensions; dim++) {
- out <<= 1;
- out += (mask & 1);
- mask >>= 1;
- }
- return out;
-}
-
// This Op only supports 1-4D cases and since we use the reference 4D
// implementation, the 1-3D tensors are mapped to 4D.
const int kMaxDim = 4;
@@ -198,30 +187,31 @@
std::vector<int32_t> stops;
std::vector<int32_t> strides;
- for (int idx = op_context.dims - 1; idx >= 0; --idx) {
- starts.emplace_back(GetTensorData<int32_t>(op_context.begin)[idx]);
- stops.emplace_back(GetTensorData<int32_t>(op_context.end)[idx]);
- strides.emplace_back(GetTensorData<int32_t>(op_context.strides)[idx]);
- }
-
for (int i = op_context.dims; i < kMaxDim; i++) {
starts.emplace_back(0);
stops.emplace_back(1);
strides.emplace_back(1);
}
- int begin_mask =
- ReverseMaskBits(op_context.params->begin_mask, op_context.dims);
- int end_mask = ReverseMaskBits(op_context.params->end_mask, op_context.dims);
- int shrink_axis_mask =
- ReverseMaskBits(op_context.params->shrink_axis_mask, op_context.dims);
+ for (int idx = 0; idx < op_context.dims; ++idx) {
+ starts.emplace_back(GetTensorData<int32_t>(op_context.begin)[idx]);
+ stops.emplace_back(GetTensorData<int32_t>(op_context.end)[idx]);
+ strides.emplace_back(GetTensorData<int32_t>(op_context.strides)[idx]);
+ }
-#define TF_LITE_STRIDED_SLICE(kernel_type, data_type) \
- kernel_type::StridedSlice( \
- GetTensorData<data_type>(op_context.input), \
- GetTensorDims(op_context.input), begin_mask, end_mask, shrink_axis_mask, \
- starts, stops, strides, GetTensorData<data_type>(op_context.output), \
- GetTensorDims(op_context.output))
+ int begin_mask = op_context.params->begin_mask << (4 - op_context.dims);
+ int end_mask = op_context.params->end_mask << (4 - op_context.dims);
+ int shrink_axis_mask = op_context.params->shrink_axis_mask
+ << (4 - op_context.dims);
+ TF_LITE_ENSURE_EQ(context, starts.size(), 4);
+ auto op_params = ::tflite::strided_slice::BuildStridedSliceParams(
+ begin_mask, end_mask, shrink_axis_mask, starts, stops, strides);
+
+#define TF_LITE_STRIDED_SLICE(kernel_type, data_type) \
+ kernel_type::StridedSlice(op_params, GetTensorShape(op_context.input), \
+ GetTensorData<data_type>(op_context.input), \
+ GetTensorShape(op_context.output), \
+ GetTensorData<data_type>(op_context.output))
switch (op_context.input->type) {
case kTfLiteFloat32:
diff --git a/tensorflow/contrib/lite/kernels/test_util.cc b/tensorflow/contrib/lite/kernels/test_util.cc
index 0fdb0a3..05a7c23 100644
--- a/tensorflow/contrib/lite/kernels/test_util.cc
+++ b/tensorflow/contrib/lite/kernels/test_util.cc
@@ -122,7 +122,7 @@
CHECK(interpreter_->AllocateTensors() == kTfLiteOk)
<< "Cannot allocate tensors";
- interpreter_->ResetVariableTensorsToZero();
+ interpreter_->ResetVariableTensors();
}
void SingleOpModel::Invoke() { CHECK(interpreter_->Invoke() == kTfLiteOk); }
diff --git a/tensorflow/contrib/lite/kernels/transpose.cc b/tensorflow/contrib/lite/kernels/transpose.cc
index 9535996..e42a304 100644
--- a/tensorflow/contrib/lite/kernels/transpose.cc
+++ b/tensorflow/contrib/lite/kernels/transpose.cc
@@ -92,26 +92,19 @@
TF_LITE_ENSURE_OK(context, ResizeOutputTensor(context, &op_context));
}
- // Reverse the permuted axes and convert to 4D due to the way Dims are
- // constructed in GetTensorDims.
const int* perm_data = GetTensorData<int32_t>(op_context.perm);
const int size = op_context.perm->dims->data[0];
- const int kOutputDimensionNum = 4;
- int reversed_perm[kOutputDimensionNum];
-
- for (int output_k = 0, input_k = size - 1; output_k < size;
- ++output_k, --input_k) {
- reversed_perm[output_k] = size - perm_data[input_k] - 1;
- }
- for (int k = size; k < kOutputDimensionNum; ++k) {
- reversed_perm[k] = k;
+ TransposeParams params;
+ params.perm_count = size;
+ for (int i = 0; i < size; ++i) {
+ params.perm[i] = perm_data[i];
}
#define TF_LITE_TRANSPOSE(type, scalar) \
- type::Transpose(GetTensorData<scalar>(op_context.input), \
- GetTensorDims(op_context.input), \
- GetTensorData<scalar>(op_context.output), \
- GetTensorDims(op_context.output), reversed_perm)
+ type::Transpose(params, GetTensorShape(op_context.input), \
+ GetTensorData<scalar>(op_context.input), \
+ GetTensorShape(op_context.output), \
+ GetTensorData<scalar>(op_context.output))
switch (op_context.input->type) {
case kTfLiteFloat32:
diff --git a/tensorflow/contrib/lite/kernels/transpose_test.cc b/tensorflow/contrib/lite/kernels/transpose_test.cc
index 337bc14..79ef0a7 100644
--- a/tensorflow/contrib/lite/kernels/transpose_test.cc
+++ b/tensorflow/contrib/lite/kernels/transpose_test.cc
@@ -51,21 +51,21 @@
reversed_perms[k] = k;
}
- // Make input and output dims (i.e. reversed shape and dest_shape).
- Dims<4> input_dims = GetTensorDims(shape);
- Dims<4> output_dims;
- for (int i = 0; i < 4; i++) {
- output_dims.sizes[i] = input_dims.sizes[reversed_perms[i]];
- }
- output_dims.strides[0] = 1;
- for (int k = 1; k < 4; k++) {
- output_dims.strides[k] =
- output_dims.strides[k - 1] * output_dims.sizes[k - 1];
+ // Make input and output shapes.
+ const RuntimeShape input_shape = GetTensorShape(shape);
+ RuntimeShape output_shape(perms.size());
+ for (int i = 0; i < perms.size(); i++) {
+ output_shape.SetDim(i, input_shape.Dims(perms[i]));
}
- reference_ops::Transpose<float>(input.data(), input_dims,
- input_transposed->data(), output_dims,
- reversed_perms);
+ TransposeParams params;
+ params.perm_count = perms.size();
+ for (int i = 0; i < perms.size(); ++i) {
+ params.perm[i] = perms[i];
+ }
+
+ reference_ops::Transpose<float>(params, input_shape, input.data(),
+ output_shape, input_transposed->data());
}
TEST(TransposeTest, TestRefOps1D) {
diff --git a/tensorflow/contrib/lite/kernels/unpack.cc b/tensorflow/contrib/lite/kernels/unpack.cc
index 9ff06f8..a7d3a9b 100644
--- a/tensorflow/contrib/lite/kernels/unpack.cc
+++ b/tensorflow/contrib/lite/kernels/unpack.cc
@@ -88,10 +88,13 @@
template <typename T>
void UnpackImpl(TfLiteContext* context, TfLiteNode* node,
const TfLiteTensor* input, int output_count, int axis) {
+ tflite::UnpackParams op_params;
+ op_params.axis = axis;
+ op_params.num_split = output_count;
VectorOfTensors<T> all_outputs(*context, *node->outputs);
- reference_ops::Unpack<T>(axis, GetTensorData<T>(input), GetTensorDims(input),
- NumDimensions(input), output_count,
- all_outputs.data(), **all_outputs.dims());
+ reference_ops::Unpack<T>(op_params, GetTensorShape(input),
+ GetTensorData<T>(input), **all_outputs.shapes(),
+ all_outputs.data());
}
TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
diff --git a/tensorflow/contrib/lite/model.cc b/tensorflow/contrib/lite/model.cc
index ea2817b..d50c345 100644
--- a/tensorflow/contrib/lite/model.cc
+++ b/tensorflow/contrib/lite/model.cc
@@ -27,8 +27,8 @@
#ifndef TFLITE_MCU
#include "tensorflow/contrib/lite/nnapi_delegate.h"
#endif
-#if defined(TFLITE_EXTENDED)
-#include "tensorflow/contrib/lite/delegates/eager/delegate.h"
+#if defined(TFLITE_FLEX)
+#include "tensorflow/contrib/lite/delegates/flex/delegate.h"
#endif
#include "tensorflow/contrib/lite/version.h"
@@ -450,8 +450,8 @@
}
(**interpreter).SetVariables(std::move(variables));
-#if defined(TFLITE_EXTENDED)
- if (auto delegate = EagerDelegate::Create()) {
+#if defined(TFLITE_FLEX)
+ if (auto delegate = FlexDelegate::Create()) {
(**interpreter)
.ModifyGraphWithDelegate(std::move(delegate),
/*allow_dynamic_tensors=*/true);
diff --git a/tensorflow/contrib/lite/models/smartreply/demo/app/src/main/BUILD b/tensorflow/contrib/lite/models/smartreply/demo/app/src/main/BUILD
index f18a2ca..2e5033d 100644
--- a/tensorflow/contrib/lite/models/smartreply/demo/app/src/main/BUILD
+++ b/tensorflow/contrib/lite/models/smartreply/demo/app/src/main/BUILD
@@ -20,6 +20,7 @@
android_binary(
name = "SmartReplyDemo",
srcs = glob(["java/**/*.java"]),
+ aapt_version = "aapt",
assets = [":assets"],
assets_dir = "",
custom_package = "com.example.android.smartreply",
diff --git a/tensorflow/contrib/lite/python/BUILD b/tensorflow/contrib/lite/python/BUILD
index 57e1290..916788f 100644
--- a/tensorflow/contrib/lite/python/BUILD
+++ b/tensorflow/contrib/lite/python/BUILD
@@ -144,7 +144,7 @@
name = "convert_saved_model",
srcs = ["convert_saved_model.py"],
srcs_version = "PY2AND3",
- visibility = ["//visibility:public"],
+ visibility = ["//tensorflow/contrib/lite:__subpackages__"],
deps = [
":convert",
"//tensorflow/contrib/saved_model:saved_model_py",
diff --git a/tensorflow/contrib/lite/python/convert.py b/tensorflow/contrib/lite/python/convert.py
index 1f48a82..613a153 100644
--- a/tensorflow/contrib/lite/python/convert.py
+++ b/tensorflow/contrib/lite/python/convert.py
@@ -67,12 +67,12 @@
# Convert model using TOCO such that only unsupported operations are
# represented as TensorFlow ops.
# WARNING: Experimental interface, subject to change.
- TOCO_EXTENDED = "TOCO_EXTENDED"
+ TOCO_FLEX = "TOCO_FLEX"
# Convert model using TOCO such that all operations are represented as
# TensorFlow ops.
# WARNING: Experimental interface, subject to change.
- TOCO_EXTENDED_ALL = "TOCO_EXTENDED_ALL"
+ TOCO_FLEX_ALL = "TOCO_FLEX_ALL"
def __str__(self):
return self.value
@@ -240,11 +240,11 @@
if dump_graphviz_dir:
toco.dump_graphviz_dir = dump_graphviz_dir
toco.dump_graphviz_include_video = dump_graphviz_video
- if converter_mode == ConverterMode.TOCO_EXTENDED:
- toco.allow_eager_ops = True
- elif converter_mode == ConverterMode.TOCO_EXTENDED_ALL:
- toco.allow_eager_ops = True
- toco.force_eager_ops = True
+ if converter_mode == ConverterMode.TOCO_FLEX:
+ toco.allow_flex_ops = True
+ elif converter_mode == ConverterMode.TOCO_FLEX_ALL:
+ toco.allow_flex_ops = True
+ toco.force_flex_ops = True
model = _model_flags_pb2.ModelFlags()
model.change_concat_input_ranges = change_concat_input_ranges
@@ -343,13 +343,14 @@
return data
-@deprecation.deprecated(None, "Use `lite.TocoConverter` instead.")
+@deprecation.deprecated(None, "Use `lite.TFLiteConverter` instead.")
def toco_convert(input_data, input_tensors, output_tensors, *args, **kwargs):
- """"Convert a model using TOCO.
+ """Convert a model using TOCO.
Typically this function is used to convert from TensorFlow GraphDef to TFLite.
Conversion can be customized by providing arguments that are forwarded to
- `build_toco_convert_protos` (see documentation for details).
+ `build_toco_convert_protos` (see documentation for details). This function has
+ been deprecated. Please use `lite.TFLiteConverter` instead.
Args:
input_data: Input data (i.e. often `sess.graph_def`),
diff --git a/tensorflow/contrib/lite/python/convert_saved_model.py b/tensorflow/contrib/lite/python/convert_saved_model.py
index 1553464..d18b60d 100644
--- a/tensorflow/contrib/lite/python/convert_saved_model.py
+++ b/tensorflow/contrib/lite/python/convert_saved_model.py
@@ -44,7 +44,7 @@
dtype)
-def _get_meta_graph_def(saved_model_dir, tag_set):
+def get_meta_graph_def(saved_model_dir, tag_set):
"""Validate saved_model and extract MetaGraphDef.
Args:
@@ -61,7 +61,7 @@
return loader.load(sess, tag_set, saved_model_dir)
-def _get_signature_def(meta_graph, signature_key):
+def get_signature_def(meta_graph, signature_key):
"""Get the signature def from meta_graph with given signature_key.
Args:
@@ -86,7 +86,7 @@
return signature_def_map[signature_key]
-def _get_inputs_outputs(signature_def):
+def get_inputs_outputs(signature_def):
"""Get inputs and outputs from SignatureDef.
Args:
@@ -236,9 +236,9 @@
input_arrays or output_arrays are not valid.
"""
# Read SignatureDef.
- meta_graph = _get_meta_graph_def(saved_model_dir, tag_set)
- signature_def = _get_signature_def(meta_graph, signature_key)
- inputs, outputs = _get_inputs_outputs(signature_def)
+ meta_graph = get_meta_graph_def(saved_model_dir, tag_set)
+ signature_def = get_signature_def(meta_graph, signature_key)
+ inputs, outputs = get_inputs_outputs(signature_def)
# Check SavedModel for assets directory.
collection_def = meta_graph.collection_def
diff --git a/tensorflow/contrib/lite/python/interpreter.py b/tensorflow/contrib/lite/python/interpreter.py
index 1be61fe..5700bf7 100644
--- a/tensorflow/contrib/lite/python/interpreter.py
+++ b/tensorflow/contrib/lite/python/interpreter.py
@@ -253,5 +253,5 @@
self._ensure_safe()
self._interpreter.Invoke()
- def reset_all_variables_to_zero(self):
- return self._interpreter.ResetVariableTensorsToZero()
+ def reset_all_variables(self):
+ return self._interpreter.ResetVariableTensors()
diff --git a/tensorflow/contrib/lite/python/interpreter_wrapper/interpreter_wrapper.cc b/tensorflow/contrib/lite/python/interpreter_wrapper/interpreter_wrapper.cc
index 9ab05f3..418f19a 100644
--- a/tensorflow/contrib/lite/python/interpreter_wrapper/interpreter_wrapper.cc
+++ b/tensorflow/contrib/lite/python/interpreter_wrapper/interpreter_wrapper.cc
@@ -466,9 +466,9 @@
error_msg);
}
-PyObject* InterpreterWrapper::ResetVariableTensorsToZero() {
+PyObject* InterpreterWrapper::ResetVariableTensors() {
TFLITE_PY_ENSURE_VALID_INTERPRETER();
- TFLITE_PY_CHECK(interpreter_->ResetVariableTensorsToZero());
+ TFLITE_PY_CHECK(interpreter_->ResetVariableTensors());
Py_RETURN_NONE;
}
diff --git a/tensorflow/contrib/lite/python/interpreter_wrapper/interpreter_wrapper.h b/tensorflow/contrib/lite/python/interpreter_wrapper/interpreter_wrapper.h
index 641dd93..f5ca81e 100644
--- a/tensorflow/contrib/lite/python/interpreter_wrapper/interpreter_wrapper.h
+++ b/tensorflow/contrib/lite/python/interpreter_wrapper/interpreter_wrapper.h
@@ -65,7 +65,7 @@
PyObject* TensorQuantization(int i) const;
PyObject* SetTensor(int i, PyObject* value);
PyObject* GetTensor(int i) const;
- PyObject* ResetVariableTensorsToZero();
+ PyObject* ResetVariableTensors();
// Returns a reference to tensor index i as a numpy array. The base_object
// should be the interpreter object providing the memory.
diff --git a/tensorflow/contrib/lite/python/lite.py b/tensorflow/contrib/lite/python/lite.py
index 2be2445..09365f1 100644
--- a/tensorflow/contrib/lite/python/lite.py
+++ b/tensorflow/contrib/lite/python/lite.py
@@ -17,6 +17,7 @@
EXPERIMENTAL: APIs here are unstable and likely to change without notice.
@@TocoConverter
+@@TFLiteConverter
@@toco_convert
@@toco_convert_protos
@@Interpreter
@@ -62,9 +63,10 @@
from tensorflow.python.lib.io import file_io as _file_io
from tensorflow.python.saved_model import signature_constants as _signature_constants
from tensorflow.python.saved_model import tag_constants as _tag_constants
+from tensorflow.python.util import deprecation as _deprecation
-class TocoConverter(object):
+class TFLiteConverter(object):
"""Convert a TensorFlow model into `output_format` using TOCO.
This is used to convert from a TensorFlow GraphDef or SavedModel into either a
@@ -121,22 +123,22 @@
```python
# Converting a GraphDef from session.
- converter = lite.TocoConverter.from_session(sess, in_tensors, out_tensors)
+ converter = lite.TFLiteConverter.from_session(sess, in_tensors, out_tensors)
tflite_model = converter.convert()
open("converted_model.tflite", "wb").write(tflite_model)
# Converting a GraphDef from file.
- converter = lite.TocoConverter.from_frozen_graph(
+ converter = lite.TFLiteConverter.from_frozen_graph(
graph_def_file, input_arrays, output_arrays)
tflite_model = converter.convert()
open("converted_model.tflite", "wb").write(tflite_model)
# Converting a SavedModel.
- converter = lite.TocoConverter.from_saved_model(saved_model_dir)
+ converter = lite.TFLiteConverter.from_saved_model(saved_model_dir)
tflite_model = converter.convert()
# Converting a tf.keras model.
- converter = lite.TocoConverter.from_keras_model_file(keras_model)
+ converter = lite.TFLiteConverter.from_keras_model_file(keras_model)
tflite_model = converter.convert()
```
"""
@@ -147,10 +149,9 @@
output_tensors,
input_arrays_with_shape=None,
output_arrays=None):
- """Constructor for TocoConverter.
+ """Constructor for TFLiteConverter.
Args:
-
graph_def: Frozen TensorFlow GraphDef.
input_tensors: List of input tensors. Type and shape are computed using
`foo.get_shape()` and `foo.dtype`.
@@ -158,8 +159,8 @@
input_arrays_with_shape: Tuple of strings representing input tensor names
and list of integers representing input shapes
(e.g., [("foo" : [1, 16, 16, 3])]). Use only when graph cannot be loaded
- into TensorFlow and when `input_tensors` and `output_tensors` are None.
- (default None)
+ into TensorFlow and when `input_tensors` and `output_tensors` are
+ None. (default None)
output_arrays: List of output tensors to freeze graph with. Use only when
graph cannot be loaded into TensorFlow and when `input_tensors` and
`output_tensors` are None. (default None)
@@ -195,7 +196,7 @@
@classmethod
def from_session(cls, sess, input_tensors, output_tensors):
- """Creates a TocoConverter class from a TensorFlow Session.
+ """Creates a TFLiteConverter class from a TensorFlow Session.
Args:
sess: TensorFlow Session.
@@ -204,7 +205,7 @@
output_tensors: List of output tensors (only .name is used from this).
Returns:
- TocoConverter class.
+ TFLiteConverter class.
"""
graph_def = _freeze_graph(sess, output_tensors)
return cls(graph_def, input_tensors, output_tensors)
@@ -215,7 +216,7 @@
input_arrays,
output_arrays,
input_shapes=None):
- """Creates a TocoConverter class from a file containing a frozen GraphDef.
+ """Creates a TFLiteConverter class from a file containing a frozen GraphDef.
Args:
graph_def_file: Full filepath of file containing frozen GraphDef.
@@ -224,10 +225,10 @@
input_shapes: Dict of strings representing input tensor names to list of
integers representing input shapes (e.g., {"foo" : [1, 16, 16, 3]}).
Automatically determined when input shapes is None (e.g., {"foo" :
- None}). (default None)
+ None}). (default None)
Returns:
- TocoConverter class.
+ TFLiteConverter class.
Raises:
IOError:
@@ -310,7 +311,7 @@
output_arrays=None,
tag_set=None,
signature_key=None):
- """Creates a TocoConverter class from a SavedModel.
+ """Creates a TFLiteConverter class from a SavedModel.
Args:
saved_model_dir: SavedModel directory to convert.
@@ -319,7 +320,7 @@
input_shapes: Dict of strings representing input tensor names to list of
integers representing input shapes (e.g., {"foo" : [1, 16, 16, 3]}).
Automatically determined when input shapes is None (e.g., {"foo" :
- None}). (default None)
+ None}). (default None)
output_arrays: List of output tensors to freeze graph with. Uses output
arrays from SignatureDef when none are provided. (default None)
tag_set: Set of tags identifying the MetaGraphDef within the SavedModel to
@@ -328,7 +329,7 @@
(default DEFAULT_SERVING_SIGNATURE_DEF_KEY)
Returns:
- TocoConverter class.
+ TFLiteConverter class.
"""
if tag_set is None:
tag_set = set([_tag_constants.SERVING])
@@ -346,7 +347,7 @@
input_arrays=None,
input_shapes=None,
output_arrays=None):
- """Creates a TocoConverter class from a tf.keras model file.
+ """Creates a TFLiteConverter class from a tf.keras model file.
Args:
model_file: Full filepath of HDF5 file containing the tf.keras model.
@@ -355,12 +356,12 @@
input_shapes: Dict of strings representing input tensor names to list of
integers representing input shapes (e.g., {"foo" : [1, 16, 16, 3]}).
Automatically determined when input shapes is None (e.g., {"foo" :
- None}). (default None)
+ None}). (default None)
output_arrays: List of output tensors to freeze graph with. Uses output
arrays from SignatureDef when none are provided. (default None)
Returns:
- TocoConverter class.
+ TFLiteConverter class.
"""
_keras.backend.clear_session()
_keras.backend.set_learning_phase(False)
@@ -502,6 +503,59 @@
tensor.set_shape(shape)
+class TocoConverter(object):
+ """Convert a TensorFlow model into `output_format` using TOCO.
+
+ This class has been deprecated. Please use `lite.TFLiteConverter` instead.
+ """
+
+ @classmethod
+ @_deprecation.deprecated(None,
+ "Use `lite.TFLiteConverter.from_session` instead.")
+ def from_session(cls, sess, input_tensors, output_tensors):
+ """Creates a TocoConverter class from a TensorFlow Session."""
+ return TFLiteConverter.from_session(sess, input_tensors, output_tensors)
+
+ @classmethod
+ @_deprecation.deprecated(
+ None, "Use `lite.TFLiteConverter.from_frozen_graph` instead.")
+ def from_frozen_graph(cls,
+ graph_def_file,
+ input_arrays,
+ output_arrays,
+ input_shapes=None):
+ """Creates a TocoConverter class from a file containing a frozen graph."""
+ return TFLiteConverter.from_frozen_graph(graph_def_file, input_arrays,
+ output_arrays, input_shapes)
+
+ @classmethod
+ @_deprecation.deprecated(
+ None, "Use `lite.TFLiteConverter.from_saved_model` instead.")
+ def from_saved_model(cls,
+ saved_model_dir,
+ input_arrays=None,
+ input_shapes=None,
+ output_arrays=None,
+ tag_set=None,
+ signature_key=None):
+ """Creates a TocoConverter class from a SavedModel."""
+ return TFLiteConverter.from_saved_model(saved_model_dir, input_arrays,
+ input_shapes, output_arrays,
+ tag_set, signature_key)
+
+ @classmethod
+ @_deprecation.deprecated(
+ None, "Use `lite.TFLiteConverter.from_keras_model_file` instead.")
+ def from_keras_model_file(cls,
+ model_file,
+ input_arrays=None,
+ input_shapes=None,
+ output_arrays=None):
+ """Creates a TocoConverter class from a tf.keras model file."""
+ return TFLiteConverter.from_keras_model_file(model_file, input_arrays,
+ input_shapes, output_arrays)
+
+
def _is_frozen_graph(sess):
"""Determines if the graph is frozen.
diff --git a/tensorflow/contrib/lite/python/lite_test.py b/tensorflow/contrib/lite/python/lite_test.py
index f112ed5..d243a49 100644
--- a/tensorflow/contrib/lite/python/lite_test.py
+++ b/tensorflow/contrib/lite/python/lite_test.py
@@ -50,18 +50,18 @@
# `output_arrays` is not defined.
with self.assertRaises(ValueError) as error:
- lite.TocoConverter(
+ lite.TFLiteConverter(
None, None, [], input_arrays_with_shape=[('input', [3, 9])])
self.assertEqual(message, str(error.exception))
# `input_arrays_with_shape` is not defined.
with self.assertRaises(ValueError) as error:
- lite.TocoConverter(None, [], None, output_arrays=['output'])
+ lite.TFLiteConverter(None, [], None, output_arrays=['output'])
self.assertEqual(message, str(error.exception))
# Tests valid constructors using a dummy value for the GraphDef.
def testValidConstructor(self):
- converter = lite.TocoConverter(
+ converter = lite.TFLiteConverter(
None,
None,
None,
@@ -76,7 +76,7 @@
'The batch size cannot be set for this model. Please use '
'input_shapes parameter.', str(error.exception))
- converter = lite.TocoConverter(None, ['input_tensor'], ['output_tensor'])
+ converter = lite.TFLiteConverter(None, ['input_tensor'], ['output_tensor'])
self.assertTrue(converter._has_valid_tensors())
@@ -89,7 +89,8 @@
sess = session.Session()
# Convert model and ensure model is not None.
- converter = lite.TocoConverter.from_session(sess, [in_tensor], [out_tensor])
+ converter = lite.TFLiteConverter.from_session(sess, [in_tensor],
+ [out_tensor])
tflite_model = converter.convert()
self.assertTrue(tflite_model)
@@ -121,7 +122,7 @@
sess = session.Session()
# Convert model and ensure model is not None.
- converter = lite.TocoConverter.from_session(
+ converter = lite.TFLiteConverter.from_session(
sess, [in_tensor_1, in_tensor_2], [out_tensor])
converter.inference_type = lite_constants.QUANTIZED_UINT8
converter.quantized_input_stats = {
@@ -166,7 +167,7 @@
sess = session.Session()
# Convert model and ensure model is not None.
- converter = lite.TocoConverter.from_session(
+ converter = lite.TFLiteConverter.from_session(
sess, [in_tensor_1, in_tensor_2], [out_tensor])
converter.inference_type = lite_constants.QUANTIZED_UINT8
converter.quantized_input_stats = {'inputA': (0., 1.)} # mean, std_dev
@@ -182,7 +183,8 @@
sess = session.Session()
# Test invalid shape. None after 1st dimension.
- converter = lite.TocoConverter.from_session(sess, [in_tensor], [out_tensor])
+ converter = lite.TFLiteConverter.from_session(sess, [in_tensor],
+ [out_tensor])
with self.assertRaises(ValueError) as error:
converter.convert()
self.assertEqual('Provide an input shape for input array \'Placeholder\'.',
@@ -195,7 +197,8 @@
sess = session.Session()
# Test invalid shape. None after 1st dimension.
- converter = lite.TocoConverter.from_session(sess, [in_tensor], [out_tensor])
+ converter = lite.TFLiteConverter.from_session(sess, [in_tensor],
+ [out_tensor])
with self.assertRaises(ValueError) as error:
converter.convert()
self.assertEqual(
@@ -210,7 +213,8 @@
sess = session.Session()
# Convert model and ensure model is not None.
- converter = lite.TocoConverter.from_session(sess, [in_tensor], [out_tensor])
+ converter = lite.TFLiteConverter.from_session(sess, [in_tensor],
+ [out_tensor])
tflite_model = converter.convert()
self.assertTrue(tflite_model)
@@ -242,7 +246,8 @@
sess.run(_global_variables_initializer())
# Convert model and ensure model is not None.
- converter = lite.TocoConverter.from_session(sess, [in_tensor], [out_tensor])
+ converter = lite.TFLiteConverter.from_session(sess, [in_tensor],
+ [out_tensor])
tflite_model = converter.convert()
self.assertTrue(tflite_model)
@@ -272,7 +277,8 @@
sess = session.Session()
# Convert model and ensure model is not None.
- converter = lite.TocoConverter.from_session(sess, [in_tensor], [out_tensor])
+ converter = lite.TFLiteConverter.from_session(sess, [in_tensor],
+ [out_tensor])
converter.output_format = lite_constants.GRAPHVIZ_DOT
graphviz_output = converter.convert()
self.assertTrue(graphviz_output)
@@ -285,7 +291,8 @@
sess = session.Session()
# Convert model and ensure model is not None.
- converter = lite.TocoConverter.from_session(sess, [in_tensor], [out_tensor])
+ converter = lite.TFLiteConverter.from_session(sess, [in_tensor],
+ [out_tensor])
graphviz_dir = self.get_temp_dir()
converter.dump_graphviz_dir = graphviz_dir
tflite_model = converter.convert()
@@ -299,7 +306,8 @@
self.assertTrue(num_items_graphviz)
# Convert model and ensure model is not None.
- converter = lite.TocoConverter.from_session(sess, [in_tensor], [out_tensor])
+ converter = lite.TFLiteConverter.from_session(sess, [in_tensor],
+ [out_tensor])
graphviz_dir = self.get_temp_dir()
converter.dump_graphviz_dir = graphviz_dir
converter.dump_graphviz_video = True
@@ -317,7 +325,8 @@
sess = session.Session()
# Convert model and ensure model is not None.
- converter = lite.TocoConverter.from_session(sess, [in_tensor], [out_tensor])
+ converter = lite.TFLiteConverter.from_session(sess, [in_tensor],
+ [out_tensor])
converter.inference_input_type = lite_constants.QUANTIZED_UINT8
converter.quantized_input_stats = {'Placeholder': (0., 1.)} # mean, std_dev
tflite_model = converter.convert()
@@ -347,7 +356,8 @@
sess = session.Session()
# Convert model and ensure model is not None.
- converter = lite.TocoConverter.from_session(sess, [in_tensor], [out_tensor])
+ converter = lite.TFLiteConverter.from_session(sess, [in_tensor],
+ [out_tensor])
converter.inference_type = lite_constants.QUANTIZED_UINT8
converter.quantized_input_stats = {'Placeholder': (0., 1.)} # mean, std_dev
converter.default_ranges_stats = (0, 6) # min, max
@@ -387,13 +397,13 @@
sess = session.Session()
# Convert float model.
- float_converter = lite.TocoConverter.from_session(sess, [in_tensor_1],
- [out_tensor])
+ float_converter = lite.TFLiteConverter.from_session(sess, [in_tensor_1],
+ [out_tensor])
float_tflite = float_converter.convert()
self.assertTrue(float_tflite)
# Convert quantized weights model.
- quantized_converter = lite.TocoConverter.from_session(
+ quantized_converter = lite.TFLiteConverter.from_session(
sess, [in_tensor_1], [out_tensor])
quantized_converter.post_training_quantize = True
quantized_tflite = quantized_converter.convert()
@@ -402,15 +412,16 @@
# Ensure that the quantized weights tflite model is smaller.
self.assertTrue(len(quantized_tflite) < len(float_tflite))
- def testExtendedMode(self):
+ def testFlexMode(self):
in_tensor = array_ops.placeholder(
shape=[1, 16, 16, 3], dtype=dtypes.float32)
out_tensor = in_tensor + in_tensor
sess = session.Session()
# Convert model and ensure model is not None.
- converter = lite.TocoConverter.from_session(sess, [in_tensor], [out_tensor])
- converter.converter_mode = lite.ConverterMode.TOCO_EXTENDED_ALL
+ converter = lite.TFLiteConverter.from_session(sess, [in_tensor],
+ [out_tensor])
+ converter.converter_mode = lite.ConverterMode.TOCO_FLEX_ALL
tflite_model = converter.convert()
self.assertTrue(tflite_model)
@@ -421,9 +432,25 @@
interpreter.allocate_tensors()
self.assertIn(
'Regular TensorFlow ops are not supported by this interpreter. Make '
- 'sure you invoke the Eager delegate before inference.',
+ 'sure you invoke the Flex delegate before inference.',
str(error.exception))
+ def testFloatTocoConverter(self):
+ """Tests deprecated test TocoConverter."""
+ in_tensor = array_ops.placeholder(
+ shape=[1, 16, 16, 3], dtype=dtypes.float32)
+ out_tensor = in_tensor + in_tensor
+ sess = session.Session()
+
+ # Convert model and ensure model is not None.
+ converter = lite.TocoConverter.from_session(sess, [in_tensor], [out_tensor])
+ tflite_model = converter.convert()
+ self.assertTrue(tflite_model)
+
+ # Ensure the interpreter is able to load.
+ interpreter = Interpreter(model_content=tflite_model)
+ interpreter.allocate_tensors()
+
class FromFrozenGraphFile(test_util.TensorFlowTestCase):
@@ -439,8 +466,8 @@
sess.close()
# Convert model and ensure model is not None.
- converter = lite.TocoConverter.from_frozen_graph(graph_def_file,
- ['Placeholder'], ['add'])
+ converter = lite.TFLiteConverter.from_frozen_graph(graph_def_file,
+ ['Placeholder'], ['add'])
tflite_model = converter.convert()
self.assertTrue(tflite_model)
@@ -474,7 +501,7 @@
sess.close()
# Convert model and ensure model is not None.
- converter = lite.TocoConverter.from_frozen_graph(
+ converter = lite.TFLiteConverter.from_frozen_graph(
graph_def_file, ['Placeholder'], ['add'],
input_shapes={'Placeholder': [1, 16, 16, 3]})
tflite_model = converter.convert()
@@ -503,8 +530,8 @@
# Ensure the graph with variables cannot be converted.
with self.assertRaises(ValueError) as error:
- lite.TocoConverter.from_frozen_graph(graph_def_file, ['Placeholder'],
- ['add'])
+ lite.TFLiteConverter.from_frozen_graph(graph_def_file, ['Placeholder'],
+ ['add'])
self.assertEqual('Please freeze the graph using freeze_graph.py.',
str(error.exception))
@@ -520,8 +547,8 @@
sess.close()
# Convert model and ensure model is not None.
- converter = lite.TocoConverter.from_frozen_graph(graph_def_file,
- ['Placeholder'], ['add'])
+ converter = lite.TFLiteConverter.from_frozen_graph(graph_def_file,
+ ['Placeholder'], ['add'])
tflite_model = converter.convert()
self.assertTrue(tflite_model)
@@ -545,8 +572,8 @@
def testInvalidFileNotFound(self):
with self.assertRaises(IOError) as error:
- lite.TocoConverter.from_frozen_graph('invalid_file', ['Placeholder'],
- ['add'])
+ lite.TFLiteConverter.from_frozen_graph('invalid_file', ['Placeholder'],
+ ['add'])
self.assertEqual('File \'invalid_file\' does not exist.',
str(error.exception))
@@ -558,8 +585,8 @@
# Attempts to convert the invalid model.
with self.assertRaises(IOError) as error:
- lite.TocoConverter.from_frozen_graph(graph_def_file, ['Placeholder'],
- ['add'])
+ lite.TFLiteConverter.from_frozen_graph(graph_def_file, ['Placeholder'],
+ ['add'])
self.assertEqual(
'Unable to parse input file \'{}\'.'.format(graph_def_file),
str(error.exception))
@@ -580,7 +607,7 @@
# Tests the object detection model that cannot be loaded in TensorFlow.
self._initObjectDetectionArgs()
- converter = lite.TocoConverter.from_frozen_graph(
+ converter = lite.TFLiteConverter.from_frozen_graph(
self._graph_def_file, self._input_arrays, self._output_arrays,
self._input_shapes)
converter.allow_custom_ops = True
@@ -621,7 +648,7 @@
# Missing `input_shapes`.
with self.assertRaises(ValueError) as error:
- lite.TocoConverter.from_frozen_graph(
+ lite.TFLiteConverter.from_frozen_graph(
self._graph_def_file, self._input_arrays, self._output_arrays)
self.assertEqual('input_shapes must be defined for this model.',
str(error.exception))
@@ -632,7 +659,7 @@
# `input_shapes` does not contain the names in `input_arrays`.
with self.assertRaises(ValueError) as error:
- lite.TocoConverter.from_frozen_graph(
+ lite.TFLiteConverter.from_frozen_graph(
self._graph_def_file,
self._input_arrays,
self._output_arrays,
@@ -641,6 +668,27 @@
'input_shapes must contain a value for each item in input_array.',
str(error.exception))
+ def testFloatTocoConverter(self):
+ in_tensor = array_ops.placeholder(
+ shape=[1, 16, 16, 3], dtype=dtypes.float32)
+ _ = in_tensor + in_tensor
+ sess = session.Session()
+
+ # Write graph to file.
+ graph_def_file = os.path.join(self.get_temp_dir(), 'model.pb')
+ write_graph(sess.graph_def, '', graph_def_file, False)
+ sess.close()
+
+ # Convert model and ensure model is not None.
+ converter = lite.TocoConverter.from_frozen_graph(graph_def_file,
+ ['Placeholder'], ['add'])
+ tflite_model = converter.convert()
+ self.assertTrue(tflite_model)
+
+ # Ensure the model is able to load.
+ interpreter = Interpreter(model_content=tflite_model)
+ interpreter.allocate_tensors()
+
class FromSavedModelTest(test_util.TensorFlowTestCase):
@@ -663,7 +711,7 @@
saved_model_dir = self._createSavedModel(shape=[1, 16, 16, 3])
# Convert model and ensure model is not None.
- converter = lite.TocoConverter.from_saved_model(saved_model_dir)
+ converter = lite.TFLiteConverter.from_saved_model(saved_model_dir)
tflite_model = converter.convert()
self.assertTrue(tflite_model)
@@ -693,7 +741,7 @@
"""Test a SavedModel, with None in input tensor's shape."""
saved_model_dir = self._createSavedModel(shape=[None, 16, 16, 3])
- converter = lite.TocoConverter.from_saved_model(saved_model_dir)
+ converter = lite.TFLiteConverter.from_saved_model(saved_model_dir)
tflite_model = converter.convert()
self.assertTrue(tflite_model)
@@ -724,7 +772,7 @@
"""Test a SavedModel ordering of input arrays."""
saved_model_dir = self._createSavedModel(shape=[1, 16, 16, 3])
- converter = lite.TocoConverter.from_saved_model(
+ converter = lite.TFLiteConverter.from_saved_model(
saved_model_dir, input_arrays=['inputB', 'inputA'])
tflite_model = converter.convert()
self.assertTrue(tflite_model)
@@ -757,7 +805,7 @@
saved_model_dir = self._createSavedModel(shape=[1, 16, 16, 3])
# Check case where input shape is given.
- converter = lite.TocoConverter.from_saved_model(
+ converter = lite.TFLiteConverter.from_saved_model(
saved_model_dir,
input_arrays=['inputA'],
input_shapes={'inputA': [1, 16, 16, 3]})
@@ -766,12 +814,25 @@
self.assertTrue(tflite_model)
# Check case where input shape is None.
- converter = lite.TocoConverter.from_saved_model(
+ converter = lite.TFLiteConverter.from_saved_model(
saved_model_dir, input_arrays=['inputA'], input_shapes={'inputA': None})
tflite_model = converter.convert()
self.assertTrue(tflite_model)
+ def testSimpleModelTocoConverter(self):
+ """Test a SavedModel with deprecated TocoConverter."""
+ saved_model_dir = self._createSavedModel(shape=[1, 16, 16, 3])
+
+ # Convert model and ensure model is not None.
+ converter = lite.TocoConverter.from_saved_model(saved_model_dir)
+ tflite_model = converter.convert()
+ self.assertTrue(tflite_model)
+
+ # Ensure the model is able to load.
+ interpreter = Interpreter(model_content=tflite_model)
+ interpreter.allocate_tensors()
+
class FromKerasFile(test_util.TensorFlowTestCase):
@@ -805,7 +866,7 @@
"""Test a Sequential tf.keras model with default inputs."""
keras_file = self._getSequentialModel()
- converter = lite.TocoConverter.from_keras_model_file(keras_file)
+ converter = lite.TFLiteConverter.from_keras_model_file(keras_file)
tflite_model = converter.convert()
self.assertTrue(tflite_model)
@@ -845,13 +906,13 @@
# Invalid input array raises error.
with self.assertRaises(ValueError) as error:
- lite.TocoConverter.from_keras_model_file(
+ lite.TFLiteConverter.from_keras_model_file(
keras_file, input_arrays=['invalid-input'])
self.assertEqual("Invalid tensors 'invalid-input' were found.",
str(error.exception))
# Valid input array.
- converter = lite.TocoConverter.from_keras_model_file(
+ converter = lite.TFLiteConverter.from_keras_model_file(
keras_file, input_arrays=['dense_input'])
tflite_model = converter.convert()
os.remove(keras_file)
@@ -863,13 +924,13 @@
# Passing in shape of invalid input array has no impact as long as all input
# arrays have a shape.
- converter = lite.TocoConverter.from_keras_model_file(
+ converter = lite.TFLiteConverter.from_keras_model_file(
keras_file, input_shapes={'invalid-input': [2, 3]})
tflite_model = converter.convert()
self.assertTrue(tflite_model)
# Passing in shape of valid input array.
- converter = lite.TocoConverter.from_keras_model_file(
+ converter = lite.TFLiteConverter.from_keras_model_file(
keras_file, input_shapes={'dense_input': [2, 3]})
tflite_model = converter.convert()
os.remove(keras_file)
@@ -890,13 +951,13 @@
# Invalid output array raises error.
with self.assertRaises(ValueError) as error:
- lite.TocoConverter.from_keras_model_file(
+ lite.TFLiteConverter.from_keras_model_file(
keras_file, output_arrays=['invalid-output'])
self.assertEqual("Invalid tensors 'invalid-output' were found.",
str(error.exception))
# Valid output array.
- converter = lite.TocoConverter.from_keras_model_file(
+ converter = lite.TFLiteConverter.from_keras_model_file(
keras_file, output_arrays=['time_distributed/Reshape_1'])
tflite_model = converter.convert()
os.remove(keras_file)
@@ -926,7 +987,7 @@
os.close(fd)
# Convert to TFLite model.
- converter = lite.TocoConverter.from_keras_model_file(keras_file)
+ converter = lite.TFLiteConverter.from_keras_model_file(keras_file)
tflite_model = converter.convert()
self.assertTrue(tflite_model)
@@ -991,7 +1052,7 @@
os.close(fd)
# Convert to TFLite model.
- converter = lite.TocoConverter.from_keras_model_file(keras_file)
+ converter = lite.TFLiteConverter.from_keras_model_file(keras_file)
tflite_model = converter.convert()
self.assertTrue(tflite_model)
@@ -1052,7 +1113,7 @@
os.close(fd)
# Convert to TFLite model.
- converter = lite.TocoConverter.from_keras_model_file(keras_file)
+ converter = lite.TFLiteConverter.from_keras_model_file(keras_file)
tflite_model = converter.convert()
self.assertTrue(tflite_model)
@@ -1086,6 +1147,18 @@
np.testing.assert_almost_equal(tflite_result, keras_result, 5)
os.remove(keras_file)
+ def testSequentialModelTocoConverter(self):
+ """Test a Sequential tf.keras model with deprecated TocoConverter."""
+ keras_file = self._getSequentialModel()
+
+ converter = lite.TocoConverter.from_keras_model_file(keras_file)
+ tflite_model = converter.convert()
+ self.assertTrue(tflite_model)
+
+ # Ensure the model is able to load.
+ interpreter = Interpreter(model_content=tflite_model)
+ interpreter.allocate_tensors()
+
if __name__ == '__main__':
test.main()
diff --git a/tensorflow/contrib/lite/python/tflite_convert.py b/tensorflow/contrib/lite/python/tflite_convert.py
index c0ff7f3..d6d9052 100644
--- a/tensorflow/contrib/lite/python/tflite_convert.py
+++ b/tensorflow/contrib/lite/python/tflite_convert.py
@@ -40,13 +40,13 @@
def _get_toco_converter(flags):
- """Makes a TocoConverter object based on the flags provided.
+ """Makes a TFLiteConverter object based on the flags provided.
Args:
flags: argparse.Namespace object containing TFLite flags.
Returns:
- TocoConverter object.
+ TFLiteConverter object.
Raises:
ValueError: Invalid flags.
@@ -68,17 +68,17 @@
"output_arrays": output_arrays
}
- # Create TocoConverter.
+ # Create TFLiteConverter.
if flags.graph_def_file:
- converter_fn = lite.TocoConverter.from_frozen_graph
+ converter_fn = lite.TFLiteConverter.from_frozen_graph
converter_kwargs["graph_def_file"] = flags.graph_def_file
elif flags.saved_model_dir:
- converter_fn = lite.TocoConverter.from_saved_model
+ converter_fn = lite.TFLiteConverter.from_saved_model
converter_kwargs["saved_model_dir"] = flags.saved_model_dir
converter_kwargs["tag_set"] = _parse_set(flags.saved_model_tag_set)
converter_kwargs["signature_key"] = flags.saved_model_signature_key
elif flags.keras_model_file:
- converter_fn = lite.TocoConverter.from_keras_model_file
+ converter_fn = lite.TFLiteConverter.from_keras_model_file
converter_kwargs["model_file"] = flags.keras_model_file
else:
raise ValueError("--graph_def_file, --saved_model_dir, or "
diff --git a/tensorflow/contrib/lite/testing/BUILD b/tensorflow/contrib/lite/testing/BUILD
index a4736bf..f0bfec2 100644
--- a/tensorflow/contrib/lite/testing/BUILD
+++ b/tensorflow/contrib/lite/testing/BUILD
@@ -13,6 +13,7 @@
load(
"//tensorflow:tensorflow.bzl",
"tf_cc_test",
+ "py_test",
)
[gen_zip_test(
@@ -163,7 +164,7 @@
":test_runner",
"//tensorflow/contrib/lite:builtin_op_data",
"//tensorflow/contrib/lite:framework",
- "//tensorflow/contrib/lite/delegates/eager:delegate",
+ "//tensorflow/contrib/lite/delegates/flex:delegate",
"//tensorflow/contrib/lite/kernels:builtin_ops",
],
)
@@ -362,4 +363,32 @@
],
)
+py_binary(
+ name = "model_coverage_lib",
+ srcs = ["//tensorflow/contrib/lite/testing:model_coverage/model_coverage_lib.py"],
+ srcs_version = "PY2AND3",
+ tags = ["no_pip"],
+ visibility = ["//tensorflow/contrib/lite:__subpackages__"],
+ deps = [
+ "//tensorflow/contrib/lite/python:lite",
+ "//tensorflow/python:platform",
+ ],
+)
+
+py_test(
+ name = "model_coverage_lib_test",
+ srcs = ["//tensorflow/contrib/lite/testing:model_coverage/model_coverage_lib_test.py"],
+ srcs_version = "PY2AND3",
+ tags = [
+ "no_oss",
+ "no_pip",
+ "no_windows",
+ "notap",
+ ],
+ deps = [
+ ":model_coverage_lib",
+ "//tensorflow/python:client_testlib",
+ ],
+)
+
tflite_portable_test_suite()
diff --git a/tensorflow/contrib/lite/testing/generate_examples.py b/tensorflow/contrib/lite/testing/generate_examples.py
index 014c80b..18036fa 100644
--- a/tensorflow/contrib/lite/testing/generate_examples.py
+++ b/tensorflow/contrib/lite/testing/generate_examples.py
@@ -81,9 +81,9 @@
action="store_true",
help="Include intermediate graphdefs in the output zip files.")
parser.add_argument(
- "--run_with_extended",
+ "--run_with_flex",
action="store_true",
- help="Whether the TFLite Extended converter is being used.")
+ help="Whether the TFLite Flex converter is being used.")
RANDOM_SEED = 342
TEST_INPUT_DEPTH = 3
@@ -339,11 +339,11 @@
graphdef_file.flush()
# TODO(aselle): Switch this to subprocess at some point.
- if "pb2lite" in bin_path and FLAGS.run_with_extended:
+ if "pb2lite" in bin_path and FLAGS.run_with_flex:
opts = ("--input_arrays={0} --output_arrays={1}".format(
",".join(input_arrays), ",".join(output_tensors)))
- elif FLAGS.run_with_extended:
- opts += " --allow_eager_ops --force_eager_ops"
+ elif FLAGS.run_with_flex:
+ opts += " --allow_flex_ops --force_flex_ops"
cmd = ("%s --input_file=%s --output_file=%s %s > %s 2>&1" %
(bin_path, graphdef_file.name, output_file.name, opts,
stdout_file.name))
@@ -3333,7 +3333,7 @@
# list of valid conversion modes is defined in
# generated_test_conversion_modes() in build_def.bzl.
test_function = ("make_%s_tests" % (out.replace(".zip", "").replace(
- "pb2lite", "").replace("toco-extended", "").rstrip("_")))
+ "pb2lite", "").replace("toco-flex", "").rstrip("_")))
if test_function not in globals():
raise RuntimeError("Can't find a test function to create %r. Tried %r" %
(out, test_function))
diff --git a/tensorflow/contrib/lite/testing/model_coverage/model_coverage_lib.py b/tensorflow/contrib/lite/testing/model_coverage/model_coverage_lib.py
new file mode 100644
index 0000000..5ca57d0
--- /dev/null
+++ b/tensorflow/contrib/lite/testing/model_coverage/model_coverage_lib.py
@@ -0,0 +1,249 @@
+# 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.
+# ==============================================================================
+"""Functions to test TFLite models."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from tensorflow.contrib.lite.python import convert_saved_model as _convert_saved_model
+from tensorflow.contrib.lite.python import lite as _lite
+from tensorflow.core.framework import graph_pb2 as _graph_pb2
+from tensorflow.python import keras as _keras
+from tensorflow.python.client import session as _session
+from tensorflow.python.framework.importer import import_graph_def as _import_graph_def
+from tensorflow.python.lib.io import file_io as _file_io
+from tensorflow.python.saved_model import signature_constants as _signature_constants
+from tensorflow.python.saved_model import tag_constants as _tag_constants
+
+
+def _convert(converter, **kwargs):
+ """Converts the model.
+
+ Args:
+ converter: TocoConverter object.
+ **kwargs: Additional arguments to be passed into the converter. Supported
+ flags are {"converter_mode", "post_training_quant"}.
+
+ Returns:
+ The converted TFLite model in serialized format.
+ """
+ if "converter_mode" in kwargs:
+ converter.converter_mode = kwargs["converter_mode"]
+ if "post_training_quantize" in kwargs:
+ converter.post_training_quantize = kwargs["post_training_quantize"]
+ return converter.convert()
+
+
+def _generate_random_input_data(tflite_model, seed=None):
+ """Generates input data based on the input tensors in the TFLite model.
+
+ Args:
+ tflite_model: Serialized TensorFlow Lite model.
+ seed: Integer seed for the random generator. (default None)
+
+ Returns:
+ List of np.ndarray.
+ """
+ interpreter = _lite.Interpreter(model_content=tflite_model)
+ interpreter.allocate_tensors()
+ input_details = interpreter.get_input_details()
+
+ if seed:
+ np.random.seed(seed=seed)
+ return [
+ np.array(
+ np.random.random_sample(input_tensor["shape"]),
+ dtype=input_tensor["dtype"]) for input_tensor in input_details
+ ]
+
+
+def _evaluate_tflite_model(tflite_model, input_data):
+ """Returns evaluation of input data on TFLite model.
+
+ Args:
+ tflite_model: Serialized TensorFlow Lite model.
+ input_data: List of np.ndarray.
+
+ Returns:
+ List of np.ndarray.
+ """
+ interpreter = _lite.Interpreter(model_content=tflite_model)
+ interpreter.allocate_tensors()
+
+ input_details = interpreter.get_input_details()
+ output_details = interpreter.get_output_details()
+
+ for input_tensor, tensor_data in zip(input_details, input_data):
+ interpreter.set_tensor(input_tensor["index"], tensor_data)
+
+ interpreter.invoke()
+ output_data = [
+ interpreter.get_tensor(output_tensor["index"])
+ for output_tensor in output_details
+ ]
+ return output_data
+
+
+def evaluate_frozen_graph(filename, input_arrays, output_arrays):
+ """Returns a function that evaluates the frozen graph on input data.
+
+ Args:
+ filename: Full filepath of file containing frozen GraphDef.
+ input_arrays: List of input tensors to freeze graph with.
+ output_arrays: List of output tensors to freeze graph with.
+
+ Returns:
+ Lambda function ([np.ndarray data] : [np.ndarray result]).
+ """
+ with _session.Session().as_default() as sess:
+ with _file_io.FileIO(filename, "rb") as f:
+ file_content = f.read()
+
+ graph_def = _graph_pb2.GraphDef()
+ graph_def.ParseFromString(file_content)
+ _import_graph_def(graph_def, name="")
+
+ inputs = _convert_saved_model.get_tensors_from_tensor_names(
+ sess.graph, input_arrays)
+ outputs = _convert_saved_model.get_tensors_from_tensor_names(
+ sess.graph, output_arrays)
+
+ return lambda input_data: sess.run(outputs, dict(zip(inputs, input_data)))
+
+
+def evaluate_saved_model(directory, tag_set, signature_key):
+ """Returns a function that evaluates the SavedModel on input data.
+
+ Args:
+ directory: SavedModel directory to convert.
+ tag_set: Set of tags identifying the MetaGraphDef within the SavedModel to
+ analyze. All tags in the tag set must be present.
+ signature_key: Key identifying SignatureDef containing inputs and outputs.
+
+ Returns:
+ Lambda function ([np.ndarray data] : [np.ndarray result]).
+ """
+ with _session.Session().as_default() as sess:
+ if tag_set is None:
+ tag_set = set([_tag_constants.SERVING])
+ if signature_key is None:
+ signature_key = _signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
+
+ meta_graph = _convert_saved_model.get_meta_graph_def(directory, tag_set)
+ signature_def = _convert_saved_model.get_signature_def(
+ meta_graph, signature_key)
+ inputs, outputs = _convert_saved_model.get_inputs_outputs(signature_def)
+
+ return lambda input_data: sess.run(outputs, dict(zip(inputs, input_data)))
+
+
+def evaluate_keras_model(filename):
+ """Returns a function that evaluates the tf.keras model on input data.
+
+ Args:
+ filename: Full filepath of HDF5 file containing the tf.keras model.
+
+ Returns:
+ Lambda function ([np.ndarray data] : [np.ndarray result]).
+ """
+ keras_model = _keras.models.load_model(filename)
+ return lambda input_data: [keras_model.predict(input_data)]
+
+
+# TODO(nupurgarg): Make this function a parameter to test_frozen_graph (and
+# related functions) in order to make it easy to use different data generators.
+def compare_models_random_data(tflite_model, tf_eval_func, tolerance=5):
+ """Compares TensorFlow and TFLite models with random data.
+
+ Args:
+ tflite_model: Serialized TensorFlow Lite model.
+ tf_eval_func: Lambda function that takes in input data and outputs the
+ results of the TensorFlow model ([np.ndarray data] : [np.ndarray result]).
+ tolerance: Decimal place to check accuracy to.
+ """
+ input_data = _generate_random_input_data(tflite_model)
+ tf_results = tf_eval_func(input_data)
+ tflite_results = _evaluate_tflite_model(tflite_model, input_data)
+ for tf_result, tflite_result in zip(tf_results, tflite_results):
+ np.testing.assert_almost_equal(tf_result, tflite_result, tolerance)
+
+
+def test_frozen_graph(filename,
+ input_arrays,
+ output_arrays,
+ input_shapes=None,
+ **kwargs):
+ """Validates the TensorFlow frozen graph converts to a TFLite model.
+
+ Converts the TensorFlow frozen graph to TFLite and checks the accuracy of the
+ model on random data.
+
+ Args:
+ filename: Full filepath of file containing frozen GraphDef.
+ input_arrays: List of input tensors to freeze graph with.
+ output_arrays: List of output tensors to freeze graph with.
+ input_shapes: Dict of strings representing input tensor names to list of
+ integers representing input shapes (e.g., {"foo" : [1, 16, 16, 3]}).
+ Automatically determined when input shapes is None (e.g., {"foo" : None}).
+ (default None)
+ **kwargs: Additional arguments to be passed into the converter.
+ """
+ converter = _lite.TocoConverter.from_frozen_graph(filename, input_arrays,
+ output_arrays, input_shapes)
+ tflite_model = _convert(converter, **kwargs)
+
+ tf_eval_func = evaluate_frozen_graph(filename, input_arrays, output_arrays)
+ compare_models_random_data(tflite_model, tf_eval_func)
+
+
+def test_saved_model(directory, tag_set=None, signature_key=None, **kwargs):
+ """Validates the TensorFlow SavedModel converts to a TFLite model.
+
+ Converts the TensorFlow SavedModel to TFLite and checks the accuracy of the
+ model on random data.
+
+ Args:
+ directory: SavedModel directory to convert.
+ tag_set: Set of tags identifying the MetaGraphDef within the SavedModel to
+ analyze. All tags in the tag set must be present.
+ signature_key: Key identifying SignatureDef containing inputs and outputs.
+ **kwargs: Additional arguments to be passed into the converter.
+ """
+ converter = _lite.TocoConverter.from_saved_model(directory, tag_set,
+ signature_key)
+ tflite_model = _convert(converter, **kwargs)
+
+ tf_eval_func = evaluate_saved_model(directory, tag_set, signature_key)
+ compare_models_random_data(tflite_model, tf_eval_func)
+
+
+def test_keras_model(filename, **kwargs):
+ """Validates the tf.keras model converts to a TFLite model.
+
+ Converts the tf.keras model to TFLite and checks the accuracy of the model on
+ random data.
+
+ Args:
+ filename: Full filepath of HDF5 file containing the tf.keras model.
+ **kwargs: Additional arguments to be passed into the converter.
+ """
+ converter = _lite.TocoConverter.from_keras_model_file(filename)
+ tflite_model = _convert(converter, **kwargs)
+
+ tf_eval_func = evaluate_keras_model(filename)
+ compare_models_random_data(tflite_model, tf_eval_func)
diff --git a/tensorflow/contrib/lite/testing/model_coverage/model_coverage_lib_test.py b/tensorflow/contrib/lite/testing/model_coverage/model_coverage_lib_test.py
new file mode 100644
index 0000000..1498f86
--- /dev/null
+++ b/tensorflow/contrib/lite/testing/model_coverage/model_coverage_lib_test.py
@@ -0,0 +1,130 @@
+# 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 model_coverage_lib.py."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import tempfile
+
+from tensorflow.contrib.lite.python import lite
+from tensorflow.contrib.lite.testing.model_coverage import model_coverage_lib as model_coverage
+from tensorflow.python import keras
+from tensorflow.python.client import session
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+from tensorflow.python.saved_model import saved_model
+from tensorflow.python.training.training_util import write_graph
+
+
+class EvaluateFrozenGraph(test.TestCase):
+
+ def _saveFrozenGraph(self, sess):
+ graph_def_file = os.path.join(self.get_temp_dir(), 'model.pb')
+ write_graph(sess.graph_def, '', graph_def_file, False)
+ return graph_def_file
+
+ def testFloat(self):
+ with session.Session().as_default() as sess:
+ in_tensor = array_ops.placeholder(
+ shape=[1, 16, 16, 3], dtype=dtypes.float32)
+ _ = in_tensor + in_tensor
+ filename = self._saveFrozenGraph(sess)
+
+ model_coverage.test_frozen_graph(filename, ['Placeholder'], ['add'])
+
+ def testMultipleOutputs(self):
+ with session.Session().as_default() as sess:
+ in_tensor_1 = array_ops.placeholder(
+ shape=[1, 16], dtype=dtypes.float32, name='inputA')
+ in_tensor_2 = array_ops.placeholder(
+ shape=[1, 16], dtype=dtypes.float32, name='inputB')
+
+ weight = constant_op.constant(-1.0, shape=[16, 16])
+ bias = constant_op.constant(-1.0, shape=[16])
+ layer = math_ops.matmul(in_tensor_1, weight) + bias
+ _ = math_ops.reduce_mean(math_ops.square(layer - in_tensor_2))
+ filename = self._saveFrozenGraph(sess)
+
+ model_coverage.test_frozen_graph(filename, ['inputA', 'inputB'],
+ ['add', 'Mean'])
+
+
+class EvaluateSavedModel(test.TestCase):
+
+ def testFloat(self):
+ saved_model_dir = os.path.join(self.get_temp_dir(), 'simple_savedmodel')
+ with session.Session().as_default() as sess:
+ in_tensor_1 = array_ops.placeholder(
+ shape=[1, 16, 16, 3], dtype=dtypes.float32, name='inputB')
+ in_tensor_2 = array_ops.placeholder(
+ shape=[1, 16, 16, 3], dtype=dtypes.float32, name='inputA')
+ out_tensor = in_tensor_1 + in_tensor_2
+
+ inputs = {'x': in_tensor_1, 'y': in_tensor_2}
+ outputs = {'z': out_tensor}
+ saved_model.simple_save(sess, saved_model_dir, inputs, outputs)
+ model_coverage.test_saved_model(saved_model_dir)
+
+
+class EvaluateKerasModel(test.TestCase):
+
+ def _getSingleInputKerasModel(self):
+ """Returns single input Sequential tf.keras model."""
+ keras.backend.clear_session()
+
+ xs = [-1, 0, 1, 2, 3, 4]
+ ys = [-3, -1, 1, 3, 5, 7]
+
+ model = keras.Sequential([keras.layers.Dense(units=1, input_shape=[1])])
+ model.compile(optimizer='sgd', loss='mean_squared_error')
+ model.train_on_batch(xs, ys)
+ return model
+
+ def _saveKerasModel(self, model):
+ try:
+ fd, keras_file = tempfile.mkstemp('.h5')
+ keras.models.save_model(model, keras_file)
+ finally:
+ os.close(fd)
+ return keras_file
+
+ def testFloat(self):
+ model = self._getSingleInputKerasModel()
+ keras_file = self._saveKerasModel(model)
+
+ model_coverage.test_keras_model(keras_file)
+
+ def testPostTrainingQuantize(self):
+ model = self._getSingleInputKerasModel()
+ keras_file = self._saveKerasModel(model)
+
+ model_coverage.test_keras_model(keras_file, post_training_quantize=True)
+
+ def testConverterMode(self):
+ model = self._getSingleInputKerasModel()
+ keras_file = self._saveKerasModel(model)
+
+ model_coverage.test_keras_model(
+ keras_file, converter_mode=lite.ConverterMode.TOCO_FLEX)
+
+
+if __name__ == '__main__':
+ test.main()
diff --git a/tensorflow/contrib/lite/testing/tflite_diff_flags.h b/tensorflow/contrib/lite/testing/tflite_diff_flags.h
index 3874bc3..ad889a2 100644
--- a/tensorflow/contrib/lite/testing/tflite_diff_flags.h
+++ b/tensorflow/contrib/lite/testing/tflite_diff_flags.h
@@ -57,7 +57,7 @@
"[optional] Number of full runs in each pass."),
tensorflow::Flag("delegate", &values.delegate,
"[optional] Delegate to use for executing ops. Must be "
- "`{\"\", EAGER}`"),
+ "`{\"\", FLEX}`"),
};
bool no_inputs = *argc == 1;
@@ -70,7 +70,7 @@
values.input_layer_shape.empty() || values.output_layer.empty()) {
fprintf(stderr, "%s", tensorflow::Flags::Usage(argv[0], flags).c_str());
return {};
- } else if (!(values.delegate == "" || values.delegate == "EAGER")) {
+ } else if (!(values.delegate == "" || values.delegate == "FLEX")) {
fprintf(stderr, "%s", tensorflow::Flags::Usage(argv[0], flags).c_str());
return {};
}
diff --git a/tensorflow/contrib/lite/testing/tflite_diff_util.h b/tensorflow/contrib/lite/testing/tflite_diff_util.h
index f679921..28b14bd 100644
--- a/tensorflow/contrib/lite/testing/tflite_diff_util.h
+++ b/tensorflow/contrib/lite/testing/tflite_diff_util.h
@@ -45,7 +45,7 @@
// second pass does multiple inferences back to back.
int num_runs_per_pass;
// Path to the delegate library to be loaded in order to execute ops. Must be
- // `{"", EAGER}`.
+ // `{"", FLEX}`.
string delegate;
};
diff --git a/tensorflow/contrib/lite/testing/tflite_driver.cc b/tensorflow/contrib/lite/testing/tflite_driver.cc
index 1836eb5..ef49e6f 100644
--- a/tensorflow/contrib/lite/testing/tflite_driver.cc
+++ b/tensorflow/contrib/lite/testing/tflite_driver.cc
@@ -17,7 +17,7 @@
#include <iostream>
#include "tensorflow/contrib/lite/builtin_op_data.h"
-#include "tensorflow/contrib/lite/delegates/eager/delegate.h"
+#include "tensorflow/contrib/lite/delegates/flex/delegate.h"
#include "tensorflow/contrib/lite/testing/split.h"
namespace tflite {
@@ -138,8 +138,8 @@
TfLiteDriver::TfLiteDriver(bool use_nnapi, const string& delegate_name)
: use_nnapi_(use_nnapi) {
- if (delegate_name == "EAGER") {
- delegate_ = EagerDelegate::Create();
+ if (delegate_name == "FLEX") {
+ delegate_ = FlexDelegate::Create();
}
}
@@ -301,7 +301,7 @@
}
void TfLiteDriver::ResetLSTMStateTensors() {
- interpreter_->ResetVariableTensorsToZero();
+ interpreter_->ResetVariableTensors();
}
} // namespace testing
diff --git a/tensorflow/contrib/lite/testing/tflite_driver.h b/tensorflow/contrib/lite/testing/tflite_driver.h
index aed35f8..dc2a4e5 100644
--- a/tensorflow/contrib/lite/testing/tflite_driver.h
+++ b/tensorflow/contrib/lite/testing/tflite_driver.h
@@ -17,7 +17,7 @@
#include <map>
-#include "tensorflow/contrib/lite/delegates/eager/delegate.h"
+#include "tensorflow/contrib/lite/delegates/flex/delegate.h"
#include "tensorflow/contrib/lite/interpreter.h"
#include "tensorflow/contrib/lite/kernels/register.h"
#include "tensorflow/contrib/lite/model.h"
@@ -53,7 +53,7 @@
class Expectation;
- std::unique_ptr<EagerDelegate> delegate_;
+ std::unique_ptr<FlexDelegate> delegate_;
bool use_nnapi_ = false;
std::unique_ptr<FlatBufferModel> model_;
std::unique_ptr<Interpreter> interpreter_;
diff --git a/tensorflow/contrib/lite/toco/args.h b/tensorflow/contrib/lite/toco/args.h
index f14dbc2..2699ac7 100644
--- a/tensorflow/contrib/lite/toco/args.h
+++ b/tensorflow/contrib/lite/toco/args.h
@@ -248,9 +248,9 @@
Arg<int64> dedupe_array_min_size_bytes = Arg<int64>(64);
Arg<bool> split_tflite_lstm_inputs = Arg<bool>(true);
// WARNING: Experimental interface, subject to change
- Arg<bool> allow_eager_ops = Arg<bool>(false);
+ Arg<bool> allow_flex_ops = Arg<bool>(false);
// WARNING: Experimental interface, subject to change
- Arg<bool> force_eager_ops = Arg<bool>(false);
+ Arg<bool> force_flex_ops = Arg<bool>(false);
};
} // namespace toco
diff --git a/tensorflow/contrib/lite/toco/g3doc/python_api.md b/tensorflow/contrib/lite/toco/g3doc/python_api.md
index 910fa4c..8c31c3d 100644
--- a/tensorflow/contrib/lite/toco/g3doc/python_api.md
+++ b/tensorflow/contrib/lite/toco/g3doc/python_api.md
@@ -39,13 +39,18 @@
is `tf.contrib.lite.TocoConverter`. The API for calling the Python intepreter is
`tf.contrib.lite.Interpreter`.
+**NOTE**: As of TensorFlow 1.12, the API for converting TensorFlow models to
+TFLite will be renamed to `TFLiteConverter`. `TFLiteConverter` is semantically
+identically to `TocoConverter`. The API is available at
+`tf.contrib.lite.TFLiteConverter` as of the Sept 26 `tf-nightly`.
+
`TocoConverter` provides class methods based on the original format of the
model. `TocoConverter.from_session()` is available for GraphDefs.
`TocoConverter.from_saved_model()` is available for SavedModels.
`TocoConverter.from_keras_model_file()` is available for `tf.Keras` files.
-Example usages for simple float-point models are shown in [Basic
-Examples](#basic). Examples usages for more complex models is shown in [Complex
-Examples](#complex).
+Example usages for simple float-point models are shown in
+[Basic Examples](#basic). Examples usages for more complex models is shown in
+[Complex Examples](#complex).
**NOTE**: Currently, `TocoConverter` will cause a fatal error to the Python
interpreter when the conversion fails. This will be remedied as soon as
diff --git a/tensorflow/contrib/lite/toco/import_tensorflow.cc b/tensorflow/contrib/lite/toco/import_tensorflow.cc
index e02d000..5eaf6e2 100644
--- a/tensorflow/contrib/lite/toco/import_tensorflow.cc
+++ b/tensorflow/contrib/lite/toco/import_tensorflow.cc
@@ -2123,9 +2123,9 @@
Model* model = new Model;
internal::ConverterMapType converter_map;
- // This is used for the TFLite "Full Eager Mode" conversion. All the ops are
+ // This is used for the TFLite "Full Flex Mode" conversion. All the ops are
// imported as `TensorFlowUnsupportedOperator`, and later all these ops are
- // converted to TFLite Eager ops.
+ // converted to TFLite Flex ops.
if (!tf_import_flags.import_all_ops_as_unsupported) {
converter_map = internal::GetTensorFlowNodeConverterMap();
}
diff --git a/tensorflow/contrib/lite/toco/import_tensorflow.h b/tensorflow/contrib/lite/toco/import_tensorflow.h
index 7db23f2..c5ff969 100644
--- a/tensorflow/contrib/lite/toco/import_tensorflow.h
+++ b/tensorflow/contrib/lite/toco/import_tensorflow.h
@@ -30,7 +30,7 @@
// Do not recognize any op and import all ops as
// `TensorFlowUnsupportedOperator`. This is used to populated with the
- // `force_eager_ops` flag.
+ // `force_flex_ops` flag.
bool import_all_ops_as_unsupported = false;
};
diff --git a/tensorflow/contrib/lite/toco/python/BUILD b/tensorflow/contrib/lite/toco/python/BUILD
index 33c5b16..cf97ba7 100644
--- a/tensorflow/contrib/lite/toco/python/BUILD
+++ b/tensorflow/contrib/lite/toco/python/BUILD
@@ -4,6 +4,7 @@
load("//tensorflow:tensorflow.bzl", "tf_py_wrap_cc")
load("//tensorflow:tensorflow.bzl", "tf_py_test")
+load("//tensorflow:tensorflow.bzl", "py_binary")
cc_library(
name = "toco_python_api",
diff --git a/tensorflow/contrib/lite/toco/tflite/export.cc b/tensorflow/contrib/lite/toco/tflite/export.cc
index fee10b1..0c9fac2 100644
--- a/tensorflow/contrib/lite/toco/tflite/export.cc
+++ b/tensorflow/contrib/lite/toco/tflite/export.cc
@@ -50,16 +50,16 @@
details::OperatorKey GetOperatorKey(
const ::toco::Operator& op,
const std::map<OperatorType, std::unique_ptr<BaseOperator>>& ops_by_type,
- bool allow_eager_ops) {
+ bool allow_flex_ops) {
string custom_code;
if (op.type == OperatorType::kUnsupported) {
const TensorFlowUnsupportedOperator& unsupported_op =
static_cast<const TensorFlowUnsupportedOperator&>(op);
- // TODO(b/113715895): When `allow_eager_ops` is on, for now there's no way
+ // TODO(b/113715895): When `allow_flex_ops` is on, for now there's no way
// to populate a regular custom op. We need to find a way to fix this.
- if (allow_eager_ops) {
- custom_code = string(::tflite::kEagerCustomCodePrefix) +
+ if (allow_flex_ops) {
+ custom_code = string(::tflite::kFlexCustomCodePrefix) +
unsupported_op.tensorflow_op;
} else {
custom_code = unsupported_op.tensorflow_op;
@@ -101,11 +101,11 @@
void LoadOperatorsMap(
const Model& model, OperatorsMap* operators_map,
const std::map<OperatorType, std::unique_ptr<BaseOperator>>& ops_by_type,
- bool allow_eager_ops) {
+ bool allow_flex_ops) {
// First find a list of unique operator types.
std::set<OperatorKey> keys;
for (const auto& op : model.operators) {
- keys.insert(GetOperatorKey(*op, ops_by_type, allow_eager_ops));
+ keys.insert(GetOperatorKey(*op, ops_by_type, allow_flex_ops));
}
// Now assign indices to them and fill in the map.
int index = 0;
@@ -216,7 +216,7 @@
for (const auto& op : model.operators) {
const details::OperatorKey operator_key =
- GetOperatorKey(*op, ops_by_type, params.allow_eager_ops);
+ GetOperatorKey(*op, ops_by_type, params.allow_flex_ops);
int op_index = operators_map.at(operator_key);
int op_version = operator_key.version;
@@ -281,7 +281,7 @@
}
int op_index = operators_map.at(
- GetOperatorKey(*op, ops_by_type, params.allow_eager_ops));
+ GetOperatorKey(*op, ops_by_type, params.allow_flex_ops));
auto tflite_op_it = ops_by_type.find(op->type);
BaseOperator* tflite_op = tflite_op_it == ops_by_type.end()
@@ -334,7 +334,7 @@
void Export(const Model& model, string* output_file_contents,
const ExportParams& params) {
- const auto ops_by_type = BuildOperatorByTypeMap(params.allow_eager_ops);
+ const auto ops_by_type = BuildOperatorByTypeMap(params.allow_flex_ops);
Export(model, output_file_contents, params, ops_by_type);
}
@@ -349,7 +349,7 @@
details::OperatorsMap operators_map;
details::LoadOperatorsMap(model, &operators_map, ops_by_type,
- params.allow_eager_ops);
+ params.allow_flex_ops);
std::vector<const Array*> buffers_to_write;
Array empty_array;
@@ -388,7 +388,7 @@
"the standard TensorFlow Lite runtime. If you have a custom "
"implementation for them you can disable this error with "
"--allow_custom_ops, or by setting allow_custom_ops=True "
- "when calling tf.contrib.lite.TocoConverter(). Here is a list "
+ "when calling tf.contrib.lite.TFLiteConverter(). Here is a list "
"of operators for which you will need custom implementations: "
<< absl::StrJoin(error_summary_final, ", ") << ".";
}
diff --git a/tensorflow/contrib/lite/toco/tflite/export.h b/tensorflow/contrib/lite/toco/tflite/export.h
index b070a38..29d6de4 100644
--- a/tensorflow/contrib/lite/toco/tflite/export.h
+++ b/tensorflow/contrib/lite/toco/tflite/export.h
@@ -26,7 +26,7 @@
// The parameters for exporting a TFLite model.
struct ExportParams {
bool allow_custom_ops = false;
- bool allow_eager_ops = false;
+ bool allow_flex_ops = false;
bool quantize_weights = false;
};
@@ -121,7 +121,7 @@
void LoadOperatorsMap(
const Model& model, OperatorsMap* operators_map,
const std::map<OperatorType, std::unique_ptr<BaseOperator>>& ops_by_type,
- bool allow_eager_ops);
+ bool allow_flex_ops);
} // namespace details
} // namespace tflite
diff --git a/tensorflow/contrib/lite/toco/tflite/export_test.cc b/tensorflow/contrib/lite/toco/tflite/export_test.cc
index 8d4d197..93882a9 100644
--- a/tensorflow/contrib/lite/toco/tflite/export_test.cc
+++ b/tensorflow/contrib/lite/toco/tflite/export_test.cc
@@ -105,7 +105,7 @@
details::OperatorsMap operators;
const auto ops_by_type = BuildOperatorByTypeMap();
- // TODO(ycling): Add a test for allow_eager_ops.
+ // TODO(ycling): Add a test for allow_flex_ops.
details::LoadOperatorsMap(input_model_, &operators, ops_by_type, false);
EXPECT_EQ(0, operators[details::OperatorKey(OperatorType::kAdd, "", 1)]);
EXPECT_EQ(1, operators[details::OperatorKey(OperatorType::kConv, "", 1)]);
diff --git a/tensorflow/contrib/lite/toco/tflite/operator.cc b/tensorflow/contrib/lite/toco/tflite/operator.cc
index ca2a6a1..9addbb8 100644
--- a/tensorflow/contrib/lite/toco/tflite/operator.cc
+++ b/tensorflow/contrib/lite/toco/tflite/operator.cc
@@ -1160,8 +1160,8 @@
class TensorFlowUnsupported : public BaseOperator {
public:
TensorFlowUnsupported(const string& name, OperatorType type,
- bool allow_eager_ops)
- : BaseOperator(name, type), allow_eager_ops_(allow_eager_ops) {}
+ bool allow_flex_ops)
+ : BaseOperator(name, type), allow_flex_ops_(allow_flex_ops) {}
Options Serialize(const Operator& op,
flatbuffers::FlatBufferBuilder* builder) const override {
@@ -1177,9 +1177,9 @@
std::unique_ptr<Operator> Deserialize(
const BuiltinOptions* builtin_options,
const CustomOptions* custom_options) const override {
- // Deserializing Eager ops doesn't work now.
+ // Deserializing Flex ops doesn't work now.
// TODO(ycling): Revisit and decide if we should fix the flow for importing
- // TFLite models with Eager ops.
+ // TFLite models with Flex ops.
auto op = absl::make_unique<TensorFlowUnsupportedOperator>();
if (custom_options) {
auto flexbuffer_map =
@@ -1200,13 +1200,13 @@
return std::unique_ptr<flexbuffers::Builder>();
}
- if (allow_eager_ops_) {
+ if (allow_flex_ops_) {
fbb->Vector([&]() {
fbb->String(node_def.op());
fbb->String(op.tensorflow_node_def);
});
fbb->Finish();
- LOG(INFO) << "Writing eager op: " << node_def.op();
+ LOG(INFO) << "Writing flex op: " << node_def.op();
return std::unique_ptr<flexbuffers::Builder>(fbb.release());
}
@@ -1316,13 +1316,13 @@
}
private:
- const bool allow_eager_ops_;
+ const bool allow_flex_ops_;
};
namespace {
// Build a vector containing all the known operators.
std::vector<std::unique_ptr<BaseOperator>> BuildOperatorList(
- bool allow_eager_ops = false) {
+ bool allow_flex_ops = false) {
std::vector<std::unique_ptr<BaseOperator>> ops;
using tensorflow::MakeUnique;
// Builtin Operators.
@@ -1434,7 +1434,7 @@
ops.push_back(MakeUnique<CTCBeamSearchDecoder>(
"CTC_BEAM_SEARCH_DECODER", OperatorType::kCTCBeamSearchDecoder));
ops.push_back(MakeUnique<TensorFlowUnsupported>(
- "TENSORFLOW_UNSUPPORTED", OperatorType::kUnsupported, allow_eager_ops));
+ "TENSORFLOW_UNSUPPORTED", OperatorType::kUnsupported, allow_flex_ops));
// There operators are supported by Toco, but not by TF Lite, and has no
// attributes.
@@ -1512,11 +1512,11 @@
} // namespace
std::map<OperatorType, std::unique_ptr<BaseOperator>> BuildOperatorByTypeMap(
- bool allow_eager_ops) {
+ bool allow_flex_ops) {
std::map<OperatorType, std::unique_ptr<BaseOperator>> result;
std::vector<std::unique_ptr<BaseOperator>> ops =
- BuildOperatorList(allow_eager_ops);
+ BuildOperatorList(allow_flex_ops);
for (auto& op : ops) {
result[op->type()] = std::move(op);
}
@@ -1525,11 +1525,11 @@
}
std::map<string, std::unique_ptr<BaseOperator>> BuildOperatorByNameMap(
- bool allow_eager_ops) {
+ bool allow_flex_ops) {
std::map<string, std::unique_ptr<BaseOperator>> result;
std::vector<std::unique_ptr<BaseOperator>> ops =
- BuildOperatorList(allow_eager_ops);
+ BuildOperatorList(allow_flex_ops);
for (auto& op : ops) {
result[op->name()] = std::move(op);
}
diff --git a/tensorflow/contrib/lite/toco/tflite/operator.h b/tensorflow/contrib/lite/toco/tflite/operator.h
index 702fb28..13d9f6c 100644
--- a/tensorflow/contrib/lite/toco/tflite/operator.h
+++ b/tensorflow/contrib/lite/toco/tflite/operator.h
@@ -26,15 +26,15 @@
class BaseOperator;
// Return a map contained all know TF Lite Operators, keyed by their names.
-// TODO(ycling): The pattern to propagate parameters (e.g. allow_eager_ops)
+// TODO(ycling): The pattern to propagate parameters (e.g. allow_flex_ops)
// is ugly here. Consider refactoring.
std::map<string, std::unique_ptr<BaseOperator>> BuildOperatorByNameMap(
- bool allow_eager_ops = false);
+ bool allow_flex_ops = false);
// Return a map contained all know TF Lite Operators, keyed by the type of
// their tf.mini counterparts.
std::map<OperatorType, std::unique_ptr<BaseOperator>> BuildOperatorByTypeMap(
- bool allow_eager_ops = false);
+ bool allow_flex_ops = false);
// These are the flatbuffer types for custom and builtin options.
using CustomOptions = flatbuffers::Vector<uint8_t>;
diff --git a/tensorflow/contrib/lite/toco/toco_cmdline_flags.cc b/tensorflow/contrib/lite/toco/toco_cmdline_flags.cc
index b6aebc0..cff7977 100644
--- a/tensorflow/contrib/lite/toco/toco_cmdline_flags.cc
+++ b/tensorflow/contrib/lite/toco/toco_cmdline_flags.cc
@@ -167,11 +167,11 @@
"converted float model. Model size will be reduced and there will "
"be latency improvements (at the cost of accuracy)."),
// WARNING: Experimental interface, subject to change
- Flag("allow_eager_ops", parsed_flags.allow_eager_ops.bind(),
- parsed_flags.allow_eager_ops.default_value(), ""),
+ Flag("allow_flex_ops", parsed_flags.allow_flex_ops.bind(),
+ parsed_flags.allow_flex_ops.default_value(), ""),
// WARNING: Experimental interface, subject to change
- Flag("force_eager_ops", parsed_flags.force_eager_ops.bind(),
- parsed_flags.force_eager_ops.default_value(), "")};
+ Flag("force_flex_ops", parsed_flags.force_flex_ops.bind(),
+ parsed_flags.force_flex_ops.default_value(), "")};
bool asked_for_help =
*argc == 2 && (!strcmp(argv[1], "--help") || !strcmp(argv[1], "-help"));
if (asked_for_help) {
@@ -266,15 +266,15 @@
READ_TOCO_FLAG(split_tflite_lstm_inputs, FlagRequirement::kNone);
READ_TOCO_FLAG(quantize_weights, FlagRequirement::kNone);
READ_TOCO_FLAG(post_training_quantize, FlagRequirement::kNone);
- READ_TOCO_FLAG(allow_eager_ops, FlagRequirement::kNone);
- READ_TOCO_FLAG(force_eager_ops, FlagRequirement::kNone);
+ READ_TOCO_FLAG(allow_flex_ops, FlagRequirement::kNone);
+ READ_TOCO_FLAG(force_flex_ops, FlagRequirement::kNone);
- if (parsed_toco_flags.force_eager_ops.value() &&
- !parsed_toco_flags.allow_eager_ops.value()) {
- // TODO(ycling): Consider to enforce `allow_eager_ops` when
- // `force_eager_ops` is true.
- LOG(WARNING) << "--force_eager_ops should always be used with "
- "--allow_eager_ops.";
+ if (parsed_toco_flags.force_flex_ops.value() &&
+ !parsed_toco_flags.allow_flex_ops.value()) {
+ // TODO(ycling): Consider to enforce `allow_flex_ops` when
+ // `force_flex_ops` is true.
+ LOG(WARNING) << "--force_flex_ops should always be used with "
+ "--allow_flex_ops.";
}
// Deprecated flag handling.
diff --git a/tensorflow/contrib/lite/toco/toco_flags.proto b/tensorflow/contrib/lite/toco/toco_flags.proto
index 53d60fe..ca3e644 100644
--- a/tensorflow/contrib/lite/toco/toco_flags.proto
+++ b/tensorflow/contrib/lite/toco/toco_flags.proto
@@ -190,16 +190,16 @@
// (at the cost of accuracy).
optional bool post_training_quantize = 26 [default = false];
- // When enabled, unsupported ops will be converted to TFLite Eager ops.
+ // When enabled, unsupported ops will be converted to TFLite Flex ops.
// TODO(ycling): Consider to rename the following 2 flags and don't call it
- // "Eager".
- // `allow_eager_ops` should always be used with `allow_custom_ops`.
+ // "Flex".
+ // `allow_flex_ops` should always be used with `allow_custom_ops`.
// WARNING: Experimental interface, subject to change
- optional bool allow_eager_ops = 27 [default = false];
+ optional bool allow_flex_ops = 27 [default = false];
- // When enabled, all TensorFlow ops will be converted to TFLite Eager
- // ops directly. This will force `allow_eager_ops` to true.
- // `force_eager_ops` should always be used with `allow_eager_ops`.
+ // When enabled, all TensorFlow ops will be converted to TFLite Flex
+ // ops directly. This will force `allow_flex_ops` to true.
+ // `force_flex_ops` should always be used with `allow_flex_ops`.
// WARNING: Experimental interface, subject to change
- optional bool force_eager_ops = 28 [default = false];
+ optional bool force_flex_ops = 28 [default = false];
}
diff --git a/tensorflow/contrib/lite/toco/toco_tooling.cc b/tensorflow/contrib/lite/toco/toco_tooling.cc
index a08b024..106494f 100644
--- a/tensorflow/contrib/lite/toco/toco_tooling.cc
+++ b/tensorflow/contrib/lite/toco/toco_tooling.cc
@@ -198,7 +198,7 @@
: (toco_flags.output_format() != TENSORFLOW_GRAPHDEF);
tf_import_flags.import_all_ops_as_unsupported =
- toco_flags.force_eager_ops();
+ toco_flags.force_flex_ops();
model = ImportTensorFlowGraphDef(model_flags, tf_import_flags,
input_file_contents);
@@ -409,9 +409,9 @@
case TFLITE: {
toco::tflite::ExportParams params;
- // Always allow custom ops when eager ops are allowed.
- if (toco_flags.force_eager_ops() || toco_flags.allow_eager_ops()) {
- params.allow_eager_ops = true;
+ // Always allow custom ops when flex ops are allowed.
+ if (toco_flags.force_flex_ops() || toco_flags.allow_flex_ops()) {
+ params.allow_flex_ops = true;
params.allow_custom_ops = true;
} else if (allow_custom_ops) {
params.allow_custom_ops = true;
diff --git a/tensorflow/contrib/lite/tools/benchmark/BUILD b/tensorflow/contrib/lite/tools/benchmark/BUILD
index dc97d22..502e181 100644
--- a/tensorflow/contrib/lite/tools/benchmark/BUILD
+++ b/tensorflow/contrib/lite/tools/benchmark/BUILD
@@ -36,11 +36,11 @@
)
cc_binary(
- name = "benchmark_model_plus_eager",
+ name = "benchmark_model_plus_flex",
srcs = [
"benchmark_main.cc",
],
- copts = common_copts + ["-DTFLITE_EXTENDED"],
+ copts = common_copts + ["-DTFLITE_FLEX"],
linkopts = tflite_linkopts() + select({
"//tensorflow:android": [
"-pie", # Android 5.0 and later supports only PIE
@@ -49,7 +49,7 @@
"//conditions:default": [],
}),
deps = [
- ":benchmark_tflite_model_plus_eager_lib",
+ ":benchmark_tflite_model_plus_flex_lib",
":logging",
],
)
@@ -111,19 +111,19 @@
)
cc_library(
- name = "benchmark_tflite_model_plus_eager_lib",
+ name = "benchmark_tflite_model_plus_flex_lib",
srcs = [
"benchmark_tflite_model.cc",
"logging.h",
],
hdrs = ["benchmark_tflite_model.h"],
- copts = common_copts + ["-DTFLITE_EXTENDED"],
+ copts = common_copts + ["-DTFLITE_FLEX"],
deps = [
":benchmark_model_lib",
":logging",
"//tensorflow/contrib/lite:framework",
"//tensorflow/contrib/lite:string_util",
- "//tensorflow/contrib/lite/delegates/eager:delegate",
+ "//tensorflow/contrib/lite/delegates/flex:delegate",
"//tensorflow/contrib/lite/kernels:builtin_ops",
"//tensorflow/contrib/lite/profiling:profile_summarizer",
],
diff --git a/tensorflow/contrib/lite/tools/benchmark/benchmark_tflite_model.cc b/tensorflow/contrib/lite/tools/benchmark/benchmark_tflite_model.cc
index ef4f0fa..463d599 100644
--- a/tensorflow/contrib/lite/tools/benchmark/benchmark_tflite_model.cc
+++ b/tensorflow/contrib/lite/tools/benchmark/benchmark_tflite_model.cc
@@ -23,9 +23,9 @@
#include <unordered_set>
#include <vector>
-#ifdef TFLITE_EXTENDED
-#include "tensorflow/contrib/lite/delegates/eager/delegate.h"
-#endif // TFLITE_EXTENDED
+#ifdef TFLITE_FLEX
+#include "tensorflow/contrib/lite/delegates/flex/delegate.h"
+#endif // TFLITE_FLEX
#include "tensorflow/contrib/lite/kernels/register.h"
#include "tensorflow/contrib/lite/model.h"
#include "tensorflow/contrib/lite/op_resolver.h"
@@ -305,14 +305,14 @@
interpreter->UseNNAPI(use_nnapi);
-#ifdef TFLITE_EXTENDED
- TFLITE_LOG(INFO) << "Instantiating Eager Delegate";
- delegate_ = EagerDelegate::Create();
+#ifdef TFLITE_FLEX
+ TFLITE_LOG(INFO) << "Instantiating Flex Delegate";
+ delegate_ = FlexDelegate::Create();
if (delegate_) {
interpreter->ModifyGraphWithDelegate(delegate_.get(),
/*allow_dynamic_tensors=*/true);
}
-#endif // TFLITE_EXTENDED
+#endif // TFLITE_FLEX
auto interpreter_inputs = interpreter->inputs();
diff --git a/tensorflow/contrib/lite/tools/benchmark/benchmark_tflite_model.h b/tensorflow/contrib/lite/tools/benchmark/benchmark_tflite_model.h
index 8541512..b091e18 100644
--- a/tensorflow/contrib/lite/tools/benchmark/benchmark_tflite_model.h
+++ b/tensorflow/contrib/lite/tools/benchmark/benchmark_tflite_model.h
@@ -20,9 +20,9 @@
#include <string>
#include <vector>
-#ifdef TFLITE_EXTENDED
-#include "tensorflow/contrib/lite/delegates/eager/delegate.h"
-#endif // TFLITE_EXTENDED
+#ifdef TFLITE_FLEX
+#include "tensorflow/contrib/lite/delegates/flex/delegate.h"
+#endif // TFLITE_FLEX
#include "tensorflow/contrib/lite/model.h"
#include "tensorflow/contrib/lite/profiling/profile_summarizer.h"
#include "tensorflow/contrib/lite/tools/benchmark/benchmark_model.h"
@@ -73,9 +73,9 @@
void PrepareInputsAndOutputs() override;
private:
-#ifdef TFLITE_EXTENDED
- std::unique_ptr<EagerDelegate> delegate_;
-#endif // TFLITE_EXTENDED
+#ifdef TFLITE_FLEX
+ std::unique_ptr<FlexDelegate> delegate_;
+#endif // TFLITE_FLEX
std::unique_ptr<tflite::FlatBufferModel> model;
std::unique_ptr<tflite::Interpreter> interpreter;
std::vector<InputLayerInfo> inputs;
diff --git a/tensorflow/contrib/lite/util.cc b/tensorflow/contrib/lite/util.cc
index 7950653..6aa35b5 100644
--- a/tensorflow/contrib/lite/util.cc
+++ b/tensorflow/contrib/lite/util.cc
@@ -18,9 +18,9 @@
namespace tflite {
-bool IsEagerOp(const char* custom_name) {
- return custom_name && strncmp(custom_name, kEagerCustomCodePrefix,
- strlen(kEagerCustomCodePrefix)) == 0;
+bool IsFlexOp(const char* custom_name) {
+ return custom_name && strncmp(custom_name, kFlexCustomCodePrefix,
+ strlen(kFlexCustomCodePrefix)) == 0;
}
TfLiteIntArray* ConvertVectorToTfLiteIntArray(const std::vector<int>& input) {
diff --git a/tensorflow/contrib/lite/util.h b/tensorflow/contrib/lite/util.h
index 6d81f84..31292a6 100644
--- a/tensorflow/contrib/lite/util.h
+++ b/tensorflow/contrib/lite/util.h
@@ -26,15 +26,15 @@
namespace tflite {
-// The prefix of Eager op custom code.
+// The prefix of Flex op custom code.
// This will be matched agains the `custom_code` field in `OperatorCode`
// Flatbuffer Table.
// WARNING: This is an experimental API and subject to change.
-constexpr char kEagerCustomCodePrefix[] = "Eager";
+constexpr char kFlexCustomCodePrefix[] = "Flex";
// Checks whether the prefix of the custom name indicates the operation is an
-// Eager operation.
-bool IsEagerOp(const char* custom_name);
+// Flex operation.
+bool IsFlexOp(const char* custom_name);
// Converts a `std::vector` to a `TfLiteIntArray`. The caller takes ownership
// of the returned pointer.
diff --git a/tensorflow/contrib/lite/util_test.cc b/tensorflow/contrib/lite/util_test.cc
index c5c1709..25f3ade 100644
--- a/tensorflow/contrib/lite/util_test.cc
+++ b/tensorflow/contrib/lite/util_test.cc
@@ -41,14 +41,14 @@
TfLiteIntArrayFree(output);
}
-TEST(UtilTest, IsEagerOp) {
- EXPECT_TRUE(IsEagerOp("Eager"));
- EXPECT_TRUE(IsEagerOp("EagerOp"));
- EXPECT_FALSE(IsEagerOp("eager"));
- EXPECT_FALSE(IsEagerOp("Eage"));
- EXPECT_FALSE(IsEagerOp("OpEager"));
- EXPECT_FALSE(IsEagerOp(nullptr));
- EXPECT_FALSE(IsEagerOp(""));
+TEST(UtilTest, IsFlexOp) {
+ EXPECT_TRUE(IsFlexOp("Flex"));
+ EXPECT_TRUE(IsFlexOp("FlexOp"));
+ EXPECT_FALSE(IsFlexOp("flex"));
+ EXPECT_FALSE(IsFlexOp("Fle"));
+ EXPECT_FALSE(IsFlexOp("OpFlex"));
+ EXPECT_FALSE(IsFlexOp(nullptr));
+ EXPECT_FALSE(IsFlexOp(""));
}
} // namespace
diff --git a/tensorflow/contrib/lookup/lookup_ops_test.py b/tensorflow/contrib/lookup/lookup_ops_test.py
index 89b538d..9e9345e 100644
--- a/tensorflow/contrib/lookup/lookup_ops_test.py
+++ b/tensorflow/contrib/lookup/lookup_ops_test.py
@@ -23,8 +23,8 @@
import six
from tensorflow.contrib import lookup
-from tensorflow.contrib.data.python.ops import counter
from tensorflow.python.client import session
+from tensorflow.python.data.experimental.ops import counter
from tensorflow.python.eager import context
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
diff --git a/tensorflow/contrib/makefile/Makefile b/tensorflow/contrib/makefile/Makefile
index d962a5e..36125c1 100644
--- a/tensorflow/contrib/makefile/Makefile
+++ b/tensorflow/contrib/makefile/Makefile
@@ -133,7 +133,8 @@
$(wildcard tensorflow/contrib/makefile/downloads/absl/absl/*/*/*benchmark*.cc) \
$(wildcard tensorflow/contrib/makefile/downloads/absl/absl/*/*/*/*benchmark*.cc) \
$(wildcard tensorflow/contrib/makefile/downloads/absl/absl/*/*/*/*/*benchmark*.cc) \
-tensorflow/contrib/makefile/downloads/absl/absl/synchronization/internal/mutex_nonprod.cc
+tensorflow/contrib/makefile/downloads/absl/absl/synchronization/internal/mutex_nonprod.cc \
+tensorflow/contrib/makefile/downloads/absl/absl/hash/internal/print_hash_of.cc
ABSL_CC_SRCS := $(filter-out $(ABSL_CC_EXCLUDE_SRCS), $(ABSL_CC_ALL_SRCS))
diff --git a/tensorflow/contrib/makefile/tf_op_files.txt b/tensorflow/contrib/makefile/tf_op_files.txt
index f81a908..91af933 100644
--- a/tensorflow/contrib/makefile/tf_op_files.txt
+++ b/tensorflow/contrib/makefile/tf_op_files.txt
@@ -91,6 +91,8 @@
tensorflow/core/kernels/cwise_op_squared_difference.cc
tensorflow/core/kernels/cwise_op_sub.cc
tensorflow/core/kernels/cwise_op_tanh.cc
+tensorflow/core/kernels/cwise_op_xdivy.cc
+tensorflow/core/kernels/cwise_op_xlogy.cc
tensorflow/core/kernels/cwise_ops_common.cc
tensorflow/core/kernels/data_format_ops.cc
tensorflow/core/kernels/decode_bmp_op.cc
diff --git a/tensorflow/contrib/metrics/python/ops/metric_ops_test.py b/tensorflow/contrib/metrics/python/ops/metric_ops_test.py
index 955b83b..fc64f34 100644
--- a/tensorflow/contrib/metrics/python/ops/metric_ops_test.py
+++ b/tensorflow/contrib/metrics/python/ops/metric_ops_test.py
@@ -2069,11 +2069,11 @@
num_batches = 100
labels = np.array([])
predictions = np.array([])
- tf_labels = variables.Variable(
+ tf_labels = variables.VariableV1(
array_ops.ones(batch_size, dtypes_lib.int32),
collections=[ops.GraphKeys.LOCAL_VARIABLES],
dtype=dtypes_lib.int32)
- tf_predictions = variables.Variable(
+ tf_predictions = variables.VariableV1(
array_ops.ones(batch_size),
collections=[ops.GraphKeys.LOCAL_VARIABLES],
dtype=dtypes_lib.float32)
@@ -2133,15 +2133,15 @@
labels = np.array([])
predictions = np.array([])
weights = np.array([])
- tf_labels = variables.Variable(
+ tf_labels = variables.VariableV1(
array_ops.ones(batch_size, dtypes_lib.int32),
collections=[ops.GraphKeys.LOCAL_VARIABLES],
dtype=dtypes_lib.int32)
- tf_predictions = variables.Variable(
+ tf_predictions = variables.VariableV1(
array_ops.ones(batch_size),
collections=[ops.GraphKeys.LOCAL_VARIABLES],
dtype=dtypes_lib.float32)
- tf_weights = variables.Variable(
+ tf_weights = variables.VariableV1(
array_ops.ones(batch_size),
collections=[ops.GraphKeys.LOCAL_VARIABLES],
dtype=dtypes_lib.float32)
@@ -2311,10 +2311,11 @@
num_batches = 100
labels = np.array([])
predictions = np.array([])
- tf_labels = variables.Variable(array_ops.ones(batch_size, dtypes_lib.int32),
- collections=[ops.GraphKeys.LOCAL_VARIABLES],
- dtype=dtypes_lib.int32)
- tf_predictions = variables.Variable(
+ tf_labels = variables.VariableV1(
+ array_ops.ones(batch_size, dtypes_lib.int32),
+ collections=[ops.GraphKeys.LOCAL_VARIABLES],
+ dtype=dtypes_lib.int32)
+ tf_predictions = variables.VariableV1(
array_ops.ones(batch_size),
collections=[ops.GraphKeys.LOCAL_VARIABLES],
dtype=dtypes_lib.float32)
diff --git a/tensorflow/contrib/model_pruning/python/pruning.py b/tensorflow/contrib/model_pruning/python/pruning.py
index a81abac..67e58ff 100644
--- a/tensorflow/contrib/model_pruning/python/pruning.py
+++ b/tensorflow/contrib/model_pruning/python/pruning.py
@@ -247,7 +247,8 @@
# Stores the tensorflow sparsity variable.
# Built using self._setup_sparsity() or provided externally
- self._sparsity = sparsity if sparsity else self._setup_sparsity()
+ self._sparsity = (sparsity
+ if sparsity is not None else self._setup_sparsity())
# List of tensorflow assignments ops for new masks and thresholds
self._assign_ops = []
diff --git a/tensorflow/contrib/model_pruning/python/pruning_test.py b/tensorflow/contrib/model_pruning/python/pruning_test.py
index cd3d8e7..1b6da5c 100644
--- a/tensorflow/contrib/model_pruning/python/pruning_test.py
+++ b/tensorflow/contrib/model_pruning/python/pruning_test.py
@@ -45,7 +45,7 @@
# Add global step variable to the graph
self.global_step = training_util.get_or_create_global_step()
# Add sparsity
- self.sparsity = variables.Variable(0.5, name="sparsity")
+ self.sparsity = variables.VariableV1(0.5, name="sparsity")
# Parse hparams
self.pruning_hparams = pruning.get_pruning_hparams().parse(
self.TEST_HPARAMS)
@@ -88,7 +88,7 @@
width = 10
height = 20
with self.cached_session():
- weights = variables.Variable(
+ weights = variables.VariableV1(
random_ops.random_normal([width, height], stddev=1), name="weights")
masked_weights = pruning.apply_mask(weights,
variable_scope.get_variable_scope())
@@ -99,10 +99,10 @@
def testUpdateSingleMask(self):
with self.cached_session() as session:
- weights = variables.Variable(
+ weights = variables.VariableV1(
math_ops.linspace(1.0, 100.0, 100), name="weights")
masked_weights = pruning.apply_mask(weights)
- sparsity = variables.Variable(0.5, name="sparsity")
+ sparsity = variables.VariableV1(0.5, name="sparsity")
p = pruning.Pruning(sparsity=sparsity)
p._spec.threshold_decay = 0.0
mask_update_op = p.mask_update_op()
@@ -115,8 +115,8 @@
def _blockMasking(self, hparams, weights, expected_mask):
- threshold = variables.Variable(0.0, name="threshold")
- sparsity = variables.Variable(0.5, name="sparsity")
+ threshold = variables.VariableV1(0.0, name="threshold")
+ sparsity = variables.VariableV1(0.5, name="sparsity")
test_spec = ",".join(hparams)
pruning_hparams = pruning.get_pruning_hparams().parse(test_spec)
@@ -169,7 +169,7 @@
partitioner = partitioned_variables.variable_axis_size_partitioner(40)
with self.cached_session() as session:
with variable_scope.variable_scope("", partitioner=partitioner):
- sparsity = variables.Variable(0.5, name="Sparsity")
+ sparsity = variables.VariableV1(0.5, name="Sparsity")
weights = variable_scope.get_variable(
"weights", initializer=math_ops.linspace(1.0, 100.0, 100))
masked_weights = pruning.apply_mask(
@@ -190,10 +190,10 @@
]
test_spec = ",".join(param_list)
pruning_hparams = pruning.get_pruning_hparams().parse(test_spec)
- weights = variables.Variable(
+ weights = variables.VariableV1(
math_ops.linspace(1.0, 100.0, 100), name="weights")
masked_weights = pruning.apply_mask(weights)
- sparsity = variables.Variable(0.00, name="sparsity")
+ sparsity = variables.VariableV1(0.00, name="sparsity")
# Set up pruning
p = pruning.Pruning(pruning_hparams, sparsity=sparsity)
p._spec.threshold_decay = 0.0
@@ -222,11 +222,11 @@
pruning_hparams = pruning.get_pruning_hparams().parse(test_spec)
with variable_scope.variable_scope("layer1"):
- w1 = variables.Variable(
+ w1 = variables.VariableV1(
math_ops.linspace(1.0, 100.0, 100), name="weights")
_ = pruning.apply_mask(w1)
with variable_scope.variable_scope("layer2"):
- w2 = variables.Variable(
+ w2 = variables.VariableV1(
math_ops.linspace(1.0, 100.0, 100), name="weights")
_ = pruning.apply_mask(w2)
diff --git a/tensorflow/contrib/opt/python/training/addsign_test.py b/tensorflow/contrib/opt/python/training/addsign_test.py
index 628a735..6150fa1 100644
--- a/tensorflow/contrib/opt/python/training/addsign_test.py
+++ b/tensorflow/contrib/opt/python/training/addsign_test.py
@@ -80,9 +80,9 @@
global_step = resource_variable_ops.ResourceVariable(
0, trainable=False)
else:
- var0 = variables.Variable(var0_np)
- var1 = variables.Variable(var1_np)
- global_step = variables.Variable(
+ var0 = variables.VariableV1(var0_np)
+ var1 = variables.VariableV1(var1_np)
+ global_step = variables.VariableV1(
0, trainable=False)
grads0 = constant_op.constant(grads0_np)
grads1 = constant_op.constant(grads1_np)
@@ -183,9 +183,9 @@
global_step = resource_variable_ops.ResourceVariable(
0, trainable=False)
else:
- var0 = variables.Variable(var0_np)
- var1 = variables.Variable(var1_np)
- global_step = variables.Variable(
+ var0 = variables.VariableV1(var0_np)
+ var1 = variables.VariableV1(var1_np)
+ global_step = variables.VariableV1(
0, trainable=False)
grads0_np_indices = np.array([0, 1], dtype=np.int32)
grads0 = ops.IndexedSlices(
diff --git a/tensorflow/contrib/opt/python/training/drop_stale_gradient_optimizer_test.py b/tensorflow/contrib/opt/python/training/drop_stale_gradient_optimizer_test.py
index 5323208..0a69096 100644
--- a/tensorflow/contrib/opt/python/training/drop_stale_gradient_optimizer_test.py
+++ b/tensorflow/contrib/opt/python/training/drop_stale_gradient_optimizer_test.py
@@ -61,8 +61,8 @@
graph = ops.Graph()
with graph.as_default():
global_step = training_util.create_global_step()
- var_0 = variables.Variable(0.0, name='v0')
- var_1 = variables.Variable(1.0, name='v1')
+ var_0 = variables.VariableV1(0.0, name='v0')
+ var_1 = variables.VariableV1(1.0, name='v1')
compute_gradients_queue = data_flow_ops.FIFOQueue(
-1, global_step.dtype.base_dtype, shapes=(),
name='compute_gradients_queue', shared_name='compute_gradients_queue')
diff --git a/tensorflow/contrib/opt/python/training/external_optimizer_test.py b/tensorflow/contrib/opt/python/training/external_optimizer_test.py
index 9997103..70c5f8f 100644
--- a/tensorflow/contrib/opt/python/training/external_optimizer_test.py
+++ b/tensorflow/contrib/opt/python/training/external_optimizer_test.py
@@ -69,9 +69,9 @@
class ExternalOptimizerInterfaceTest(TestCase):
def test_optimize(self):
- scalar = variables.Variable(random_ops.random_normal([]), 'scalar')
- vector = variables.Variable(random_ops.random_normal([2]), 'vector')
- matrix = variables.Variable(random_ops.random_normal([2, 3]), 'matrix')
+ scalar = variables.VariableV1(random_ops.random_normal([]), 'scalar')
+ vector = variables.VariableV1(random_ops.random_normal([2]), 'vector')
+ matrix = variables.VariableV1(random_ops.random_normal([2, 3]), 'matrix')
minimum_location = constant_op.constant(np.arange(9), dtype=dtypes.float32)
@@ -96,7 +96,7 @@
def test_callbacks(self):
vector_val = np.array([7., -2.], dtype=np.float32)
- vector = variables.Variable(vector_val, 'vector')
+ vector = variables.VariableV1(vector_val, 'vector')
minimum_location_val = np.arange(2)
minimum_location = constant_op.constant(
@@ -160,7 +160,7 @@
rtol=1e-5,
atol=1e-5,
dimension=5):
- x = variables.Variable(array_ops.zeros(dimension))
+ x = variables.VariableV1(array_ops.zeros(dimension))
optimizer = external_optimizer.ScipyOptimizerInterface(
self._objective(x), method=method, options=options)
@@ -173,7 +173,7 @@
def test_unconstrained(self):
dimension = 5
- x = variables.Variable(array_ops.zeros(dimension))
+ x = variables.VariableV1(array_ops.zeros(dimension))
optimizer = external_optimizer.ScipyOptimizerInterface(self._objective(x))
with self.cached_session() as sess:
@@ -230,7 +230,7 @@
def test_nonlinear_programming(self):
vector_initial_value = [7., 7.]
- vector = variables.Variable(vector_initial_value, 'vector')
+ vector = variables.VariableV1(vector_initial_value, 'vector')
# Make norm as small as possible.
loss = math_ops.reduce_sum(math_ops.square(vector))
@@ -249,7 +249,7 @@
def test_scalar_bounds(self):
vector_initial_value = [7., 7.]
- vector = variables.Variable(vector_initial_value, 'vector')
+ vector = variables.VariableV1(vector_initial_value, 'vector')
# Make norm as small as possible.
loss = math_ops.reduce_sum(math_ops.square(vector))
@@ -267,7 +267,7 @@
def test_vector_bounds(self):
vector_initial_value = [7., 7.]
- vector = variables.Variable(vector_initial_value, 'vector')
+ vector = variables.VariableV1(vector_initial_value, 'vector')
# Make norm as small as possible.
loss = math_ops.reduce_sum(math_ops.square(vector))
@@ -287,7 +287,7 @@
# after running optimizer.minimize().
# Bug reference: b/64065260
vector_initial_value = [7., 7.]
- vector = variables.Variable(vector_initial_value, 'vector')
+ vector = variables.VariableV1(vector_initial_value, 'vector')
loss = math_ops.reduce_sum(math_ops.square(vector))
optimizer = external_optimizer.ScipyOptimizerInterface(
@@ -301,7 +301,7 @@
def test_callbacks(self):
vector_val = np.array([7., -2.], dtype=np.float32)
- vector = variables.Variable(vector_val, 'vector')
+ vector = variables.VariableV1(vector_val, 'vector')
minimum_location_val = np.arange(2)
minimum_location = constant_op.constant(
diff --git a/tensorflow/contrib/opt/python/training/model_average_optimizer_test.py b/tensorflow/contrib/opt/python/training/model_average_optimizer_test.py
index b1fc50a..a25455e 100644
--- a/tensorflow/contrib/opt/python/training/model_average_optimizer_test.py
+++ b/tensorflow/contrib/opt/python/training/model_average_optimizer_test.py
@@ -110,10 +110,11 @@
class ModelAverageOptimizerTest(test.TestCase):
+
def _run(self, train_op, sess):
sess.run(train_op)
- def test1Workers2Period(self):
+ def disabled_test1Workers2Period(self):
num_workers = 2
steps = 2
num_ps = 1
diff --git a/tensorflow/contrib/opt/python/training/powersign_test.py b/tensorflow/contrib/opt/python/training/powersign_test.py
index 0bcf5d2..1cf9901 100644
--- a/tensorflow/contrib/opt/python/training/powersign_test.py
+++ b/tensorflow/contrib/opt/python/training/powersign_test.py
@@ -81,9 +81,9 @@
global_step = resource_variable_ops.ResourceVariable(
0, trainable=False)
else:
- var0 = variables.Variable(var0_np)
- var1 = variables.Variable(var1_np)
- global_step = variables.Variable(
+ var0 = variables.VariableV1(var0_np)
+ var1 = variables.VariableV1(var1_np)
+ global_step = variables.VariableV1(
0, trainable=False)
grads0 = constant_op.constant(grads0_np)
grads1 = constant_op.constant(grads1_np)
@@ -188,9 +188,9 @@
global_step = resource_variable_ops.ResourceVariable(
0, trainable=False)
else:
- var0 = variables.Variable(var0_np)
- var1 = variables.Variable(var1_np)
- global_step = variables.Variable(
+ var0 = variables.VariableV1(var0_np)
+ var1 = variables.VariableV1(var1_np)
+ global_step = variables.VariableV1(
0, trainable=False)
grads0_np_indices = np.array([0, 1], dtype=np.int32)
grads0 = ops.IndexedSlices(
diff --git a/tensorflow/contrib/opt/python/training/shampoo_test.py b/tensorflow/contrib/opt/python/training/shampoo_test.py
index 05bcf2c..a2fd8fb 100644
--- a/tensorflow/contrib/opt/python/training/shampoo_test.py
+++ b/tensorflow/contrib/opt/python/training/shampoo_test.py
@@ -54,9 +54,9 @@
grad_np_2 = np.random.rand(size)
with self.cached_session() as sess:
- global_step = variables.Variable(
+ global_step = variables.VariableV1(
0, dtype=dtypes.int64, use_resource=use_resource_var)
- var = variables.Variable(
+ var = variables.VariableV1(
init_var_np, dtype=dtypes.float32, use_resource=use_resource_var)
grad = constant_op.constant(grad_np, dtype=dtypes.float32)
grad_2 = constant_op.constant(grad_np_2, dtype=dtypes.float32)
@@ -105,9 +105,9 @@
grad_np_2 = np.random.rand(size[0], size[1])
with self.cached_session() as sess:
- global_step = variables.Variable(
+ global_step = variables.VariableV1(
0, dtype=dtypes.int64, use_resource=use_resource_var)
- var = variables.Variable(
+ var = variables.VariableV1(
init_var_np, dtype=dtypes.float32, use_resource=use_resource_var)
grad = constant_op.constant(grad_np, dtype=dtypes.float32)
grad_2 = constant_op.constant(grad_np_2, dtype=dtypes.float32)
@@ -164,9 +164,9 @@
grad_np_2 = np.random.rand(size[0], size[1], size[2])
with self.cached_session() as sess:
- global_step = variables.Variable(
+ global_step = variables.VariableV1(
0, dtype=dtypes.int64, use_resource=use_resource_var)
- var = variables.Variable(
+ var = variables.VariableV1(
init_var_np, dtype=dtypes.float32, use_resource=use_resource_var)
grad = constant_op.constant(grad_np, dtype=dtypes.float32)
grad_2 = constant_op.constant(grad_np_2, dtype=dtypes.float32)
@@ -254,9 +254,9 @@
grad_np_2 = np.random.rand(size)
with self.cached_session() as sess:
- global_step = variables.Variable(
+ global_step = variables.VariableV1(
0, dtype=dtypes.int64, use_resource=use_resource_var)
- var = variables.Variable(
+ var = variables.VariableV1(
init_var_np, dtype=dtypes.float32, use_resource=use_resource_var)
grad = constant_op.constant(grad_np, dtype=dtypes.float32)
grad_2 = constant_op.constant(grad_np_2, dtype=dtypes.float32)
@@ -310,9 +310,9 @@
grad_np_2 = np.random.rand(size[0], size[1])
with self.cached_session() as sess:
- global_step = variables.Variable(
+ global_step = variables.VariableV1(
0, dtype=dtypes.int64, use_resource=use_resource_var)
- var = variables.Variable(
+ var = variables.VariableV1(
init_var_np, dtype=dtypes.float32, use_resource=use_resource_var)
grad = constant_op.constant(grad_np, dtype=dtypes.float32)
grad_2 = constant_op.constant(grad_np_2, dtype=dtypes.float32)
@@ -383,9 +383,9 @@
grad_np_2 = np.random.rand(sample_size_2, size[1])
with self.cached_session() as sess:
- global_step = variables.Variable(
+ global_step = variables.VariableV1(
0, dtype=dtypes.int64, use_resource=use_resource_var)
- var = variables.Variable(
+ var = variables.VariableV1(
init_var_np, dtype=dtypes.float32, use_resource=use_resource_var)
grad = ops.IndexedSlices(
constant_op.constant(grad_np, dtype=dtypes.float32),
@@ -463,9 +463,9 @@
grad_np = np.random.rand(sample_size, size[1], size[2])
with self.cached_session() as sess:
- global_step = variables.Variable(
+ global_step = variables.VariableV1(
0, dtype=dtypes.int64, use_resource=use_resource_var)
- var = variables.Variable(
+ var = variables.VariableV1(
init_var_np, dtype=dtypes.float32, use_resource=use_resource_var)
grad = ops.IndexedSlices(
constant_op.constant(grad_np, dtype=dtypes.float32),
@@ -533,9 +533,9 @@
gbar_weight = 0.1
with self.cached_session() as sess:
- global_step = variables.Variable(
+ global_step = variables.VariableV1(
0, dtype=dtypes.int64, use_resource=use_resource_var)
- var = variables.Variable(
+ var = variables.VariableV1(
init_var_np, dtype=dtypes.float32, use_resource=use_resource_var)
grad = constant_op.constant(grad_np, dtype=dtypes.float32)
grad_2 = constant_op.constant(grad_np_2, dtype=dtypes.float32)
@@ -628,9 +628,9 @@
mat_g3 = np.zeros_like(mat_g3_a)
with self.cached_session() as sess:
- global_step = variables.Variable(
+ global_step = variables.VariableV1(
0, dtype=dtypes.int64, use_resource=use_resource_var)
- var = variables.Variable(
+ var = variables.VariableV1(
init_var_np, dtype=dtypes.float32, use_resource=use_resource_var)
grad = array_ops.placeholder(dtypes.float32, shape=size)
@@ -705,9 +705,9 @@
mat_g3 = np.zeros_like(mat_g3_a)
with self.cached_session() as sess:
- global_step = variables.Variable(
+ global_step = variables.VariableV1(
0, dtype=dtypes.int64, use_resource=use_resource_var)
- var = variables.Variable(
+ var = variables.VariableV1(
init_var_np, dtype=dtypes.float32, use_resource=use_resource_var)
grad = array_ops.placeholder(dtypes.float32, shape=size)
diff --git a/tensorflow/contrib/optimizer_v2/optimizer_v2.py b/tensorflow/contrib/optimizer_v2/optimizer_v2.py
index 6af59dc..53e27c0 100644
--- a/tensorflow/contrib/optimizer_v2/optimizer_v2.py
+++ b/tensorflow/contrib/optimizer_v2/optimizer_v2.py
@@ -30,7 +30,6 @@
from tensorflow.python.ops import gradients
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.ops import variable_scope
from tensorflow.python.ops import variables
from tensorflow.python.training import distribute as distribute_lib
@@ -965,8 +964,7 @@
# Use the processors to update the variables.
update_ops = []
for grad, var in grads_and_vars:
- update_ops.extend(distribution.unwrap(distribution.update(
- var, update, grad)))
+ update_ops.extend(distribution.update(var, update, grad, grouped=False))
# Give the child class a chance to do something after applying
# gradients
@@ -978,26 +976,24 @@
update_ops = control_flow_ops.group(update_ops)
with ops.control_dependencies([update_ops]):
- finish_updates = distribution.update_non_slot(non_slot_devices, finish)
- if finish_updates is None:
- finish_updates = update_ops
+ finish_updates = distribution.update_non_slot(
+ non_slot_devices, finish, grouped=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]
# Update `global_step` (if any).
if global_step is None:
apply_updates = distribution.group(finish_updates, name=name)
else:
- with ops.control_dependencies(distribution.unwrap(finish_updates)):
+ with ops.control_dependencies(finish_updates):
- def update_global_step(global_step):
- if isinstance(global_step, resource_variable_ops.ResourceVariable):
- return global_step.assign_add(
- ops.convert_to_tensor(1, dtype=global_step.dtype),
- read_value=False)
- else:
- return state_ops.assign_add(global_step, 1)
+ def update_global_step(global_step, name):
+ return global_step.assign_add(1, read_value=False, name=name)
- apply_updates = distribution.group(
- distribution.update(global_step, update_global_step), name=name)
+ apply_updates = distribution.update(
+ global_step, update_global_step, name)
# Add the training op to the TRAIN_OP graph collection in graph mode.
if not eager_execution:
diff --git a/tensorflow/contrib/quantize/python/fold_batch_norms.py b/tensorflow/contrib/quantize/python/fold_batch_norms.py
index e5790a6..7575b1b 100644
--- a/tensorflow/contrib/quantize/python/fold_batch_norms.py
+++ b/tensorflow/contrib/quantize/python/fold_batch_norms.py
@@ -418,10 +418,11 @@
transpose_b=layer_op.get_attr('transpose_b'),
name=new_layer_name)
elif layer_op.type == 'DepthwiseConv2dNative':
+ # We don't copy dilation rate because we reuse the input SpaceToBatch
+ # and create our own BatchToSpace operation below.
conv = nn.depthwise_conv2d(
input_tensor,
weight_tensor,
- rate=layer_op.get_attr('dilations'),
strides=layer_op.get_attr('strides'),
padding=layer_op.get_attr('padding'),
name=new_layer_name)
diff --git a/tensorflow/contrib/quantize/python/quantize.py b/tensorflow/contrib/quantize/python/quantize.py
index 5e63d33..afb9de8 100644
--- a/tensorflow/contrib/quantize/python/quantize.py
+++ b/tensorflow/contrib/quantize/python/quantize.py
@@ -461,8 +461,8 @@
return self._bias_add_op
-def _FollowedByFakeQuant(tensor):
- """Returns True if the tensor is followed by a FakeQuant."""
+def _GetFollowingFakeQuantOp(tensor):
+ """Returns the following FakeQuant op if it exists else None."""
fake_quant_ops = set([
'FakeQuantWithMinMaxVars', 'FakeQuantWithMinMaxArgs',
'FakeQuantWithMinMaxVarsPerChannel'
@@ -472,11 +472,11 @@
while consumers:
c = consumers.pop()
if c.type in fake_quant_ops:
- return True
+ return c
elif c.type in pass_through_ops:
for output in c.outputs:
consumers.extend(output.consumers())
- return False
+ return None
def _InsertQuantOp(context,
@@ -559,44 +559,77 @@
# Prevent ops from being quantized multiple times. Bypass ops can sometimes
# overlap between multiple matches, so we need to ensure that we don't
# add duplicate FakeQuant operations.
- if _FollowedByFakeQuant(inputs):
- return
+ fake_quant_op = _GetFollowingFakeQuantOp(inputs)
- if moving_avg:
- quant = (
- quant_ops.MovingAvgQuantize(
- inputs,
- init_min=init_min,
- init_max=init_max,
- ema_decay=ema_decay,
- is_training=is_training,
- num_bits=bits,
- narrow_range=narrow_range,
- vars_collection=vars_collection,
- name_prefix=name_prefix))
+ # If we find that we are attempting to insert a fake quant op following
+ # a fake quant, we skip inserting a fake quant op
+
+ if fake_quant_op is None:
+ if moving_avg:
+ quant = (
+ quant_ops.MovingAvgQuantize(
+ inputs,
+ init_min=init_min,
+ init_max=init_max,
+ ema_decay=ema_decay,
+ is_training=is_training,
+ num_bits=bits,
+ narrow_range=narrow_range,
+ vars_collection=vars_collection,
+ name_prefix=name_prefix))
+ else:
+ quant = (
+ quant_ops.LastValueQuantize(
+ inputs,
+ init_min=init_min,
+ init_max=init_max,
+ is_training=is_training,
+ num_bits=bits,
+ narrow_range=narrow_range,
+ vars_collection=vars_collection,
+ name_prefix=name_prefix))
+
+ if quant_delay and quant_delay > 0:
+ activate_quant = math_ops.greater_equal(
+ common.CreateOrGetQuantizationStep(),
+ quant_delay,
+ name=name_prefix + '/activate_quant')
+ quant = control_flow_ops.cond(
+ activate_quant,
+ lambda: quant,
+ lambda: inputs,
+ name=name_prefix + '/delayed_quant')
else:
- quant = (
- quant_ops.LastValueQuantize(
- inputs,
- init_min=init_min,
- init_max=init_max,
- is_training=is_training,
- num_bits=bits,
- narrow_range=narrow_range,
- vars_collection=vars_collection,
- name_prefix=name_prefix))
+ # If a fake quant op is present already, make sure that
+ # any downstream use of the tensor reroutes to the appropriate quantized
+ # tensor. If there is no quant_delay, this is simply the output of the
+ # fake quant op. If there is a quant delay, we reroute to the output
+ # of the delayed quant operation, which inserts quantization only after
+ # a specified quant_delay
- if quant_delay and quant_delay > 0:
- activate_quant = math_ops.greater_equal(
- common.CreateOrGetQuantizationStep(),
- quant_delay,
- name=name_prefix + '/activate_quant')
- quant = control_flow_ops.cond(
- activate_quant,
- lambda: quant,
- lambda: inputs,
- name=name_prefix + '/delayed_quant')
+ quant = fake_quant_op.outputs[0]
+ if quant_delay and quant_delay > 0:
+ name_prefix = '/'.join(quant.name.split('/')[:-1])
+ quant = quant.graph.get_tensor_by_name(name_prefix +
+ '/delayed_quant/Merge:0')
+ pruned_consumer_set = set()
+ for consumer in consumers:
+ fake_quant_dest_op = _GetFollowingFakeQuantOp(consumer.outputs[0])
+ if (fake_quant_dest_op is None or
+ fake_quant_dest_op.name != fake_quant_op.name):
+ pruned_consumer_set.add(consumer)
+ consumers = pruned_consumer_set
+ # If we have
+ # input->pass_through->fake_quant
+ # there is nothing to reroute.
+ #
+ # If we have
+ # input-> pass_through->fake_quant
+ # |-> consumer
+ # Then we reroute such that:
+ # input-> pass_through->fake_quant
+ # |-> consumer
if consumers:
tensors_modified_count = common.RerouteTensor(
quant, inputs, can_modify=consumers)
diff --git a/tensorflow/contrib/quantize/python/quantize_graph_test.py b/tensorflow/contrib/quantize/python/quantize_graph_test.py
index e80d218..a9fc6c3 100644
--- a/tensorflow/contrib/quantize/python/quantize_graph_test.py
+++ b/tensorflow/contrib/quantize/python/quantize_graph_test.py
@@ -27,6 +27,7 @@
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn_ops
+from tensorflow.python.ops import template
from tensorflow.python.platform import googletest
@@ -306,6 +307,42 @@
# No ops should be inserted or removed.
self.assertEqual(op_names_before_rewrite, op_names_after_rewrite)
+ def testWithSharedWeights(self):
+
+ self._RunTestOverAllRewrites(self._TestWithSharedWeights)
+ self._RunTestOverTrainingRewrites(self._TestRewriteWithSharedWeights)
+
+ def _TestRewriteWithSharedWeights(self, rewrite_fn, quant_delay=1):
+ self._TestWithSharedWeights(rewrite_fn, quant_delay)
+
+ def _TestWithSharedWeights(self, rewrite_fn, quant_delay=None):
+ with ops.Graph().as_default() as g:
+ conv = template.make_template('shared_weights_conv', self._ConvLayer)
+ conv()
+ conv()
+ if quant_delay is None:
+ rewrite_fn()
+ else:
+ rewrite_fn(quant_delay=quant_delay)
+
+ conv_ops = [op for op in g.get_operations() if op.type == 'Conv2D']
+ weights_quants = [
+ op for op in g.get_operations()
+ if 'weights_quant' in op.name and op.type == 'FakeQuantWithMinMaxVars'
+ ]
+ # Check that the shared weights variable is not quantized multiple times
+ self.assertTrue(len(weights_quants) == 1)
+ weights_quant_tensor = weights_quants[0].outputs[0]
+ if quant_delay:
+ delayed_weights_quants = [
+ op for op in g.get_operations()
+ if 'weights_quant' in op.name and op.type == 'Merge'
+ ]
+ self.assertTrue(len(delayed_weights_quants) == 1)
+ weights_quant_tensor = delayed_weights_quants[0].outputs[0]
+ # Check that the Conv2D operations get the quantized weights
+ self.assertTrue(all(weights_quant_tensor in op.inputs for op in conv_ops))
+
def _ConvLayer(
self, input_tensor=None, scope='test', pre_activation_bypass=False,
post_activation_bypass=False):
diff --git a/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_test.py b/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_test.py
index bf699db..f31ad53 100644
--- a/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_test.py
+++ b/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_test.py
@@ -163,8 +163,8 @@
def __init__(self, batch_size, state_size):
super(TestStateSaverWithCounters, self).__init__(batch_size, state_size)
- self._num_state_calls = variables_lib.Variable(0)
- self._num_save_state_calls = variables_lib.Variable(0)
+ self._num_state_calls = variables_lib.VariableV1(0)
+ self._num_save_state_calls = variables_lib.VariableV1(0)
def state(self, name):
with ops_lib.control_dependencies(
diff --git a/tensorflow/contrib/session_bundle/exporter_test.py b/tensorflow/contrib/session_bundle/exporter_test.py
index 86df425..68419ff 100644
--- a/tensorflow/contrib/session_bundle/exporter_test.py
+++ b/tensorflow/contrib/session_bundle/exporter_test.py
@@ -64,10 +64,10 @@
# v2 is an unsaved variable derived from v0 and v1. It is used to
# exercise the ability to run an init op when restoring a graph.
with sess.graph.device("/cpu:0"):
- v0 = variables.Variable(10, name="v0")
+ v0 = variables.VariableV1(10, name="v0")
with sess.graph.device("/cpu:1"):
- v1 = variables.Variable(20, name="v1")
- v2 = variables.Variable(1, name="v2", trainable=False, collections=[])
+ v1 = variables.VariableV1(20, name="v1")
+ v2 = variables.VariableV1(1, name="v2", trainable=False, collections=[])
assign_v2 = state_ops.assign(v2, math_ops.add(v0, v1))
init_op = control_flow_ops.group(assign_v2, name="init_op")
diff --git a/tensorflow/contrib/stateless/BUILD b/tensorflow/contrib/stateless/BUILD
index dcbef28..a217397 100644
--- a/tensorflow/contrib/stateless/BUILD
+++ b/tensorflow/contrib/stateless/BUILD
@@ -9,19 +9,13 @@
load("//tensorflow:tensorflow.bzl", "cuda_py_test")
load("//tensorflow:tensorflow.bzl", "tf_gen_op_wrapper_py")
-tf_gen_op_wrapper_py(
- name = "stateless_random_ops",
- out = "gen_stateless_random_ops.py", # cmake chokes without this
- deps = ["//tensorflow/core:stateless_random_ops_op_lib"],
-)
-
py_library(
name = "stateless",
srcs = ["__init__.py"],
srcs_version = "PY2AND3",
deps = [
- ":stateless_random_ops",
"//tensorflow/python:framework_ops",
+ "//tensorflow/python:stateless_random_ops_gen",
"//tensorflow/python:util",
],
)
diff --git a/tensorflow/contrib/stateless/__init__.py b/tensorflow/contrib/stateless/__init__.py
index 0cca40f..fe23fe0 100644
--- a/tensorflow/contrib/stateless/__init__.py
+++ b/tensorflow/contrib/stateless/__init__.py
@@ -32,10 +32,11 @@
from __future__ import division
from __future__ import print_function
-# pylint: disable=wildcard-import
-from tensorflow.contrib.stateless.gen_stateless_random_ops import *
-
from tensorflow.python.framework import ops
+
+# pylint: disable=wildcard-import
+from tensorflow.python.ops.gen_stateless_random_ops import *
+
from tensorflow.python.util.all_util import remove_undocumented
ops.NotDifferentiable("StatelessMultinomial")
diff --git a/tensorflow/contrib/tensor_forest/client/random_forest.py b/tensorflow/contrib/tensor_forest/client/random_forest.py
index 0042d37..6e3bfbb 100644
--- a/tensorflow/contrib/tensor_forest/client/random_forest.py
+++ b/tensorflow/contrib/tensor_forest/client/random_forest.py
@@ -446,6 +446,10 @@
Returns:
A `TensorForestEstimator` instance.
"""
+ # Override default number of trainers if config is provided.
+ if num_trainers == 1 and config is not None:
+ num_trainers = max(1, config.num_worker_replicas)
+
super(TensorForestEstimator, self).__init__(
model_fn=get_model_fn(
params.fill(),
@@ -564,6 +568,10 @@
local_eval=False):
"""See TensorForestEstimator.__init__."""
model_fns = []
+ # Override default number of trainers if config is provided.
+ if num_trainers == 1 and config is not None:
+ num_trainers = max(1, config.num_worker_replicas)
+
for i in range(len(params_list)):
params = params_list[i].fill()
model_fns.append(
@@ -709,6 +717,11 @@
Returns:
A `TensorForestEstimator` instance.
"""
+ # Override default number of trainers if config is provided.
+ if num_trainers == 1 and config is not None:
+ num_trainers = max(1, config.num_worker_replicas)
+ if trainer_id == 0 and config is not None:
+ trainer_id = config.global_id_in_cluster
super(CoreTensorForestEstimator, self).__init__(
model_fn=get_model_fn(
diff --git a/tensorflow/contrib/tensor_forest/python/kernel_tests/scatter_add_ndim_op_test.py b/tensorflow/contrib/tensor_forest/python/kernel_tests/scatter_add_ndim_op_test.py
index 1c4e18d..0b02bdc 100644
--- a/tensorflow/contrib/tensor_forest/python/kernel_tests/scatter_add_ndim_op_test.py
+++ b/tensorflow/contrib/tensor_forest/python/kernel_tests/scatter_add_ndim_op_test.py
@@ -27,7 +27,7 @@
class ScatterAddNdimTest(test_util.TensorFlowTestCase):
def test1dim(self):
- input_data = variables.Variable(
+ input_data = variables.VariableV1(
[1., 2., 3., 4., 5., 6., 7., 8., 9., 10., 11., 12.])
indices = [[1], [10]]
updates = [100., 200.]
@@ -40,8 +40,8 @@
input_data.eval())
def test3dim(self):
- input_data = variables.Variable([[[1., 2., 3.], [4., 5., 6.]],
- [[7., 8., 9.], [10., 11., 12.]]])
+ input_data = variables.VariableV1([[[1., 2., 3.], [4., 5., 6.]],
+ [[7., 8., 9.], [10., 11., 12.]]])
indices = [[0, 0, 1], [1, 1, 2]]
updates = [100., 200.]
@@ -53,7 +53,7 @@
def testNoUpdates(self):
init_val = [[[1., 2., 3.], [4., 5., 6.]], [[7., 8., 9.], [10., 11., 12.]]]
- input_data = variables.Variable(init_val)
+ input_data = variables.VariableV1(init_val)
indices = []
updates = []
@@ -64,7 +64,7 @@
def testBadInput(self):
init_val = [[[1., 2., 3.], [4., 5., 6.]], [[7., 8., 9.], [10., 11., 12.]]]
- input_data = variables.Variable(init_val)
+ input_data = variables.VariableV1(init_val)
indices = [[0, 0, 1], [1, 1, 2]]
updates = [100.]
with self.cached_session():
@@ -75,8 +75,8 @@
self.assertAllEqual(init_val, input_data.eval())
def testIncompleteIndices(self):
- input_data = variables.Variable([[[1., 2., 3.], [4., 5., 6.]],
- [[7., 8., 9.], [10., 11., 12.]]])
+ input_data = variables.VariableV1([[[1., 2., 3.], [4., 5., 6.]],
+ [[7., 8., 9.], [10., 11., 12.]]])
indices = [[0, 0], [1, 1]]
updates = [[100., 200., 300.], [400., 500., 600.]]
diff --git a/tensorflow/contrib/tensorrt/python/trt_convert_test.py b/tensorflow/contrib/tensorrt/python/trt_convert_test.py
index f3a1ef0..52cb0bd 100644
--- a/tensorflow/contrib/tensorrt/python/trt_convert_test.py
+++ b/tensorflow/contrib/tensorrt/python/trt_convert_test.py
@@ -94,7 +94,7 @@
with g.device("/GPU:0"):
inp = array_ops.placeholder(
dtype=dtypes.float32, shape=[None, 1, 1], name="input")
- var = variables.Variable([[[1.0]]], dtype=dtypes.float32, name="v1")
+ var = variables.VariableV1([[[1.0]]], dtype=dtypes.float32, name="v1")
add = inp + var.value()
mul = inp * add
add = mul + add
diff --git a/tensorflow/contrib/timeseries/examples/BUILD b/tensorflow/contrib/timeseries/examples/BUILD
index 21c0c30..5779721 100644
--- a/tensorflow/contrib/timeseries/examples/BUILD
+++ b/tensorflow/contrib/timeseries/examples/BUILD
@@ -1,4 +1,5 @@
load("//tensorflow:tensorflow.bzl", "py_test")
+load("//tensorflow:tensorflow.bzl", "py_binary")
package(
default_visibility = ["//tensorflow:internal"],
diff --git a/tensorflow/contrib/timeseries/python/timeseries/head_test.py b/tensorflow/contrib/timeseries/python/timeseries/head_test.py
index 647455ae..04d17bc 100644
--- a/tensorflow/contrib/timeseries/python/timeseries/head_test.py
+++ b/tensorflow/contrib/timeseries/python/timeseries/head_test.py
@@ -104,7 +104,7 @@
"ticker":
array_ops.reshape(
math_ops.cast(
- variables.Variable(
+ variables.VariableV1(
name="ticker",
initial_value=0,
dtype=dtypes.int64,
diff --git a/tensorflow/contrib/tpu/BUILD b/tensorflow/contrib/tpu/BUILD
index 4e0b612..0c4bdab 100644
--- a/tensorflow/contrib/tpu/BUILD
+++ b/tensorflow/contrib/tpu/BUILD
@@ -36,6 +36,27 @@
)
py_library(
+ name = "async_checkpoint",
+ srcs = ["python/tpu/async_checkpoint.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ "//tensorflow/python:array_ops",
+ "//tensorflow/python:control_flow_ops",
+ "//tensorflow/python:framework_for_generated_wrappers",
+ "//tensorflow/python:init_ops",
+ "//tensorflow/python:math_ops",
+ "//tensorflow/python:platform",
+ "//tensorflow/python:state_ops",
+ "//tensorflow/python:summary",
+ "//tensorflow/python:summary_ops_v2",
+ "//tensorflow/python:training",
+ "//tensorflow/python:variable_scope",
+ "//tensorflow/python:variables",
+ "//tensorflow/python/estimator:estimator_py",
+ ],
+)
+
+py_library(
name = "tpu_estimator",
srcs = [
"python/tpu/error_handling.py",
@@ -46,6 +67,7 @@
],
srcs_version = "PY2AND3",
deps = [
+ ":async_checkpoint",
":tpu_lib",
"//tensorflow/compiler/xla/experimental/xla_sharding",
"//tensorflow/compiler/xla/python_api:xla_shape",
@@ -81,6 +103,9 @@
],
deps = [
"//tensorflow/contrib/tpu/proto:tpu_embedding_configuration_proto_cc",
+ "//tensorflow/contrib/tpu/utils:tpu_embedding_optimization_parameters_utils",
+ "//tensorflow/contrib/tpu/utils:tpu_embedding_output_layout_utils",
+ "//tensorflow/core:lib",
"//tensorflow/core:lib_proto_parsing",
"//tensorflow/core:protos_all_cc",
],
@@ -100,12 +125,17 @@
],
deps = [
"//tensorflow/contrib/tpu/proto:tpu_embedding_configuration_proto_cc",
+ "//tensorflow/contrib/tpu/utils:tpu_embedding_optimization_parameters_utils",
+ "//tensorflow/contrib/tpu/utils:tpu_embedding_output_layout_utils",
"//tensorflow/core:lib_proto_parsing",
],
)
tf_gen_op_wrapper_py(
name = "tpu_ops",
+ hidden = [
+ "SendTPUEmbeddingGradients",
+ ],
deps = [
":cross_replica_ops_op_lib",
":heartbeat_ops_op_lib",
@@ -225,7 +255,10 @@
":tpu_py",
"//tensorflow/contrib/cluster_resolver:tpu_cluster_resolver_py",
"//tensorflow/contrib/tpu/proto:compilation_result_proto_py",
+ "//tensorflow/contrib/tpu/proto:optimization_parameters_proto_py",
"//tensorflow/contrib/tpu/proto:topology_proto_py",
+ "//tensorflow/contrib/tpu/proto:tpu_embedding_configuration_proto_py",
+ "//tensorflow/contrib/tpu/proto:tpu_embedding_output_layout_proto_py",
"//tensorflow/core:protos_all_py",
"//tensorflow/python:array_ops",
"//tensorflow/python:control_flow_ops",
diff --git a/tensorflow/contrib/tpu/__init__.py b/tensorflow/contrib/tpu/__init__.py
index 3c0456d..6ce6b77 100644
--- a/tensorflow/contrib/tpu/__init__.py
+++ b/tensorflow/contrib/tpu/__init__.py
@@ -55,6 +55,9 @@
@@TPUDistributionStrategy
@@keras_to_tpu_model
+
+@@AsyncCheckpointSaverHook
+@@TPUInMemoryEvalHook
"""
from __future__ import absolute_import
@@ -64,6 +67,7 @@
# pylint: disable=wildcard-import,unused-import
from tensorflow.contrib.tpu.python import profiler
from tensorflow.contrib.tpu.python.ops.tpu_ops import *
+from tensorflow.contrib.tpu.python.tpu.async_checkpoint import *
from tensorflow.contrib.tpu.python.tpu.bfloat16 import *
from tensorflow.contrib.tpu.python.tpu.device_assignment import *
from tensorflow.contrib.tpu.python.tpu.keras_support import tpu_model as keras_to_tpu_model
diff --git a/tensorflow/contrib/tpu/ops/tpu_embedding_ops.cc b/tensorflow/contrib/tpu/ops/tpu_embedding_ops.cc
index 18b9893..ef2f8dd 100644
--- a/tensorflow/contrib/tpu/ops/tpu_embedding_ops.cc
+++ b/tensorflow/contrib/tpu/ops/tpu_embedding_ops.cc
@@ -14,10 +14,15 @@
==============================================================================*/
#include "tensorflow/contrib/tpu/proto/tpu_embedding_configuration.pb.h"
+#include "tensorflow/contrib/tpu/utils/tpu_embedding_optimization_parameters_utils.h"
+#include "tensorflow/contrib/tpu/utils/tpu_embedding_output_layout_utils.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/common_shape_fns.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"
#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/lib/strings/stringprintf.h"
namespace tensorflow {
@@ -41,7 +46,7 @@
// 5. TPUEmbeddingActivations, when used with appropriate Python libraries,
// enables the automatic differentiation of models that use embeddings.
// 6. TPUEmbeddingSendGradients takes a list of Tensors (of the same shapes
-// as those returned by TPUEmbeddingReceivActivations) containing gradients
+// as those returned by TPUEmbeddingReceiveActivations) containing gradients
// to use in updating the embedding tables.
// 7. Before saving a checkpoint, use the TPUEmbeddingRetrieve Op to update
// the Graph's embedding table Variables from the updated tables in the
@@ -53,230 +58,354 @@
// saving a checkpoint, the model must Retrieve the parameters back into the
// host CPU memory.
-REGISTER_OP("TPUEmbeddingLoadGradientDescentParameters")
- .Input("parameters: float32")
- .Attr("tpu_embedding_config: string")
- .Attr("table_id: int >= 0")
- .Attr("num_hosts: int >= 1")
- .Attr("host_id: int >= 0")
- .SetIsStateful()
- .SetShapeFn(shape_inference::UnknownShape)
- .Doc(R"doc(
-Load an embedding table shard into TPU memory for use with GradientDescent.
+namespace {
-TPU embeddings use dedicated per-optimizer Ops for loading and retrieving
-trainable variables and optimizer state from TPU memory. This op enables
-functionality equivalent to GradientDescentOptimizer.
+void RegisterPerTableLoadAndRetrieveOps();
-parameters: The shard of the embedding table resident on the host executing this
- op. For single-TPU models, this is the entire embedding table.
-tpu_embedding_config: Serialized TPUEmbeddingConfiguration proto.
-table_id: The id of the table specified in the tpu_embedding_config.
-num_hosts: The number of CPU hosts in the distributed training job.
-host_id: Which CPU host in the distributed training job will execute this op.
-)doc");
-
-namespace tpu_embedding_config_util {
-
-Status GradientDescentShapes(shape_inference::InferenceContext *c) {
- string config_string;
- TF_RETURN_IF_ERROR(c->GetAttr("tpu_embedding_config", &config_string));
- tpu::TPUEmbeddingConfiguration config;
- if (!config.ParseFromString(config_string)) {
- return errors::InvalidArgument("Malformed tpu_embedding_config.");
+class RegisterPerTableLoadAndRetrieveOpsOnConstruction {
+ public:
+ RegisterPerTableLoadAndRetrieveOpsOnConstruction() {
+ RegisterPerTableLoadAndRetrieveOps();
}
+};
- int table_id;
- TF_RETURN_IF_ERROR(c->GetAttr("table_id", &table_id));
- int64 num_tables = config.table_descriptor_size();
- if (table_id >= num_tables) {
- return errors::InvalidArgument("Table id >= num_tables");
+// Object whose constructor does registrations.
+RegisterPerTableLoadAndRetrieveOpsOnConstruction
+ register_per_table_load_and_retrieve_ops_var;
+
+Status RegisterPerTableLoadOpsForAlgorithmBody(
+ tpu::OptimizationAlgorithm alg, bool is_debug_op,
+ OpRegistrationData* op_reg_data) {
+ tpu::GradientAccumulationSupport grad_accum_support;
+ TF_CHECK_OK(GetGradientAccumulationSupport(alg, &grad_accum_support));
+
+ std::vector<tpu::StateVariableSpecification> state_variable_specs;
+ TF_CHECK_OK(GetOptimizationAlgorithmStateVariables(
+ alg,
+ grad_accum_support == tpu::GradientAccumulationSupport::kSupported &&
+ is_debug_op,
+ &state_variable_specs));
+ auto* op_def = &op_reg_data->op_def;
+ op_def->set_name(
+ strings::StrCat("LoadTPUEmbedding", GetOptimizationAlgorithmName(alg),
+ "Parameters", (is_debug_op ? "GradAccumDebug" : "")));
+ // It is important for the order of the inputs to the op defined here
+ // to match the order in input_names because the indexes are used in
+ // the combining transformation.
+ for (const auto& parameter : state_variable_specs) {
+ if (parameter.has_user_defined() || is_debug_op) {
+ auto* arg = op_def->add_input_arg();
+ arg->set_name(parameter.name());
+ arg->set_description(
+ strings::StrCat("Value of ", parameter.name(), " used in the ",
+ GetOptimizationAlgorithmFriendlyName(alg),
+ " optimization algorithm."));
+ arg->set_type(DT_FLOAT);
+ }
}
- int64 width = config.table_descriptor(table_id).dimension();
- int64 num_rows = config.table_descriptor(table_id).vocabulary_size();
-
- TF_RETURN_IF_ERROR(c->set_output("parameters", {c->Matrix(num_rows, width)}));
+ {
+ auto* table_id_attr = op_def->add_attr();
+ table_id_attr->set_name("table_id");
+ table_id_attr->set_type("int");
+ table_id_attr->set_has_minimum(true);
+ table_id_attr->set_minimum(-1);
+ table_id_attr->mutable_default_value()->set_i(-1);
+ }
+ {
+ auto* table_name_attr = op_def->add_attr();
+ table_name_attr->set_name("table_name");
+ table_name_attr->set_type("string");
+ table_name_attr->mutable_default_value()->set_s("");
+ }
+ {
+ auto* num_shards_attr = op_def->add_attr();
+ num_shards_attr->set_name("num_shards");
+ num_shards_attr->set_type("int");
+ }
+ {
+ auto* shard_id_attr = op_def->add_attr();
+ shard_id_attr->set_name("shard_id");
+ shard_id_attr->set_type("int");
+ }
+ op_def->set_summary("Load embedding parameters for a single table.");
+ string parameter_descriptions;
+ for (const auto& parameter : state_variable_specs) {
+ if (parameter.has_user_defined() || is_debug_op) {
+ strings::Appendf(¶meter_descriptions,
+ R"(
+%s: A tensor containing the initial embedding table %s to use in embedding
+lookups using the %s optimization algorithm.)",
+ parameter.name().c_str(), parameter.name().c_str(),
+ GetOptimizationAlgorithmFriendlyName(alg).c_str());
+ }
+ }
+ op_def->set_description(strings::Printf(R"doc(
+An op that loads optimization parameters into HBM for embedding. Must be
+preceded by a ConfigureTPUEmbeddingHost op that sets up the correct
+embedding table configuration. For example, this op is used to install
+parameters that are loaded from a checkpoint before a training loop is
+executed.
+%s
+table_name: Name of this table; must match a name in the
+ TPUEmbeddingConfiguration proto (overrides table_id).
+num_shards: Number of shards into which the embedding tables are divided.
+shard_id: Identifier of shard for this operation.
+table_id: Index of this table in the EmbeddingLayerConfiguration proto
+ (deprecated).
+)doc",
+ parameter_descriptions.c_str()));
+ op_def->set_is_commutative(false);
+ op_def->set_is_aggregate(false);
+ op_def->set_is_stateful(true);
+ auto shape_inference_function =
+ [state_variable_specs,
+ is_debug_op](shape_inference::InferenceContext* c) -> Status {
+ int table_id;
+ TF_RETURN_IF_ERROR(c->GetAttr("table_id", &table_id));
+ string table_name;
+ TF_RETURN_IF_ERROR(c->GetAttr("table_name", &table_name));
+ // Exactly one must be non-default.
+ if ((table_id >= 0) == (!table_name.empty())) {
+ return errors::InvalidArgument(
+ "exactly one of table_id or table_name must be non-default");
+ }
+ int num_shards;
+ TF_RETURN_IF_ERROR(c->GetAttr("num_shards", &num_shards));
+ int shard_id;
+ TF_RETURN_IF_ERROR(c->GetAttr("shard_id", &shard_id));
+ const int user_param_count =
+ std::count_if(state_variable_specs.begin(), state_variable_specs.end(),
+ [&](const tpu::StateVariableSpecification& sv) {
+ return sv.has_user_defined() || is_debug_op;
+ });
+ std::vector<shape_inference::ShapeHandle> inputs(user_param_count);
+ int input_index = 0;
+ for (int i = 0; i < state_variable_specs.size(); ++i) {
+ if (state_variable_specs[i].has_user_defined() || is_debug_op) {
+ std::vector<shape_inference::ShapeHandle> input_temp;
+ TF_RETURN_IF_ERROR(
+ c->input(state_variable_specs[i].name(), &input_temp));
+ if (input_temp.size() != 1) {
+ return errors::InvalidArgument("each input to be rank 1");
+ }
+ inputs[input_index] = input_temp[0];
+ ++input_index;
+ }
+ }
+ // Verify shapes have rank 2 and are compatible when they are
+ // required to be valid.
+ shape_inference::ShapeHandle parameter_shape;
+ TF_RETURN_IF_ERROR(c->WithRank(inputs[0], 2, ¶meter_shape));
+ for (int j = 1; j < user_param_count; ++j) {
+ shape_inference::ShapeHandle accumulator_j_shape;
+ TF_RETURN_IF_ERROR(c->WithRank(inputs[j], 2, &accumulator_j_shape));
+ shape_inference::ShapeHandle merged;
+ TF_RETURN_IF_ERROR(
+ c->Merge(parameter_shape, accumulator_j_shape, &merged));
+ }
+ return Status::OK();
+ };
+ op_reg_data->shape_inference_fn = shape_inference_function;
return Status::OK();
}
-} // namespace tpu_embedding_config_util
+Status RegisterPerTableRetrieveOpsForAlgorithmBody(
+ tpu::OptimizationAlgorithm alg, bool is_debug_op,
+ OpRegistrationData* op_reg_data) {
+ tpu::GradientAccumulationSupport grad_accum_support;
+ TF_CHECK_OK(GetGradientAccumulationSupport(alg, &grad_accum_support));
-REGISTER_OP("TPUEmbeddingRetrieveGradientDescentParameters")
- .Output("parameters: float32")
- .Attr("tpu_embedding_config: string")
- .Attr("table_id: int")
- .Attr("num_hosts: int")
- .Attr("host_id: int")
- .SetIsStateful()
- .SetShapeFn(tpu_embedding_config_util::GradientDescentShapes)
- .Doc(R"doc(
-Retrieve an embedding table shard from TPU memory.
+ std::vector<tpu::StateVariableSpecification> state_variable_specs;
+ TF_CHECK_OK(GetOptimizationAlgorithmStateVariables(
+ alg,
+ grad_accum_support == tpu::GradientAccumulationSupport::kSupported &&
+ is_debug_op,
+ &state_variable_specs));
-TPU embeddings use dedicated per-optimizer Ops for loading and retrieving
-trainable variables and optimizer state from TPU memory. This op enables
-functionality equivalent to GradientDescentOptimizer.
-
-tpu_embedding_config: Serialized TPUEmbeddingConfiguration proto.
-table_id: The id of the table specified in tpu_embedding_config.
-num_hosts: The number of CPU hosts in the distributed training job.
-host_id: Which CPU host in the distributed training job will execute this op.
-)doc");
-
-REGISTER_OP("TPUEmbeddingLoadAdagradParameters")
- .Input("parameters: float32")
- .Input("accumulators: float32")
- .Attr("tpu_embedding_config: string")
- .Attr("table_id: int >= 0")
- .Attr("num_hosts: int >= 1")
- .Attr("host_id: int >= 0")
- .SetIsStateful()
- .SetShapeFn(shape_inference::UnknownShape)
- .Doc(R"doc(
-Load an embedding table shard into TensorNode memories for use with Adagrad.
-
-TPU embeddings use dedicated per-optimizer Ops for loading and retrieving
-trainable variables and optimizer state from TPU memory. This op enables
-functionality equivalent to AdagradOptimizer.
-
-parameters: The shard of the embedding table resident on the host executing this
- op. For single-TPU models, this is the entire embedding table.
-accumulators: Shard of the Adagrad accumulators resident on the host executing
- this op.
-tpu_embedding_config: Serialized TPUEmbeddingConfiguration proto.
-table_id: The id of the table specified in the embedding_config.
-num_hosts: The number of CPU hosts in the distributed training job.
-host_id: Which CPU host in the distributed training job will execute this op.
-)doc");
-
-namespace tpu_embedding_config_util {
-
-Status AdagradShapes(shape_inference::InferenceContext *c) {
- string config_string;
- TF_RETURN_IF_ERROR(c->GetAttr("tpu_embedding_config", &config_string));
- tpu::TPUEmbeddingConfiguration config;
- if (!config.ParseFromString(config_string)) {
- return errors::InvalidArgument("Malformed tpu_embedding_config.");
+ auto* op_def = &op_reg_data->op_def;
+ op_def->set_name(strings::StrCat(
+ "RetrieveTPUEmbedding", tpu::GetOptimizationAlgorithmName(alg),
+ "Parameters", (is_debug_op ? "GradAccumDebug" : "")));
+ // It is important for the order of the outputs of the op defined here
+ // to match the order in output_names because the indexes are used in
+ // the combining transformation.
+ for (const auto& parameter : state_variable_specs) {
+ if (parameter.has_user_defined() || is_debug_op) {
+ auto* arg = op_def->add_output_arg();
+ arg->set_name(parameter.name());
+ arg->set_description(
+ strings::StrCat("Parameter ", parameter.name(), " updated by the ",
+ tpu::GetOptimizationAlgorithmFriendlyName(alg),
+ " optimization algorithm."));
+ arg->set_type(DT_FLOAT);
+ }
}
-
- int table_id;
- TF_RETURN_IF_ERROR(c->GetAttr("table_id", &table_id));
- int64 num_tables = config.table_descriptor_size();
- if (table_id >= num_tables) {
- return errors::InvalidArgument("Table id >= num_tables");
+ {
+ auto* table_id_attr = op_def->add_attr();
+ table_id_attr->set_name("table_id");
+ table_id_attr->set_type("int");
+ table_id_attr->set_has_minimum(true);
+ table_id_attr->set_minimum(-1);
+ table_id_attr->mutable_default_value()->set_i(-1);
}
- int64 width = config.table_descriptor(table_id).dimension();
- int64 num_rows = config.table_descriptor(table_id).vocabulary_size();
-
- TF_RETURN_IF_ERROR(c->set_output("parameters", {c->Matrix(num_rows, width)}));
- TF_RETURN_IF_ERROR(
- c->set_output("accumulators", {c->Matrix(num_rows, width)}));
+ {
+ auto* table_name_attr = op_def->add_attr();
+ table_name_attr->set_name("table_name");
+ table_name_attr->set_type("string");
+ table_name_attr->mutable_default_value()->set_s("");
+ }
+ {
+ auto* num_shards_attr = op_def->add_attr();
+ num_shards_attr->set_name("num_shards");
+ num_shards_attr->set_type("int");
+ }
+ {
+ auto* shard_id_attr = op_def->add_attr();
+ shard_id_attr->set_name("shard_id");
+ shard_id_attr->set_type("int");
+ }
+ op_def->set_summary("Retrieve embedding parameters for a single table.");
+ string parameter_descriptions;
+ for (const auto& param : state_variable_specs) {
+ if (param.has_user_defined() || is_debug_op) {
+ strings::Appendf(¶meter_descriptions,
+ R"(
+%s: A tensor containing the embedding table %s to store with the
+parameters from embedding updates using the %s optimization algorithm.)",
+ param.name().c_str(), param.name().c_str(),
+ tpu::GetOptimizationAlgorithmFriendlyName(alg).c_str());
+ }
+ }
+ op_def->set_description(strings::Printf(R"doc(
+An op that retrieves optimization parameters from embedding to host
+memory. Must be preceded by a ConfigureTPUEmbeddingHost op that sets up
+the correct embedding table configuration. For example, this op is
+used to retrieve updated parameters before saving a checkpoint.
+%s
+table_name: Name of this table; must match a name in the
+ TPUEmbeddingConfiguration proto (overrides table_id).
+num_shards: Number of shards into which the embedding tables are divided.
+shard_id: Identifier of shard for this operation.
+table_id: Index of this table in the EmbeddingLayerConfiguration proto
+ (deprecated).
+)doc",
+ parameter_descriptions.c_str()));
+ op_def->set_is_commutative(false);
+ op_def->set_is_aggregate(false);
+ op_def->set_is_stateful(true);
+ auto shape_inference_function =
+ [state_variable_specs,
+ is_debug_op](shape_inference::InferenceContext* c) -> Status {
+ int table_id;
+ TF_RETURN_IF_ERROR(c->GetAttr("table_id", &table_id));
+ string table_name;
+ TF_RETURN_IF_ERROR(c->GetAttr("table_name", &table_name));
+ // Exactly one must be non-default.
+ if ((table_id >= 0) == (!table_name.empty())) {
+ return errors::InvalidArgument(
+ "exactly one of table_id or table_name must be non-default");
+ }
+ int num_shards;
+ TF_RETURN_IF_ERROR(c->GetAttr("num_shards", &num_shards));
+ int shard_id;
+ TF_RETURN_IF_ERROR(c->GetAttr("shard_id", &shard_id));
+ for (int j = 0; j < state_variable_specs.size(); ++j) {
+ if (state_variable_specs[j].has_user_defined() || is_debug_op) {
+ auto shape = c->MakeShape(
+ std::vector<shape_inference::DimensionHandle>(2, c->UnknownDim()));
+ TF_RETURN_IF_ERROR(
+ c->set_output(state_variable_specs[j].name(),
+ std::vector<shape_inference::ShapeHandle>(1, shape)));
+ }
+ }
+ return Status::OK();
+ };
+ op_reg_data->shape_inference_fn = shape_inference_function;
return Status::OK();
}
-} // namespace tpu_embedding_config_util
-
-REGISTER_OP("TPUEmbeddingRetrieveAdagradParameters")
- .Output("parameters: float32")
- .Output("accumulators: float32")
- .Attr("tpu_embedding_config: string")
- .Attr("table_id: int >= 0")
- .Attr("num_hosts: int >= 1")
- .Attr("host_id: int >= 0")
- .SetIsStateful()
- .SetShapeFn(tpu_embedding_config_util::AdagradShapes)
- .Doc(R"doc(
-Retrieve an embedding table shard from TPU memory.
-
-TPU embeddings use dedicated per-optimizer Ops for loading and retrieving
-trainable variables and optimizer state from TPU memory. This op enables
-functionality equivalent to AdagradOptimizer.
-
-tpu_embedding_config: Serialized TPUEmbeddingConfiguration proto.
-table_id: The id of the table specified in the embedding_config_json.
-num_hosts: The number of CPU hosts in the distributed training job.
-host_id: Which CPU host in the distributed training job will execute this op.
-)doc");
-
-REGISTER_OP("TPUEmbeddingEnqueueSparseBatch")
- .Input("sample_indices: num_tables * int32")
- .Input("embedding_indices: num_tables * int32")
- .Input("aggregation_weights: num_tables * float32")
- .Attr("num_tables: int")
- .Attr("device_ordinal: int = -1")
- .SetIsStateful()
- .SetShapeFn(shape_inference::UnknownShape)
- .Doc(R"doc(
-An op that feeds a batch of embedding indices and weights to the TPU.
-
-Embedding lookups are equivalent to sparse-dense matrix multiplications: the
-sparse matrix contains nonzeros in column j in order to retrieve row j from the
-embedding table.
-
-The three Tensor list arguments (sample_indices, embedding_indices, and
-aggregation_weights) represent these sparse matrices in COO format. The Tensor
-lists each have one entry for each embedding table specified in the model.
-For the kth embedding table, the three Tensors at position k in the list
-specify a COO-format sparse matrix. For the kth table, the row indices,
-column indices, and nonzero values of the COO sparse matrix are specified by
-sample_indices[k], embedding_indices[k], and aggregation_weights[k],
-respectively. Entries must be sorted by row index, then by column index.
-
-There should be at most one TPUEmbeddingEnqueueSparseBatch op in a signle
-training step per TPU shard.
-
-sample_indices: A list of rank 1 Tensors specifying row indices of the COO
- sparse matrix representing the embedding lookups for each table.
-embedding_indices: A list of rank 1 Tensors specifying column indices of the
- COO sparse matrix representing the embedding lookups for each table.
-aggregation_weights: A list of rank 1 Tensors specifying the nonzero values
- of the COO sparse matrix representing the embedding lookups for each table.
-device_ordinal: The TPU device to use. This should be -1 when the Op
- is running on a TPU device, and >= 0 when the Op is running on the CPU
- device.
-)doc");
-
-namespace tpu_embedding_config_util {
-
-Status ActivationShapes(shape_inference::InferenceContext *c) {
- string config_string;
- TF_RETURN_IF_ERROR(c->GetAttr("tpu_embedding_config", &config_string));
- tpu::TPUEmbeddingConfiguration config;
- if (!config.ParseFromString(config_string)) {
- return errors::InvalidArgument("Malformed tpu_embedding_config.");
+void RegisterPerTableLoadAndRetrieveOps() {
+ // Load ops
+ for (tpu::OptimizationAlgorithm alg : tpu::GetOptimizationAlgorithms()) {
+ OpRegistry::Global()->Register(
+ [alg](OpRegistrationData* op_reg_data) -> Status {
+ return RegisterPerTableLoadOpsForAlgorithmBody(alg, false,
+ op_reg_data);
+ });
+ tpu::GradientAccumulationSupport grad_accum_support;
+ TF_CHECK_OK(GetGradientAccumulationSupport(alg, &grad_accum_support));
+ if (grad_accum_support == tpu::GradientAccumulationSupport::kSupported) {
+ // TODO(gkurian): Condition this on being used internally within Google.
+ OpRegistry::Global()->Register(
+ [alg](OpRegistrationData* op_reg_data) -> Status {
+ return RegisterPerTableLoadOpsForAlgorithmBody(alg, true,
+ op_reg_data);
+ });
+ }
}
- int64 batch_size = config.batch_size_per_tensor_core();
- int64 num_tables = config.table_descriptor_size();
- for (int table_id = 0; table_id < num_tables; ++table_id) {
- int64 width = config.table_descriptor(table_id).dimension();
- int64 num_features = config.table_descriptor(table_id).vocabulary_size();
- c->set_output(table_id, c->Matrix(batch_size * num_features, width));
+ // Retrieve ops
+ for (tpu::OptimizationAlgorithm alg : tpu::GetOptimizationAlgorithms()) {
+ OpRegistry::Global()->Register(
+ [alg](OpRegistrationData* op_reg_data) -> Status {
+ return RegisterPerTableRetrieveOpsForAlgorithmBody(alg, false,
+ op_reg_data);
+ });
+ tpu::GradientAccumulationSupport grad_accum_support;
+ TF_CHECK_OK(GetGradientAccumulationSupport(alg, &grad_accum_support));
+ if (grad_accum_support == tpu::GradientAccumulationSupport::kSupported) {
+ // TODO(gkurian): Condition this on being used internally within Google.
+ OpRegistry::Global()->Register(
+ [alg](OpRegistrationData* op_reg_data) -> Status {
+ return RegisterPerTableRetrieveOpsForAlgorithmBody(alg, true,
+ op_reg_data);
+ });
+ }
}
- return Status::OK();
}
-} // namespace tpu_embedding_config_util
+} // namespace
-REGISTER_OP("TPUEmbeddingReceiveActivations")
- .Output("outputs: num_tables * float")
- .Attr("num_tables: int >= 1")
- .Attr("tpu_embedding_config: string")
+REGISTER_OP("RecvTPUEmbeddingActivations")
+ .Output("outputs: num_outputs * float")
+ .Attr("num_outputs: int >= 1")
+ .Attr("config: string")
.SetIsStateful()
- .SetShapeFn(tpu_embedding_config_util::ActivationShapes)
+ .SetShapeFn([](shape_inference::InferenceContext* c) -> Status {
+ string config_string;
+ TF_RETURN_IF_ERROR(c->GetAttr("config", &config_string));
+ tpu::TPUEmbeddingConfiguration config;
+ if (!config.ParseFromString(config_string)) {
+ return errors::InvalidArgument("Malformed tpu_embedding_config.");
+ }
+ tpu::AddDefaultEmbeddingOutputLayoutIfNeeded(&config);
+ std::vector<TensorShapeProto> output_shapes;
+ TF_RETURN_IF_ERROR(ComputeOutputTensorShapes(config, &output_shapes));
+ if (c->num_outputs() != output_shapes.size()) {
+ return errors::InvalidArgument("num outputs != size of output shapes");
+ }
+ for (int i = 0; i < c->num_outputs(); ++i) {
+ shape_inference::ShapeHandle output_shape;
+ TF_RETURN_IF_ERROR(
+ c->MakeShapeFromShapeProto(output_shapes[i], &output_shape));
+ c->set_output(i, output_shape);
+ }
+ return Status::OK();
+ })
.Doc(R"doc(
An op that receives embedding activations on the TPU.
The TPU system performs the embedding lookups and aggregations specified by
-the arguments to TPUEmbeddingEnqueueSparseBatch. The results of these
-aggregations are visible to the Tensorflow Graph as the outputs of a
-TPUEmbeddingDequeueActivations Op. This op returns a list containing one
-Tensor of activations per table specified in the model. There can be at most
-one ReceieveActivations op in the TPU graph.
+the arguments to TPUEmbeddingEnqueue(Integer/Sparse/SparseTensor)Batch. The
+results of these aggregations are visible to the Tensorflow Graph as the
+outputs of a RecvTPUEmbeddingActivations op. This op returns a list containing
+one Tensor of activations per table specified in the model. There can be at
+most one RecvTPUEmbeddingActivations op in the TPU graph.
outputs: A TensorList of embedding activations containing one Tensor per
embedding table in the model.
-num_tables: The number of output activation tensors, equal to the number of
+num_outputs: The number of output activation tensors, equal to the number of
embedding tables in the model.
-tpu_embedding_config: Serialized TPUEmbeddingConfiguration proto.
+config: Serialized TPUEmbeddingConfiguration proto.
)doc");
REGISTER_OP("TPUEmbeddingActivations")
@@ -306,12 +435,27 @@
activations.
)doc");
-REGISTER_OP("TPUEmbeddingSendGradients")
- .Input("gradients: num_tables * float32")
- .Attr("num_tables: int >= 1")
- .Attr("tpu_embedding_config: string")
+REGISTER_OP("SendTPUEmbeddingGradients")
+ .Input("inputs: N * float32")
+ .Input("learning_rates: NN * float32")
+ .Attr("N: int >= 1")
+ .Attr("NN: int >= 0 = 0")
+ .Attr("config: string")
.SetIsStateful()
- .SetShapeFn(shape_inference::UnknownShape)
+ .SetShapeFn([](shape_inference::InferenceContext* c) -> Status {
+ int nn;
+ TF_RETURN_IF_ERROR(c->GetAttr("NN", &nn));
+ std::vector<shape_inference::ShapeHandle> learning_rates;
+ TF_RETURN_IF_ERROR(c->input("learning_rates", &learning_rates));
+ for (int i = 0; i < nn; ++i) {
+ // Verify that each learning_rates element is scalar
+ shape_inference::ShapeHandle learning_rates_shape;
+ TF_RETURN_IF_ERROR(
+ c->WithRank(learning_rates[i], 0, &learning_rates_shape));
+ }
+
+ return Status::OK();
+ })
.Doc(R"doc(
An op that performs gradient updates of embedding tables.
@@ -321,8 +465,120 @@
from these gradients via the optimizer specified in the configuration given
to tpu.initialize_system.
-gradients: A TensorList of gradients with which to update embedding tables.
-tpu_embedding_config: Serialized TPUEmbeddingConfiguration proto.
+inputs: A TensorList of gradients with which to update embedding tables.
+ It contains one tensor per embedding table in the model.
+learning_rates: A list of float32 scalars, one for each embedding table,
+ containing the learning rates for each table when dynamic learning rate is
+ enabled through the OptimizationParameters in TPUEmbeddingConfiguration.
+ When the learning rate is constant, the list should be empty.
+config: Serialized TPUEmbeddingConfiguration proto.
+)doc");
+
+REGISTER_OP("EnqueueTPUEmbeddingIntegerBatch")
+ .Input("batch: N * int32")
+ .Attr("N: int")
+ .Attr("device_ordinal: int = -1")
+ .SetIsStateful()
+ .SetShapeFn(shape_inference::UnknownShape)
+ .Doc(R"doc(
+An op that enqueues a list of input batch tensors to TPUEmbedding.
+
+batch: A list of 1D tensors, one for each embedding table, containing the
+ indices into the tables.
+device_ordinal: The TPU device to use. Should be >= 0 and less than the number
+ of TPU cores in the task on which the node is placed.
+)doc");
+
+REGISTER_OP("EnqueueTPUEmbeddingSparseBatch")
+ .Input("sample_indices: N * int32")
+ .Input("embedding_indices: N * int32")
+ .Input("aggregation_weights: N * float32")
+ .Attr("N: int")
+ .Attr("device_ordinal: int = -1")
+ .Attr("combiners: list(string) = []")
+ .SetIsStateful()
+ .SetShapeFn([](shape_inference::InferenceContext* c) -> Status {
+ std::vector<string> combiners;
+ TF_RETURN_IF_ERROR(c->GetAttr("combiners", &combiners));
+ int n;
+ TF_RETURN_IF_ERROR(c->GetAttr("N", &n));
+ if (!combiners.empty() && combiners.size() != n) {
+ return errors::InvalidArgument("Invalid length of combiners. Have ",
+ combiners.size(), " but expected 0 or ",
+ n);
+ }
+
+ return Status::OK();
+ })
+ .Doc(R"doc(
+An op that enqueues TPUEmbedding input indices from a SparseTensor.
+
+This Op eases the porting of code that uses embedding_lookup_sparse(),
+although some Python preprocessing of the SparseTensor arguments to
+embedding_lookup_sparse() is required to produce the arguments to this Op,
+since only a single EnqueueTPUEmbeddingSparseBatch Op is allowed per training
+step.
+
+The tensors at corresponding positions in the three input lists
+must have the same shape, i.e. rank 1 with dim_size() equal to the total
+number of lookups into the table described by the corresponding table_id.
+
+sample_indices: A list of Rank 1 Tensors specifying the training example and
+ feature to which the corresponding embedding_indices and aggregation_weights
+ values belong. sample_indices[i] must equal b * nf + f, where nf is the
+ number of features from the corresponding table, f is in [0, nf), and
+ b is in [0, batch size).
+embedding_indices: A list of Rank 1 Tensors, indices into the embedding tables.
+aggregation_weights: A list of Rank 1 Tensors containing per sample -- i.e. per
+ (training example, feature) -- aggregation weights.
+device_ordinal: The TPU device to use. Should be >= 0 and less than the number
+ of TPU cores in the task on which the node is placed.
+combiners: A list of string scalars, one for each embedding table that specify
+ how to normalize the embedding activations after weighted summation.
+ Supported combiners are 'mean', 'sum', or 'sqrtn'. It is invalid to have
+ the sum of the weights be 0 for 'mean' or the sum of the squared weights be
+ 0 for 'sqrtn'. If combiners isn't passed, the default is to use 'sum' for
+ all tables.
+)doc");
+
+REGISTER_OP("EnqueueTPUEmbeddingSparseTensorBatch")
+ .Input("sample_indices: N * int32")
+ .Input("embedding_indices: N * int32")
+ .Input("aggregation_weights: N * float32")
+ .Attr("N: int")
+ .Attr("device_ordinal: int = -1")
+ .Attr("combiners: list(string) = []")
+ .Attr("table_ids: list(int)")
+ .SetIsStateful()
+ .SetShapeFn(shape_inference::UnknownShape)
+ .Doc(R"doc(
+This Op eases the porting of code that uses tf.nn.embedding_lookup_sparse().
+
+sample_indices[i], embedding_indices[i] and aggregation_weights[i] correspond
+to ith feature. table_ids[i] indicates which embedding table to look up ith
+feature.
+
+The tensors at corresponding positions in the three input lists (sample_indices,
+embedding_indices and aggregation_weights) must have the same shape, i.e. rank 1
+with dim_size() equal to the total number of lookups into the table described by
+the corresponding feature.
+
+sample_indices: A list of Rank 1 Tensors, corresponds to sp_ids.indices[:,0] in
+ embedding_lookup_sparse().
+embedding_indices: A list of Rank 1 Tensors, corresponds to sp_ids.values
+ in embedding_lookup_sparse().
+aggregation_weights: A list of Rank 1 Tensors, corresponds to sp_weights.values
+ in embedding_lookup_sparse().
+device_ordinal: The TPU device to use. Should be >= 0 and less than the number
+ of TPU cores in the task on which the node is placed.
+combiners: A list of string scalars, one for each embedding table that specify
+ how to normalize the embedding activations after weighted summation.
+ Supported combiners are 'mean', 'sum', or 'sqrtn'. It is invalid to have
+ the sum of the weights be 0 for 'mean' or the sum of the squared weights be
+ 0 for 'sqrtn'. If combiners isn't passed, the default is to use 'sum' for
+ all tables.
+table_ids: A list of int. table_ids[i] indicates which embedding table to look
+ up ith feature in the list.
)doc");
} // namespace tensorflow
diff --git a/tensorflow/contrib/tpu/profiler/capture_tpu_profile.cc b/tensorflow/contrib/tpu/profiler/capture_tpu_profile.cc
index b498599..1c5ea2d 100644
--- a/tensorflow/contrib/tpu/profiler/capture_tpu_profile.cc
+++ b/tensorflow/contrib/tpu/profiler/capture_tpu_profile.cc
@@ -156,8 +156,7 @@
channel_args));
NewProfileSessionResponse new_session_response;
TF_QCHECK_OK(FromGrpcStatus(
- stub->NewSession(&context, new_session_request, &new_session_response)))
- << new_session_response.error_message();
+ stub->NewSession(&context, new_session_request, &new_session_response)));
std::cout << "Profile session succeed for host(s):"
<< str_util::Join(hostnames, ",") << std::endl;
@@ -238,7 +237,8 @@
MonitorResponse response;
TF_QCHECK_OK(FromGrpcStatus(stub->Monitor(&context, request, &response)));
- std::cout << "Xprof Monitoring Results (Sample " << query + 1 << "):\n\n"
+ std::cout << "Cloud TPU Monitoring Results (Sample " << query + 1
+ << "):\n\n"
<< response.data() << std::flush;
}
}
diff --git a/tensorflow/contrib/tpu/profiler/op_profile.proto b/tensorflow/contrib/tpu/profiler/op_profile.proto
index 68cf510..292108f 100644
--- a/tensorflow/contrib/tpu/profiler/op_profile.proto
+++ b/tensorflow/contrib/tpu/profiler/op_profile.proto
@@ -18,13 +18,15 @@
message Node {
string name = 1; // Semantics depend on contents.
Metrics metrics = 2; // May be omitted e.g. for fused instructions.
- repeated Node children = 3;
+ repeated Node children = 3; // Subjected to pruning.
// Details about what this node represents.
oneof contents {
InstructionCategory category = 4;
XLAInstruction xla = 5;
}
+
+ int32 num_children = 6; // Total number of children before pruning.
// A category of XLA instructions.
// name is a descriptive string, like "data formatting".
message InstructionCategory {
@@ -64,8 +66,8 @@
// - it does not reveal the peak core FLOPS of the hardware
double flops = 2;
- // The VMEM bandwidth used to load operands from HBM, as a fraction of
- // thereotical VMEM bandwidth on the specific hardware.
+ // The memory bandwidth used to load operands, as a fraction of
+ // thereotical memory bandwidth on the specific hardware.
double memory_bandwidth = 3;
double raw_time = 11; // Elapsed core-time in picoseconds.
diff --git a/tensorflow/contrib/tpu/proto/optimization_parameters.proto b/tensorflow/contrib/tpu/proto/optimization_parameters.proto
index fc13205..8529b48 100644
--- a/tensorflow/contrib/tpu/proto/optimization_parameters.proto
+++ b/tensorflow/contrib/tpu/proto/optimization_parameters.proto
@@ -22,13 +22,22 @@
}
}
+// Each optimizer's parameter proto has a link to its documentation and CPU
+// implementation (if available) for user reference.
+
+// https://www.tensorflow.org/api_docs/python/tf/train/AdagradOptimizer
+// https://github.com/tensorflow/tensorflow/blob/c19e29306ce1777456b2dbb3a14f511edf7883a8/tensorflow/core/kernels/training_ops.cc#L151
message AdagradParameters {
float initial_accumulator = 1;
}
+// https://www.tensorflow.org/api_docs/python/tf/train/GradientDescentOptimizer
+// https://github.com/tensorflow/tensorflow/blob/c19e29306ce1777456b2dbb3a14f511edf7883a8/tensorflow/core/kernels/training_ops.cc#L423
message StochasticGradientDescentParameters {
}
+// https://www.tensorflow.org/api_docs/python/tf/train/FtrlOptimizer
+// https://github.com/tensorflow/tensorflow/blob/c19e29306ce1777456b2dbb3a14f511edf7883a8/tensorflow/core/kernels/training_ops.cc#L192
message FtrlParameters {
float l1 = 1;
float l2 = 2;
@@ -41,21 +50,42 @@
// learning rate feature instead, setting the learning rate to:
// user learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
// Here, t is the current timestep.
+//
+// https://www.tensorflow.org/api_docs/python/tf/train/AdamOptimizer
// https://github.com/tensorflow/tensorflow/blob/ab51450c817674c8ff08a7ae4f8ac50cdc4bed8b/tensorflow/python/training/adam.py#L54
+//
+// Note that the code by default implements the lazy version of Adam
+// (https://www.tensorflow.org/api_docs/python/tf/contrib/opt/LazyAdamOptimizer)
+// unless the use_non_lazy_adam parameter is set, in which case it implements
+// the normal version of Adam that updates all parameters in the embedding
+// table, even for entries that are not used in the current minibatch
+// (https://www.tensorflow.org/api_docs/python/tf/contrib/opt/AdamOptimizer). If
+// use_non_lazy_adam is enabled, use_gradient_accumulation is also required in
+// order to get correct results; a warning will be printed otherwise (which may
+// change to an error in the future). If use_max_with_epsilon is set, the Adam
+// variable update formula will be changed from m / (sqrt(v) + epsilon) to
+// m / max(sqrt(v), abs(epsilon)); this option improves the performance of TPU
+// training and is not expected to harm model quality.
message AdamParameters {
float beta1 = 3;
float beta2 = 4;
float epsilon = 5;
float initial_m = 6;
float initial_v = 7;
+ bool use_non_lazy_adam = 8;
+ bool use_max_with_epsilon = 9;
}
+// https://www.tensorflow.org/api_docs/python/tf/train/MomentumOptimizer
+// https://github.com/tensorflow/tensorflow/blob/c19e29306ce1777456b2dbb3a14f511edf7883a8/tensorflow/core/kernels/training_ops.cc#L271
message MomentumParameters {
float momentum = 1;
bool use_nesterov = 2;
float initial_accum = 3;
}
+// https://www.tensorflow.org/api_docs/python/tf/train/RMSPropOptimizer
+// https://github.com/tensorflow/tensorflow/blob/c19e29306ce1777456b2dbb3a14f511edf7883a8/tensorflow/core/kernels/training_ops.cc#L356
message RmsPropParameters {
float rho = 1;
float momentum = 2;
@@ -64,6 +94,8 @@
float initial_mom = 5;
}
+// https://www.tensorflow.org/api_docs/python/tf/train/RMSPropOptimizer
+// https://github.com/tensorflow/tensorflow/blob/c19e29306ce1777456b2dbb3a14f511edf7883a8/tensorflow/core/kernels/training_ops.cc#L372
message CenteredRmsPropParameters {
float rho = 1;
float momentum = 2;
@@ -73,6 +105,7 @@
float initial_mg = 6;
}
+// Variant of algorithm in http://proceedings.mlr.press/v44/shamir15.pdf
message MdlAdagradLightParameters {
float l2 = 1;
float lr_power = 2;
@@ -91,6 +124,8 @@
float initial_benefit = 15;
}
+// https://www.tensorflow.org/api_docs/python/tf/train/RMSPropOptimizer
+// https://github.com/tensorflow/tensorflow/blob/c19e29306ce1777456b2dbb3a14f511edf7883a8/tensorflow/core/kernels/training_ops.cc#L68
message AdadeltaParameters {
float rho = 1;
float epsilon = 2;
@@ -98,6 +133,8 @@
float initial_update = 4;
}
+// https://www.tensorflow.org/api_docs/python/tf/train/RMSPropOptimizer
+// https://github.com/tensorflow/tensorflow/blob/c19e29306ce1777456b2dbb3a14f511edf7883a8/tensorflow/core/kernels/training_ops.cc#L164
message ProximalAdagradParameters {
float l1 = 1;
float l2 = 2;
diff --git a/tensorflow/contrib/tpu/python/ops/tpu_ops.py b/tensorflow/contrib/tpu/python/ops/tpu_ops.py
index a1aee69..e2e4aca 100644
--- a/tensorflow/contrib/tpu/python/ops/tpu_ops.py
+++ b/tensorflow/contrib/tpu/python/ops/tpu_ops.py
@@ -200,6 +200,33 @@
return gen_tpu_ops.infeed_dequeue_tuple(dtypes, shapes, name=name)
# pylint: enable=redefined-outer-name
+ # pylint: disable=protected-access
+ def send_tpu_embedding_gradients(inputs,
+ config,
+ learning_rates=None,
+ name=None):
+ """A placeholder op for feeding per-sample gradients to the embedding layer.
+
+ Args:
+ inputs: A TensorList of gradients with which to update embedding tables.
+ Contains one tensor per embedding table in the model.
+ config: Serialized TPUEmbeddingConfiguration proto.
+ learning_rates: A TensorList of float32 scalars, one for each embedding
+ table, containing the learning rates for each table when dynamic
+ learning rate is enabled through the OptimizationParameters in
+ TPUEmbeddingConfiguration. When the learning rate is constant, the list
+ should be empty (optional).
+ name: A name for the operation (optional).
+
+ Returns:
+ A SendTPUEmbeddingGradients operation.
+ """
+ if learning_rates is None:
+ learning_rates = []
+ return gen_tpu_ops._send_tpu_embedding_gradients(
+ inputs=inputs, learning_rates=learning_rates, config=config, name=name)
+
+
else:
# We have already built the appropriate libraries into the binary via CMake
# if we have built contrib, so we don't need this
diff --git a/tensorflow/contrib/tpu/python/tpu/async_checkpoint.py b/tensorflow/contrib/tpu/python/tpu/async_checkpoint.py
new file mode 100644
index 0000000..20b7ba0
--- /dev/null
+++ b/tensorflow/contrib/tpu/python/tpu/async_checkpoint.py
@@ -0,0 +1,202 @@
+# 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.
+# ======================================
+"""Hook for asynchronous checkpointing.
+
+This hook dispatches checkpoint writing operations in a separate thread to
+allow execution to continue on the main thread.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import threading
+import time
+
+from tensorflow.core.util.event_pb2 import SessionLog
+from tensorflow.python.framework import meta_graph
+from tensorflow.python.framework import ops
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.training import basic_session_run_hooks
+from tensorflow.python.training import training_util
+from tensorflow.python.training.session_run_hook import SessionRunArgs
+from tensorflow.python.training.summary_io import SummaryWriterCache
+
+
+class AsyncCheckpointSaverHook(basic_session_run_hooks.CheckpointSaverHook):
+ """Saves checkpoints every N steps or seconds."""
+
+ def __init__(self,
+ checkpoint_dir,
+ save_secs=None,
+ save_steps=None,
+ saver=None,
+ checkpoint_basename="model.ckpt",
+ scaffold=None,
+ listeners=None):
+ """Initializes a `CheckpointSaverHook`.
+
+ Args:
+ checkpoint_dir: `str`, base directory for the checkpoint files.
+ save_secs: `int`, save every N secs.
+ save_steps: `int`, save every N steps.
+ saver: `Saver` object, used for saving.
+ checkpoint_basename: `str`, base name for the checkpoint files.
+ scaffold: `Scaffold`, use to get saver object.
+ listeners: List of `CheckpointSaverListener` subclass instances. Used for
+ callbacks that run immediately before or after this hook saves the
+ checkpoint.
+
+ Raises:
+ ValueError: One of `save_steps` or `save_secs` should be set.
+ ValueError: At most one of `saver` or `scaffold` should be set.
+ """
+ logging.info("Create AsyncCheckpointSaverHook.")
+ if saver is not None and scaffold is not None:
+ raise ValueError("You cannot provide both saver and scaffold.")
+ self._saver = saver
+ self._save_thread = None
+ self._checkpoint_dir = checkpoint_dir
+ self._save_path = os.path.join(checkpoint_dir, checkpoint_basename)
+ self._scaffold = scaffold
+ self._timer = basic_session_run_hooks.SecondOrStepTimer(
+ every_secs=save_secs, every_steps=save_steps)
+ self._listeners = listeners or []
+ self._steps_per_run = 1
+ self._summary_writer = None
+ self._global_step_tensor = None
+
+ def _set_steps_per_run(self, steps_per_run):
+ self._steps_per_run = steps_per_run
+
+ def begin(self):
+ self._summary_writer = SummaryWriterCache.get(self._checkpoint_dir)
+ self._global_step_tensor = training_util._get_or_create_global_step_read() # pylint: disable=protected-access
+ if self._global_step_tensor is None:
+ raise RuntimeError(
+ "Global step should be created to use CheckpointSaverHook.")
+ for l in self._listeners:
+ l.begin()
+
+ def after_create_session(self, session, coord):
+ global_step = session.run(self._global_step_tensor)
+
+ # We do write graph and saver_def at the first call of before_run.
+ # We cannot do this in begin, since we let other hooks to change graph and
+ # add variables in begin. Graph is finalized after all begin calls.
+ training_util.write_graph(
+ ops.get_default_graph().as_graph_def(add_shapes=True),
+ self._checkpoint_dir, "graph.pbtxt")
+ saver_def = self._get_saver().saver_def if self._get_saver() else None
+ graph = ops.get_default_graph()
+ meta_graph_def = meta_graph.create_meta_graph_def(
+ graph_def=graph.as_graph_def(add_shapes=True), saver_def=saver_def)
+ self._summary_writer.add_graph(graph)
+ self._summary_writer.add_meta_graph(meta_graph_def)
+ # The checkpoint saved here is the state at step "global_step".
+ self._save(session, global_step)
+ self._timer.update_last_triggered_step(global_step)
+
+ def before_run(self, run_context): # pylint: disable=unused-argument
+ return SessionRunArgs(self._global_step_tensor)
+
+ def after_run(self, run_context, run_values):
+ stale_global_step = run_values.results
+ if self._timer.should_trigger_for_step(stale_global_step +
+ self._steps_per_run):
+ # get the real value after train op.
+ global_step = run_context.session.run(self._global_step_tensor)
+ if self._timer.should_trigger_for_step(global_step):
+ self._timer.update_last_triggered_step(global_step)
+ if self._save(run_context.session, global_step):
+ run_context.request_stop()
+
+ def end(self, session):
+ if self._save_thread:
+ logging.info("Waiting for any pending checkpoints to finish.")
+ self._save_thread.join()
+
+ last_step = session.run(self._global_step_tensor)
+
+ # Save the last checkpoint synchronously if needed.
+ if last_step != self._timer.last_triggered_step():
+ self._save(session, last_step, asynchronous=False)
+
+ for l in self._listeners:
+ l.end(session, last_step)
+
+ def _save(self, session, step, asynchronous=True):
+ """Saves the latest checkpoint, returns should_stop."""
+
+ # Skip saving on step 0
+ if step == 0:
+ return
+
+ def _save_fn():
+ """Run the saver process."""
+ logging.info("Saving checkpoints for %d into %s.", step, self._save_path)
+
+ start_time = time.time()
+ for l in self._listeners:
+ l.before_save(session, step)
+
+ self._get_saver().save(session, self._save_path, global_step=step)
+ self._summary_writer.add_session_log(
+ SessionLog(
+ status=SessionLog.CHECKPOINT, checkpoint_path=self._save_path),
+ step)
+ end_time = time.time()
+ logging.info("Checkpoint actual writing time: (%.3f sec)",
+ end_time - start_time)
+ logging.info("Checkpoint finished for %d into %s.", step, self._save_path)
+
+ for l in self._listeners:
+ l.before_save(session, step)
+
+ if not asynchronous:
+ _save_fn()
+ return
+
+ if self._save_thread is not None:
+ self._save_thread.join(timeout=0.1)
+ if self._save_thread.is_alive():
+ logging.info("Saver thread still in progress, skipping checkpoint.")
+ return
+
+ self._save_thread = threading.Thread(target=_save_fn)
+ self._save_thread.start()
+
+ def _get_saver(self):
+ if self._saver is not None:
+ return self._saver
+ elif self._scaffold is not None:
+ return self._scaffold.saver
+
+ # Get saver from the SAVERS collection if present.
+ collection_key = ops.GraphKeys.SAVERS
+ savers = ops.get_collection(collection_key)
+ if not savers:
+ raise RuntimeError(
+ "No items in collection {}. Please add a saver to the collection "
+ "or provide a saver or scaffold.".format(collection_key))
+ elif len(savers) > 1:
+ raise RuntimeError(
+ "More than one item in collection {}. "
+ "Please indicate which one to use by passing it to the constructor."
+ .format(collection_key))
+
+ self._saver = savers[0]
+ return savers[0]
diff --git a/tensorflow/contrib/tpu/python/tpu/datasets.py b/tensorflow/contrib/tpu/python/tpu/datasets.py
index d879170..c694e9c 100644
--- a/tensorflow/contrib/tpu/python/tpu/datasets.py
+++ b/tensorflow/contrib/tpu/python/tpu/datasets.py
@@ -18,8 +18,8 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.contrib.data.python.ops import batching
-from tensorflow.contrib.data.python.ops import interleave_ops
+from tensorflow.python.data.experimental.ops import batching
+from tensorflow.python.data.experimental.ops import interleave_ops
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import iterator_ops
from tensorflow.python.data.ops import readers
diff --git a/tensorflow/contrib/tpu/python/tpu/keras_support.py b/tensorflow/contrib/tpu/python/tpu/keras_support.py
index f67e0e6..696656e 100644
--- a/tensorflow/contrib/tpu/python/tpu/keras_support.py
+++ b/tensorflow/contrib/tpu/python/tpu/keras_support.py
@@ -325,18 +325,6 @@
return KerasCrossShardOptimizer(opt)
-def _clone_metrics(metrics):
- """Returns a copy of metrics. A copy is created for stateful metrics."""
- if metrics is None:
- return None
- with variable_scope.variable_scope(
- 'metrics', reuse=variable_scope.AUTO_REUSE):
- return [
- m.__class__.from_config(m.get_config()) if isinstance(
- m, metrics_module.Metric) else m for m in metrics
- ]
-
-
def _clone_optimizer(optimizer, config=None):
"""Returns a cloned optimizer with the provided optimizer.config or config."""
if not isinstance(optimizer, keras_optimizers.Optimizer):
@@ -817,7 +805,8 @@
return input_specs, get_next_ops
-def _inject_tpu_inputs_for_infeed(tpu_assignment, mode, input_tensors, inputs):
+def _inject_tpu_inputs_for_infeed(tpu_assignment, mode,
+ core_id_place_holder, input_tensors, inputs):
"""Append core information to the set of inputs."""
# This is used during compilation to identify the current TPU core and enable
# concatenation operations across cores.
@@ -825,8 +814,6 @@
return input_tensors, inputs
# Puts a place holder in input spec.
- core_id_place_holder = array_ops.placeholder(
- dtype=dtypes.int32, shape=[1], name='core_id')
input_tensors = [core_id_place_holder] + input_tensors
# Now fill the core id. For `num_cores` = 2, `batch_size` = 8, we fill the
@@ -874,6 +861,10 @@
self._compilation_cache = {}
self._cloned_model = None
self._cloned_optimizer = None
+ # Create a placeholder for the TPU core ID. Cache the placeholder to avoid
+ # modifying the graph for every batch.
+ self._core_id_place_holder = array_ops.placeholder(
+ dtype=dtypes.int32, shape=[1], name='core_id')
def _specialize_model(self, input_specs, infeed_manager):
"""Specialize `self.model` (a Keras model) for the given input shapes."""
@@ -960,14 +951,24 @@
optimizer=_replicated_optimizer(self._cloned_optimizer),
loss=self.model.loss,
loss_weights=self.model.loss_weights,
- metrics=_clone_metrics(self.model.metrics),
- weighted_metrics=_clone_metrics(self.model.weighted_metrics),
+ metrics=metrics_module.clone_metrics(self.model.metrics),
+ weighted_metrics=metrics_module.clone_metrics(
+ self.model.weighted_metrics),
target_tensors=tpu_targets,
)
# Compute our outfeed depending on the execution mode
if is_training:
- self._cloned_model._make_train_function()
+ if not isinstance(self._cloned_optimizer, keras_optimizers.TFOptimizer):
+ # For Keras optimizer, we try to place the variable weights on the TPU
+ # device. Keras creates optimizer variables (e.g. momentum values for
+ # the Momentum optimizer) when _make_train_function is invoked.
+ with keras_tpu_variables.replicated_variable_for_optimizer(
+ self._tpu_assignment.num_towers):
+ self._cloned_model._make_train_function()
+ else:
+ self._cloned_model._make_train_function()
+
self._outfeed_spec = [
tensor_spec.TensorSpec(tensor.shape, tensor.dtype, tensor.name)
for tensor in self._cloned_model.train_function.outputs
@@ -1141,7 +1142,8 @@
inputs = inputs[:len(input_tensors)]
input_tensors, inputs = (
_inject_tpu_inputs_for_infeed(
- self._tpu_assignment, self.execution_mode, input_tensors, inputs))
+ self._tpu_assignment, self.execution_mode,
+ self._core_id_place_holder, input_tensors, inputs))
return input_tensors, inputs
def _process_outputs(self, outfeed_outputs):
@@ -1360,13 +1362,9 @@
raise ValueError('target_tensors is not supported for TPU execution.')
self._cpu_model.compile(
- _clone_optimizer(optimizer),
- loss,
- _clone_metrics(metrics),
- loss_weights,
- sample_weight_mode,
- _clone_metrics(weighted_metrics),
- target_tensors,
+ _clone_optimizer(optimizer), loss,
+ metrics_module.clone_metrics(metrics), loss_weights, sample_weight_mode,
+ metrics_module.clone_metrics(weighted_metrics), target_tensors,
**kwargs)
super(KerasTPUModel, self).compile(optimizer, loss, metrics, loss_weights,
@@ -2122,10 +2120,10 @@
cpu_model.compile(
_clone_optimizer(model.optimizer, optimizer_config),
model.loss,
- _clone_metrics(model.metrics),
+ metrics_module.clone_metrics(model.metrics),
model.loss_weights,
model.sample_weight_mode,
- _clone_metrics(model.weighted_metrics),
+ metrics_module.clone_metrics(model.weighted_metrics),
)
if model_weights:
diff --git a/tensorflow/contrib/tpu/python/tpu/keras_tpu_variables.py b/tensorflow/contrib/tpu/python/tpu/keras_tpu_variables.py
index 170977d..004b101 100644
--- a/tensorflow/contrib/tpu/python/tpu/keras_tpu_variables.py
+++ b/tensorflow/contrib/tpu/python/tpu/keras_tpu_variables.py
@@ -25,10 +25,15 @@
import contextlib
+import numpy as np
+
from tensorflow.python.client import session as session_lib
+from tensorflow.python.framework import dtypes as dtypes_module
from tensorflow.python.framework import ops
+from tensorflow.python.keras import backend
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import gen_resource_variable_ops
+from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import variable_scope
@@ -73,7 +78,7 @@
if tpu_context is None:
return self._primary_var.handle
- return tpu_context.get_replicated_var_handle(self)
+ return tpu_context.get_replicated_var_handle(self._name, self._vars)
@contextlib.contextmanager
def _assign_dependencies(self):
@@ -285,3 +290,51 @@
return variable_scope.variable_scope(
"", custom_getter=_replicated_variable_getter)
+
+
+@contextlib.contextmanager
+def replicated_variable_for_optimizer(num_replicas):
+ """Context manager for optimizer weights. Overrides K.variable."""
+ if num_replicas == 1:
+ yield
+ return
+
+ try:
+ old_v = backend.variable
+
+ def opt_variable(value, dtype=None, name=None, constraint=None):
+ """Instantiates a variable and returns it."""
+ if dtype is None:
+ dtype = backend.floatx()
+
+ variables = []
+ for i in range(num_replicas):
+ # Keras holds the variables in optimizer class instance , so the name
+ # does not matter here. ResourceVariable constructor will find a unique
+ # name (including name=None) for each replica.
+ with ops.device("device:TPU:{}".format(i)):
+ v = resource_variable_ops.ResourceVariable(
+ value,
+ dtype=dtypes_module.as_dtype(dtype),
+ name=name,
+ constraint=constraint)
+ variables.append(v)
+ name = "replicate_{}_{}".format("variable" if name is None else name,
+ ops.uid())
+ v = ReplicatedVariable(name, variables)
+
+ # pylint: disable=protected-access
+
+ if isinstance(value, np.ndarray):
+ v._keras_shape = value.shape
+ elif hasattr(value, "shape"):
+ v._keras_shape = backend.int_shape(value)
+ v._uses_learning_phase = False
+ backend.track_variable(v)
+ return v
+
+ backend.variable = opt_variable
+ yield
+
+ finally:
+ backend.variable = old_v
diff --git a/tensorflow/contrib/tpu/python/tpu/session_support.py b/tensorflow/contrib/tpu/python/tpu/session_support.py
index 3e91e2d..05264f5 100644
--- a/tensorflow/contrib/tpu/python/tpu/session_support.py
+++ b/tensorflow/contrib/tpu/python/tpu/session_support.py
@@ -41,6 +41,29 @@
pass
+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 ops.device(device):
+ heartbeat_op = tpu_ops.worker_heartbeat('')
+ options = config_pb2.RunOptions(timeout_in_ms=5000)
+ temp_session.run(heartbeat_op, options=options)
+ except errors.InvalidArgumentError as _:
+ logging.warning('Error running heartbeat on %s', device)
+ return None
+ except errors.DeadlineExceededError as _:
+ logging.warning('Timeout connecting to %s when testing heartbeat', device)
+ return None
+
+ # If we successfully connected and pinged the worker, go ahead and construct
+ # the operation.
+ with ops.device(device):
+ return tpu_ops.worker_heartbeat(request_ph)
+
+
class WorkerHeartbeatManager(object):
"""Manages the status/heartbeat monitor for a set of workers."""
@@ -72,30 +95,27 @@
name='worker_heartbeat_request', dtype=dtypes.string)
heartbeat_ops = []
+ kept_devices = []
for device in devices:
- with ops.device(device):
- heartbeat_ops.append(tpu_ops.worker_heartbeat(request_placeholder))
+ heartbeat_op = _make_heartbeat_op(session, device, request_placeholder)
+ if heartbeat_op is not None:
+ kept_devices.append(device)
+ heartbeat_ops.append(heartbeat_op)
+ else:
+ logging.warning('Heartbeat support not available for %s', device)
- return WorkerHeartbeatManager(session, devices, heartbeat_ops,
+ return WorkerHeartbeatManager(session, kept_devices, heartbeat_ops,
request_placeholder)
- def heartbeat_supported(self):
- """Returns True if heartbeat operations are supported on all workers."""
- try:
- # Send ping to verify worker has heartbeat support.
- self.ping()
- return True
- except errors.InvalidArgumentError as _:
- return False
+ def num_workers(self):
+ return len(self._devices)
def configure(self, message):
"""Configure heartbeat manager for all devices.
Args:
message: `event_pb2.WorkerHeartbeatRequest`
-
Returns: `None`
-
"""
logging.info('Configuring worker heartbeat: %s',
text_format.MessageToString(message))
@@ -155,7 +175,7 @@
def all_worker_devices(session):
"""Return a list of devices for each worker in the system."""
devices = session.list_devices()
- return [device.name for device in devices if 'CPU' in device.name]
+ return [device.name for device in devices if ':CPU:' in device.name]
class WatchdogManager(threading.Thread):
@@ -184,7 +204,6 @@
"""Initialize a watchdog manager.
Args:
-
session: Session connected to worker devices. A cloned session and graph
will be created for managing worker pings.
devices: Set of devices to monitor. If none, all workers will be
@@ -277,16 +296,14 @@
target=training_session.sess_str, graph=self._graph)
self._workers = WorkerHeartbeatManager.from_devices(
self._session, all_worker_devices(self._session))
- self._heartbeat_supported = self._workers.heartbeat_supported()
+ self._heartbeat_supported = self._workers.num_workers() > 0
if self._heartbeat_supported:
self._workers.configure(
event_pb2.WorkerHeartbeatRequest(
shutdown_mode=event_pb2.WAIT_FOR_COORDINATOR))
else:
logging.warn(
- 'Worker heartbeats not supported by all workers. No failure '
- 'handling will be enabled.'
- )
+ 'No workers support hearbeats. Failure handling will be disabled.')
def saver(self):
if self._saver:
@@ -303,8 +320,7 @@
logging.error(
'Multiple savers in the SAVERS collection. On-demand checkpointing '
'will be disabled. Pass an explicit `saver` to the constructor to '
- 'override this behavior.'
- )
+ 'override this behavior.')
return None
return savers[0]
diff --git a/tensorflow/contrib/tpu/python/tpu/tpu.py b/tensorflow/contrib/tpu/python/tpu/tpu.py
index 712b02f..11aaa1c 100644
--- a/tensorflow/contrib/tpu/python/tpu/tpu.py
+++ b/tensorflow/contrib/tpu/python/tpu/tpu.py
@@ -155,19 +155,20 @@
self._pivot = pivot
self._replicated_vars = {}
- def get_replicated_var_handle(self, var):
+ def get_replicated_var_handle(self, name, vars_):
"""Returns a variable handle for replicated TPU variable 'var'.
This is a method used by an experimental replicated variable implementation
and is not intended as a public API.
Args:
- var: The replicated TPU variable.
+ name: The common name of the variable.
+ vars_: The replicated TPU variables.
Returns:
The handle of the TPU replicated input node.
"""
- handle = self._replicated_vars.get(var)
+ handle = self._replicated_vars.get(name)
if handle is not None:
return handle
@@ -183,10 +184,10 @@
saved_context = graph._get_control_flow_context()
graph._set_control_flow_context(self.outer_context)
handle = tpu_ops.tpu_replicated_input(
- [v.handle for v in var._vars], name=var.name + "/handle")
+ [v.handle for v in vars_], name=name + "/handle")
graph._set_control_flow_context(saved_context)
# pylint: enable=protected-access
- self._replicated_vars[var] = handle
+ self._replicated_vars[name] = handle
return handle
def report_unsupported_operations(self):
@@ -661,6 +662,10 @@
# be less confusing to clients if they knowingly choose to use resource
# variables.
# Partitioned variables is not supported (b/112311320).
+ vscope = variable_scope.get_variable_scope()
+ saved_use_resource = vscope.use_resource
+ saved_custom_getter = vscope.custom_getter
+
def custom_getter(getter, name, *args, **kwargs):
"""Variables on TPU have a few restrictions."""
partitioner = kwargs["partitioner"]
@@ -671,12 +676,10 @@
"`partitioner` that is {} for variable {}. "
"Setting `partitioner` to `None`."
.format(partitioner, name))
- return getter(name, *args, **kwargs)
-
- vscope = variable_scope.get_variable_scope()
-
- saved_use_resource = vscope.use_resource
- saved_custom_getter = vscope.custom_getter
+ if saved_custom_getter is None:
+ return getter(name, *args, **kwargs)
+ else:
+ return saved_custom_getter(getter, name, *args, **kwargs)
vscope.set_use_resource(True)
vscope.set_custom_getter(custom_getter)
diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_context.py b/tensorflow/contrib/tpu/python/tpu/tpu_context.py
index 7cfb6c3..da6bdf6 100644
--- a/tensorflow/contrib/tpu/python/tpu/tpu_context.py
+++ b/tensorflow/contrib/tpu/python/tpu/tpu_context.py
@@ -154,6 +154,20 @@
# as far as model is replicated to all cores in the system.
return self._internal_ctx.device_for_replica(replica_id)
+ @property
+ def tpu_host_placement_function(self):
+ """Returns the TPU host place function.
+
+ The place function takes host_id as the input and returns the TF device
+ for the correspoding host.
+ """
+
+ def _placement_function(host_id):
+ """Return the host device given host_id."""
+ return self._internal_ctx.tpu_host_placement_function(host_id=host_id)
+
+ return _placement_function
+
class _InternalTPUContext(object):
"""A context holds immutable states of TPU computation.
diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py
index 23c5451..3aa5b6e 100644
--- a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py
+++ b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py
@@ -231,7 +231,7 @@
`metric_fn` runs on CPU to generate metrics and `tensors` represents the
`Tensor`s transferred from TPU system to CPU host and passed to `metric_fn`.
To be precise, TPU evaluation expects a slightly different signature from the
- @{tf.estimator.Estimator}. While `EstimatorSpec.eval_metric_ops` expects a
+ `tf.estimator.Estimator`. While `EstimatorSpec.eval_metric_ops` expects a
dict, `TPUEstimatorSpec.eval_metrics` is a tuple of `metric_fn` and `tensors`.
The `tensors` could be a list of `Tensor`s or dict of names to `Tensor`s. The
`tensors` usually specify the model logits, which are transferred back from
@@ -254,7 +254,7 @@
sending tensors from TPU to CPU. To reduce the overhead, try reducing the
size of the tensors. The `tensors` are concatenated along their major (batch)
dimension, and so must be >= rank 1. The `host_call` is useful for writing
- summaries with @{tf.contrib.summary.create_file_writer}.
+ summaries with `tf.contrib.summary.create_file_writer`.
"""
def __new__(cls,
@@ -404,12 +404,17 @@
self._feed_error = None
self._finished = False
+ self._should_initialize_tpu = True
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 = [tpu.initialize_system(job=self._master_job)]
- self._finalize_ops = [tpu.shutdown_system(job=self._master_job)]
+ 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()
self._init_ops.extend(summary_writer_init_ops)
@@ -421,10 +426,10 @@
def _run_infeed(self, queue_ctx, session):
logging.info('Starting infeed thread controller.')
if self._initial_infeed_sleep_secs:
- logging.info('%s thread sleeping for %d seconds.', self._name,
+ logging.info('Infeed thread sleeping for %d seconds.',
self._initial_infeed_sleep_secs)
time.sleep(self._initial_infeed_sleep_secs)
- logging.info('%s thread starting after sleep', self._name)
+ logging.info('Infeed thread starting after sleep')
with self._rendezvous.catch_errors(source='infeed', session=session):
if self._run_infeed_loop_on_coordinator:
diff --git a/tensorflow/contrib/tpu/tpu_estimator.md b/tensorflow/contrib/tpu/tpu_estimator.md
index 639e708..b6514e1 100644
--- a/tensorflow/contrib/tpu/tpu_estimator.md
+++ b/tensorflow/contrib/tpu/tpu_estimator.md
@@ -87,7 +87,7 @@
label = tf.cast(features["label"], tf.int32)
return image, label
- dataset = tf.contrib.data.TFRecordDataset(
+ dataset = tf.data.TFRecordDataset(
filename, buffer_size=FLAGS.dataset_reader_buffer_size)
dataset = dataset.map(parser).cache().repeat().batch(batch_size)
images, labels = dataset.make_one_shot_iterator().get_next()
diff --git a/tensorflow/contrib/tpu/utils/BUILD b/tensorflow/contrib/tpu/utils/BUILD
new file mode 100644
index 0000000..c27b737
--- /dev/null
+++ b/tensorflow/contrib/tpu/utils/BUILD
@@ -0,0 +1,30 @@
+# Description: Utilities for TPU Operations
+
+licenses(["notice"]) # Apache 2.0
+
+cc_library(
+ name = "tpu_embedding_optimization_parameters_utils",
+ srcs = ["tpu_embedding_optimization_parameters_utils.cc"],
+ hdrs = ["tpu_embedding_optimization_parameters_utils.h"],
+ visibility = ["//visibility:public"],
+ deps = [
+ "//tensorflow/contrib/tpu/proto:optimization_parameters_proto_cc",
+ "//tensorflow/core:framework_headers_lib",
+ "//tensorflow/core:lib_proto_parsing",
+ "@com_google_absl//absl/base",
+ ],
+)
+
+cc_library(
+ name = "tpu_embedding_output_layout_utils",
+ srcs = ["tpu_embedding_output_layout_utils.cc"],
+ hdrs = ["tpu_embedding_output_layout_utils.h"],
+ visibility = ["//visibility:public"],
+ deps = [
+ "//tensorflow/contrib/tpu/proto:tpu_embedding_configuration_proto_cc",
+ "//tensorflow/contrib/tpu/proto:tpu_embedding_output_layout_proto_cc",
+ "//tensorflow/core:framework_headers_lib",
+ "//tensorflow/core:lib_proto_parsing",
+ "//tensorflow/core:protos_all_cc",
+ ],
+)
diff --git a/tensorflow/contrib/tpu/utils/tpu_embedding_optimization_parameters_utils.cc b/tensorflow/contrib/tpu/utils/tpu_embedding_optimization_parameters_utils.cc
new file mode 100644
index 0000000..76cb553
--- /dev/null
+++ b/tensorflow/contrib/tpu/utils/tpu_embedding_optimization_parameters_utils.cc
@@ -0,0 +1,255 @@
+/* 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/contrib/tpu/utils/tpu_embedding_optimization_parameters_utils.h"
+#include "tensorflow/core/lib/core/errors.h"
+
+namespace tensorflow {
+namespace tpu {
+
+string GetOptimizationAlgorithmName(OptimizationAlgorithm alg) {
+ switch (alg) {
+ case OptimizationAlgorithm::kAdagrad:
+ return "Adagrad";
+ case OptimizationAlgorithm::kStochasticGradientDescent:
+ return "StochasticGradientDescent";
+ case OptimizationAlgorithm::kFtrl:
+ return "FTRL";
+ case OptimizationAlgorithm::kAdam:
+ return "ADAM";
+ case OptimizationAlgorithm::kMomentum:
+ return "Momentum";
+ case OptimizationAlgorithm::kRmsProp:
+ return "RMSProp";
+ case OptimizationAlgorithm::kCenteredRmsProp:
+ return "CenteredRMSProp";
+ case OptimizationAlgorithm::kMdlAdagradLight:
+ return "MDLAdagradLight";
+ case OptimizationAlgorithm::kAdadelta:
+ return "Adadelta";
+ case OptimizationAlgorithm::kProximalAdagrad:
+ return "ProximalAdagrad";
+ case OptimizationAlgorithm::PARAMETERS_NOT_SET:
+ return "*** Not set ***";
+ }
+}
+
+string GetOptimizationAlgorithmFriendlyName(OptimizationAlgorithm alg) {
+ switch (alg) {
+ case OptimizationAlgorithm::kAdagrad:
+ return "Adagrad";
+ case OptimizationAlgorithm::kStochasticGradientDescent:
+ return "stochastic gradient descent";
+ case OptimizationAlgorithm::kFtrl:
+ return "FTRL";
+ case OptimizationAlgorithm::kAdam:
+ return "ADAM";
+ case OptimizationAlgorithm::kMomentum:
+ return "Momentum";
+ case OptimizationAlgorithm::kRmsProp:
+ return "RMSProp";
+ case OptimizationAlgorithm::kCenteredRmsProp:
+ return "centered RMSProp";
+ case OptimizationAlgorithm::kMdlAdagradLight:
+ return "MDL Adagrad Light";
+ case OptimizationAlgorithm::kAdadelta:
+ return "Adadelta";
+ case OptimizationAlgorithm::kProximalAdagrad:
+ return "proximal Adagrad";
+ case OptimizationAlgorithm::PARAMETERS_NOT_SET:
+ return "unknown (not specified)";
+ }
+}
+
+// Returns the number of optimization parameter vectors used by the optimization
+// algorithm, excluding the weights themselves and assuming no gradient
+// accumulation.
+Status GetBaseAuxiliaryParameterCount(OptimizationAlgorithm alg, int* count) {
+ switch (alg) {
+ case OptimizationAlgorithm::kAdagrad:
+ *count = 1;
+ return Status::OK();
+ case OptimizationAlgorithm::kStochasticGradientDescent:
+ *count = 0;
+ return Status::OK();
+ case OptimizationAlgorithm::kFtrl:
+ *count = 2;
+ return Status::OK();
+ case OptimizationAlgorithm::kAdam:
+ *count = 2;
+ return Status::OK();
+ case OptimizationAlgorithm::kMomentum:
+ *count = 1;
+ return Status::OK();
+ case OptimizationAlgorithm::kRmsProp:
+ *count = 2;
+ return Status::OK();
+ case OptimizationAlgorithm::kCenteredRmsProp:
+ *count = 3;
+ return Status::OK();
+ case OptimizationAlgorithm::kMdlAdagradLight:
+ *count = 3;
+ return Status::OK();
+ case OptimizationAlgorithm::kAdadelta:
+ *count = 2;
+ return Status::OK();
+ case OptimizationAlgorithm::kProximalAdagrad:
+ *count = 1;
+ return Status::OK();
+ case OptimizationAlgorithm::PARAMETERS_NOT_SET:
+ return errors::InvalidArgument("No optimization algorithm specified");
+ }
+}
+
+Status GetGradientAccumulationSupport(OptimizationAlgorithm alg,
+ GradientAccumulationSupport* support) {
+ switch (alg) {
+ case OptimizationAlgorithm::kAdagrad:
+ *support = GradientAccumulationSupport::kSupported;
+ return Status::OK();
+ case OptimizationAlgorithm::kStochasticGradientDescent:
+ *support = GradientAccumulationSupport::kUnnecessary;
+ return Status::OK();
+ default: {
+ int auxiliary_parameter_count;
+ TF_RETURN_IF_ERROR(
+ GetBaseAuxiliaryParameterCount(alg, &auxiliary_parameter_count));
+ *support = auxiliary_parameter_count + 1 <= kMaxAuxiliaryParameterCount
+ ? GradientAccumulationSupport::kSupported
+ : GradientAccumulationSupport::kNotSupported;
+ return Status::OK();
+ }
+ }
+}
+namespace {
+// Make a normal state variable specification.
+StateVariableSpecification MakeStandardStateVariableSpecification(
+ const string& name) {
+ StateVariableSpecification result;
+ result.set_name(name);
+ result.mutable_user_defined();
+ return result;
+}
+} // namespace
+
+Status GetOptimizationAlgorithmStateVariables(
+ OptimizationAlgorithm alg, bool use_gradient_accumulation,
+ std::vector<StateVariableSpecification>* state_variables) {
+ // The first parameter set is always the weights themselves.
+ state_variables->push_back(
+ MakeStandardStateVariableSpecification("parameters"));
+ // The order of the returned parameters needs to match the offsets used by
+ // the algorithm implementations in test_util.cc and
+ // address_handler_program_creator.cc.
+ switch (alg) {
+ case OptimizationAlgorithm::kAdagrad: {
+ state_variables->push_back(
+ MakeStandardStateVariableSpecification("accumulators"));
+ break;
+ }
+ case OptimizationAlgorithm::kStochasticGradientDescent: {
+ // None.
+ break;
+ }
+ case OptimizationAlgorithm::kFtrl: {
+ state_variables->push_back(
+ MakeStandardStateVariableSpecification("accumulators"));
+ state_variables->push_back(
+ MakeStandardStateVariableSpecification("linears"));
+ break;
+ }
+ case OptimizationAlgorithm::kAdam: {
+ state_variables->push_back(
+ MakeStandardStateVariableSpecification("momenta"));
+ state_variables->push_back(
+ MakeStandardStateVariableSpecification("velocities"));
+ break;
+ }
+ case OptimizationAlgorithm::kMomentum: {
+ state_variables->push_back(
+ MakeStandardStateVariableSpecification("momenta"));
+ break;
+ }
+ case OptimizationAlgorithm::kRmsProp: {
+ state_variables->push_back(MakeStandardStateVariableSpecification("ms"));
+ state_variables->push_back(MakeStandardStateVariableSpecification("mom"));
+ break;
+ }
+ case OptimizationAlgorithm::kCenteredRmsProp: {
+ state_variables->push_back(MakeStandardStateVariableSpecification("ms"));
+ state_variables->push_back(MakeStandardStateVariableSpecification("mom"));
+ state_variables->push_back(MakeStandardStateVariableSpecification("mg"));
+ break;
+ }
+ case OptimizationAlgorithm::kMdlAdagradLight: {
+ state_variables->push_back(
+ MakeStandardStateVariableSpecification("accumulators"));
+ state_variables->push_back(
+ MakeStandardStateVariableSpecification("weights"));
+ state_variables->push_back(
+ MakeStandardStateVariableSpecification("benefits"));
+ break;
+ }
+ case OptimizationAlgorithm::kAdadelta: {
+ state_variables->push_back(
+ MakeStandardStateVariableSpecification("accumulators"));
+ state_variables->push_back(
+ MakeStandardStateVariableSpecification("updates"));
+ break;
+ }
+ case OptimizationAlgorithm::kProximalAdagrad: {
+ state_variables->push_back(
+ MakeStandardStateVariableSpecification("accumulators"));
+ break;
+ }
+ case OptimizationAlgorithm::PARAMETERS_NOT_SET: {
+ return errors::InvalidArgument("No optimization algorithm specified");
+ }
+ }
+ // This needs to be last so that the save/restore ops do not need to know
+ // about gradient accumulation.
+ if (use_gradient_accumulation) {
+ StateVariableSpecification gradient_acc;
+ gradient_acc.set_name("gradient_accumulators");
+ gradient_acc.mutable_fill_with_constant()->set_initial_value(
+ kGradientAccumulatorInitialValue);
+ state_variables->push_back(std::move(gradient_acc));
+ }
+ if (state_variables->size() > kMaxAuxiliaryParameterCount + 1) {
+ return errors::InvalidArgument(
+ "Optimization algorithm", GetOptimizationAlgorithmName(alg),
+ "does not support gradient accumulation because it "
+ "already has too many other accumulators");
+ }
+ return Status::OK();
+} // namespace tpu
+
+std::vector<OptimizationAlgorithm> GetOptimizationAlgorithms() {
+ return {
+ OptimizationAlgorithm::kAdagrad,
+ OptimizationAlgorithm::kStochasticGradientDescent,
+ OptimizationAlgorithm::kFtrl,
+ OptimizationAlgorithm::kAdam,
+ OptimizationAlgorithm::kMomentum,
+ OptimizationAlgorithm::kRmsProp,
+ OptimizationAlgorithm::kCenteredRmsProp,
+ OptimizationAlgorithm::kMdlAdagradLight,
+ OptimizationAlgorithm::kAdadelta,
+ OptimizationAlgorithm::kProximalAdagrad,
+ };
+}
+
+} // namespace tpu
+} // namespace tensorflow
diff --git a/tensorflow/contrib/tpu/utils/tpu_embedding_optimization_parameters_utils.h b/tensorflow/contrib/tpu/utils/tpu_embedding_optimization_parameters_utils.h
new file mode 100644
index 0000000..81d5026
--- /dev/null
+++ b/tensorflow/contrib/tpu/utils/tpu_embedding_optimization_parameters_utils.h
@@ -0,0 +1,90 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CONTRIB_TPU_UTILS_TPU_EMBEDDING_OPTIMIZATION_PARAMETERS_UTILS_H_
+#define TENSORFLOW_CONTRIB_TPU_UTILS_TPU_EMBEDDING_OPTIMIZATION_PARAMETERS_UTILS_H_
+
+#include <string>
+#include "absl/base/casts.h"
+#include "tensorflow/contrib/tpu/proto/optimization_parameters.pb.h"
+#include "tensorflow/core/lib/core/status.h"
+
+namespace tensorflow {
+namespace tpu {
+
+using OptimizationAlgorithm = OptimizationParameters::ParametersCase;
+
+// Returns the name of the optimization algorithm.
+string GetOptimizationAlgorithmName(OptimizationAlgorithm alg);
+
+// Returns a user-friendly name for the optimization algorithm.
+string GetOptimizationAlgorithmFriendlyName(OptimizationAlgorithm alg);
+
+// Returns all supported optimization algorithms.
+std::vector<OptimizationAlgorithm> GetOptimizationAlgorithms();
+
+enum class GradientAccumulationSupport {
+ // Accumulation cannot be used with this optimizer.
+ kNotSupported,
+
+ // Accumulation is unnecessary because optimizer application is commutative.
+ kUnnecessary,
+
+ // Accumulation is allowed and changes optimizer behavior.
+ kSupported,
+};
+
+// Returns the number of optimization parameter vectors used by the optimization
+// algorithm, excluding the weights themselves and assuming no gradient
+// accumulation.
+Status GetBaseAuxiliaryParameterCount(OptimizationAlgorithm alg, int *count);
+
+// Returns whether (and how) an optimization algorithm supports gradient
+// accumulation.
+Status GetGradientAccumulationSupport(OptimizationAlgorithm alg,
+ GradientAccumulationSupport *support);
+
+// Returns the parameter specifications for the optimization algorithm (the main
+// parameters first, followed by any auxiliary parameters such as Adagrad
+// accumulators).
+Status GetOptimizationAlgorithmStateVariables(
+ OptimizationAlgorithm alg, bool use_gradient_accumulation,
+ std::vector<StateVariableSpecification> *state_variables);
+
+// Maximum value of auxiliar_parameter_count for any optimization algorithm.
+static constexpr int kMaxAuxiliaryParameterCount = 3;
+
+// Fill value for gradient accumulators. This is a denormal so that it will be
+// flushed to zero on the current TPU platforms and needs to continue to have
+// the following properties in the future:
+//
+// 1. Does not have the same bit pattern as a zero and can be distinguished from
+// it using integer operations.
+// 2. Treated as zero by floating-point arithmetic operations (at least addition
+// and subtraction).
+// 3. Cannot be produced by any floating-point arithmetic operation, including
+// those involving itself.
+//
+// It does not need to compare equal or not equal to zero in floating point. We
+// need to use a non-zero value here because some optimization algorithms are
+// not no-ops on zero gradients, so we need to distinguish an accumulated
+// gradient of zero from one that has been cleared after its gradients have
+// already been applied to the parameters and accumulators.
+const float kGradientAccumulatorInitialValue = absl::bit_cast<float, uint32>(1);
+
+} // namespace tpu
+} // namespace tensorflow
+
+#endif // TENSORFLOW_CONTRIB_TPU_UTILS_TPU_EMBEDDING_OPTIMIZATION_PARAMETERS_UTILS_H_
diff --git a/tensorflow/contrib/tpu/utils/tpu_embedding_output_layout_utils.cc b/tensorflow/contrib/tpu/utils/tpu_embedding_output_layout_utils.cc
new file mode 100644
index 0000000..8480ec4
--- /dev/null
+++ b/tensorflow/contrib/tpu/utils/tpu_embedding_output_layout_utils.cc
@@ -0,0 +1,98 @@
+/* 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/contrib/tpu/utils/tpu_embedding_output_layout_utils.h"
+#include "tensorflow/contrib/tpu/proto/tpu_embedding_output_layout.pb.h"
+#include "tensorflow/core/lib/core/errors.h"
+
+namespace tensorflow {
+namespace tpu {
+
+void AddDefaultEmbeddingOutputLayoutIfNeeded(
+ TPUEmbeddingConfiguration* config) {
+ if (config->has_output_layout()) {
+ // Model or previous step has already filled this in.
+ return;
+ }
+
+ TPUEmbeddingOutputLayout* layout = config->mutable_output_layout();
+ // Create output tensors.
+ for (const auto& table : config->table_descriptor()) {
+ TPUEmbeddingOutputLayout::EmbeddingOutputTensor* output =
+ layout->add_output();
+ TPUEmbeddingOutputLayout::TwoDOutputTensor* two_d = output->mutable_two_d();
+ two_d->set_dim1_size(table.dimension());
+ two_d->set_dim0_size_per_sample(table.num_features());
+ }
+
+ // Create table output locations.
+ for (int table_id = 0; table_id < config->table_descriptor_size();
+ ++table_id) {
+ TPUEmbeddingOutputLayout::TableDescriptor* output_table =
+ layout->add_table();
+ const auto& table = config->table_descriptor(table_id);
+ for (int feature_index = 0; feature_index < table.num_features();
+ ++feature_index) {
+ TPUEmbeddingOutputLayout::FeatureDescriptor* output_feature =
+ output_table->add_feature();
+ TPUEmbeddingOutputLayout::OutputLocation* output_location =
+ output_feature->add_output_location();
+ output_location->set_tensor_index(table_id);
+ output_location->set_dim0_offset(feature_index);
+ output_location->set_dim1_offset(0);
+ }
+ }
+}
+
+Status ComputeOutputTensorShapes(const TPUEmbeddingConfiguration& config,
+ std::vector<TensorShapeProto>* shapes) {
+ if (!config.has_output_layout()) {
+ return errors::InvalidArgument(
+ "TPUEmbeddingConfiguration is missing output layout.");
+ }
+ const TPUEmbeddingOutputLayout& layout = config.output_layout();
+ int batch_size = config.batch_size_per_tensor_core();
+
+ for (int i = 0; i < layout.output_size(); ++i) {
+ const auto& output = layout.output(i);
+ TensorShapeProto shape;
+ switch (output.output_format_case()) {
+ case TPUEmbeddingOutputLayout::EmbeddingOutputTensor::OutputFormatCase::
+ kTwoD: {
+ auto* dim0 = shape.add_dim();
+ dim0->set_size(output.two_d().dim0_size_per_sample() * batch_size);
+ auto* dim1 = shape.add_dim();
+ dim1->set_size(output.two_d().dim1_size());
+ break;
+ }
+ case TPUEmbeddingOutputLayout::EmbeddingOutputTensor::OutputFormatCase::
+ OUTPUT_FORMAT_NOT_SET: {
+ return errors::InvalidArgument(
+ "Output layout in TPUEmbeddingConfiguration has unset embedding "
+ "output tensor format.");
+ }
+ default: {
+ return errors::InvalidArgument(
+ "Output layout in TPUEmbeddingConfiguration has invalid or "
+ "unhandled embedding output tensor format.");
+ }
+ }
+ shapes->push_back(shape);
+ }
+ return Status::OK();
+}
+
+} // namespace tpu
+} // namespace tensorflow
diff --git a/tensorflow/contrib/tpu/utils/tpu_embedding_output_layout_utils.h b/tensorflow/contrib/tpu/utils/tpu_embedding_output_layout_utils.h
new file mode 100644
index 0000000..c10fbee
--- /dev/null
+++ b/tensorflow/contrib/tpu/utils/tpu_embedding_output_layout_utils.h
@@ -0,0 +1,38 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CONTRIB_TPU_UTILS_TPU_EMBEDDING_OUTPUT_LAYOUT_UTILS_H_
+#define TENSORFLOW_CONTRIB_TPU_UTILS_TPU_EMBEDDING_OUTPUT_LAYOUT_UTILS_H_
+
+#include "tensorflow/contrib/tpu/proto/tpu_embedding_configuration.pb.h"
+#include "tensorflow/core/framework/tensor_shape.pb.h"
+#include "tensorflow/core/lib/core/status.h"
+
+namespace tensorflow {
+namespace tpu {
+
+// Creates a default output layout for compatibility if none was provided by the
+// model.
+void AddDefaultEmbeddingOutputLayoutIfNeeded(TPUEmbeddingConfiguration* config);
+
+// Computes the shape of the output tensors from an output layout.
+Status ComputeOutputTensorShapes(
+ const TPUEmbeddingConfiguration& config,
+ std::vector<tensorflow::TensorShapeProto>* shapes);
+
+} // namespace tpu
+} // namespace tensorflow
+
+#endif // TENSORFLOW_CONTRIB_TPU_UTILS_TPU_EMBEDDING_OUTPUT_LAYOUT_UTILS_H_
diff --git a/tensorflow/contrib/training/BUILD b/tensorflow/contrib/training/BUILD
index ddf8365..00295f5 100644
--- a/tensorflow/contrib/training/BUILD
+++ b/tensorflow/contrib/training/BUILD
@@ -295,7 +295,6 @@
tags = ["notsan"],
deps = [
":training_py",
- "//tensorflow/contrib/data/python/kernel_tests/serialization:dataset_serialization_test_base",
"//tensorflow/python:client_testlib",
"//tensorflow/python:framework_for_generated_wrappers",
"//tensorflow/python:gradients",
@@ -305,6 +304,7 @@
"//tensorflow/python:training",
"//tensorflow/python:variables",
"//tensorflow/python/data",
+ "//tensorflow/python/data/experimental/kernel_tests/serialization:dataset_serialization_test_base",
"//third_party/py/numpy",
],
)
@@ -313,6 +313,5 @@
name = "protos_all",
srcs = glob(["**/*.proto"]),
cc_api_version = 2,
- java_api_version = 2,
visibility = ["//visibility:public"],
)
diff --git a/tensorflow/contrib/training/python/training/device_setter_test.py b/tensorflow/contrib/training/python/training/device_setter_test.py
index 20746d9..3bb2dce 100644
--- a/tensorflow/contrib/training/python/training/device_setter_test.py
+++ b/tensorflow/contrib/training/python/training/device_setter_test.py
@@ -98,10 +98,10 @@
cluster=_CLUSTER_SPEC,
ps_strategy=device_setter_lib.GreedyLoadBalancingStrategy(
2, device_setter_lib.byte_size_load_fn))):
- u = variables.Variable(array_ops.zeros([2, 2]))
- v = variables.Variable(array_ops.zeros([2, 1]))
- w = variables.Variable(array_ops.zeros([2, 2]))
- x = variables.Variable(array_ops.zeros([1, 3]))
+ u = variables.VariableV1(array_ops.zeros([2, 2]))
+ v = variables.VariableV1(array_ops.zeros([2, 1]))
+ w = variables.VariableV1(array_ops.zeros([2, 2]))
+ x = variables.VariableV1(array_ops.zeros([1, 3]))
a = v + w
self.assertDeviceEqual("/job:ps/task:0", u.device)
self.assertDeviceEqual("/job:ps/task:0", u.initializer.device)
diff --git a/tensorflow/contrib/training/python/training/tensor_queue_dataset_test.py b/tensorflow/contrib/training/python/training/tensor_queue_dataset_test.py
index d9b0511..c1657fe 100644
--- a/tensorflow/contrib/training/python/training/tensor_queue_dataset_test.py
+++ b/tensorflow/contrib/training/python/training/tensor_queue_dataset_test.py
@@ -19,8 +19,8 @@
import numpy as np
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
from tensorflow.contrib.training.python.training import tensor_queue_dataset as tqd
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
diff --git a/tensorflow/core/BUILD b/tensorflow/core/BUILD
index bc0bfb7..0aae29d 100644
--- a/tensorflow/core/BUILD
+++ b/tensorflow/core/BUILD
@@ -149,6 +149,7 @@
"tf_cuda_tests_tags",
)
load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda")
+load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda_is_configured")
load("@io_bazel_rules_closure//closure:defs.bzl", "closure_proto_library")
load(
"//third_party/mkl:build_defs.bzl",
@@ -238,7 +239,6 @@
srcs = [],
cc_api_version = 2,
default_header = True,
- java_api_version = 2,
js_api_version = 2,
protodeps = [
":protos_all_proto",
@@ -271,6 +271,12 @@
visibility = ["//visibility:public"],
)
+java_proto_library(
+ name = "example_java_proto",
+ visibility = ["//visibility:public"],
+ deps = [":example_protos"],
+)
+
closure_proto_library(
name = "example_protos_closure",
visibility = ["//visibility:public"],
@@ -707,14 +713,11 @@
cc_library(
name = "feature_util",
srcs = ["example/feature_util.cc"],
- hdrs = [
- "example/feature_util.h",
- "platform/types.h",
- ],
+ hdrs = ["example/feature_util.h"],
visibility = ["//visibility:public"],
deps = [
":core_stringpiece",
- ":platform_protobuf",
+ ":lib_proto_parsing",
":protos_all_cc",
],
)
@@ -1041,6 +1044,7 @@
"dataset_ops",
"decode_proto_ops",
"encode_proto_ops",
+ "experimental_dataset_ops",
"function_ops",
"functional_ops",
"image_ops",
@@ -1057,7 +1061,6 @@
"random_grad",
"random_ops",
"remote_fused_graph_ops",
- "resource_variable_ops",
"rpc_ops",
"scoped_allocator_ops",
"sdca_ops",
@@ -1099,6 +1102,14 @@
deps = ["//tensorflow/core/kernels:debug_ops"],
)
+tf_gen_op_libs(
+ is_external = False,
+ op_lib_names = [
+ "resource_variable_ops",
+ ],
+ deps = [":lib"],
+)
+
# And one for all user ops
cc_library(
name = "user_ops_op_lib",
@@ -1164,6 +1175,7 @@
":dataset_ops_op_lib",
":decode_proto_ops_op_lib",
":encode_proto_ops_op_lib",
+ ":experimental_dataset_ops_op_lib",
":function_ops_op_lib",
":functional_ops_op_lib",
":image_ops_op_lib",
@@ -1363,6 +1375,7 @@
"//tensorflow/core/kernels:mkl_pooling_ops",
"//tensorflow/core/kernels:mkl_relu_op",
"//tensorflow/core/kernels:mkl_reshape_op",
+ "//tensorflow/core/kernels:mkl_slice_op",
"//tensorflow/core/kernels:mkl_softmax_op",
"//tensorflow/core/kernels:mkl_transpose_op",
"//tensorflow/core/kernels:mkl_tfconv_op",
@@ -2377,7 +2390,6 @@
srcs = ERROR_CODES_PROTO_SRCS,
cc_api_version = 2,
default_header = True,
- java_api_version = 2,
js_api_version = 2,
provide_cc_alias = True,
)
@@ -2398,7 +2410,6 @@
srcs = COMMON_PROTO_SRCS + ADDITIONAL_CORE_PROTO_SRCS,
cc_api_version = 2,
default_header = True,
- java_api_version = 2,
js_api_version = 2,
protodeps = [
":error_codes_proto",
@@ -2478,6 +2489,8 @@
"framework/op_segment.h",
"framework/rendezvous.h", # only needed for tests
"framework/resource_var.h",
+ "framework/run_handler.h",
+ "framework/run_handler_util.h",
"framework/tensor_reference.h",
"framework/tracking_allocator.h", # only needed for tests
"framework/unique_tensor_references.h",
@@ -2554,6 +2567,7 @@
"**/*test*",
"**/*main.cc",
"example/example_parser_configuration.*",
+ "example/feature_util.cc",
"util/reporter.cc",
"framework/fake_input.*",
"framework/op_gen_lib.*",
@@ -2583,6 +2597,7 @@
],
}),
deps = [
+ ":feature_util",
":lib",
":lib_internal",
":protos_all_proto_text",
@@ -2962,6 +2977,7 @@
":core_cpu_internal",
":device_tracer",
":framework",
+ ":framework_internal",
":graph",
":lib",
":lib_internal",
@@ -2999,7 +3015,7 @@
"platform/device_tracer.h",
],
copts = tf_copts(),
- cuda_deps = tf_additional_cupti_wrapper_deps() + tf_additional_device_tracer_cuda_deps(),
+ cuda_deps = if_cuda_is_configured(tf_additional_cupti_wrapper_deps() + tf_additional_device_tracer_cuda_deps()),
visibility = ["//visibility:private"],
deps = [
":core_cpu_internal",
@@ -3821,6 +3837,7 @@
"//tensorflow/core/kernels:mkl_pooling_ops",
"//tensorflow/core/kernels:mkl_relu_op",
"//tensorflow/core/kernels:mkl_reshape_op",
+ "//tensorflow/core/kernels:mkl_slice_op",
"//tensorflow/core/kernels:mkl_softmax_op",
"//tensorflow/core/kernels:mkl_tfconv_op",
]),
@@ -4108,6 +4125,19 @@
],
)
+tf_cc_test(
+ name = "framework_run_handler_util_test",
+ size = "small",
+ srcs = ["framework/run_handler_util_test.cc"],
+ linkstatic = tf_kernel_tests_linkstatic(),
+ deps = [
+ ":framework_internal",
+ ":lib",
+ ":test",
+ ":test_main",
+ ],
+)
+
tf_cuda_cc_test(
name = "common_runtime_direct_session_test",
size = "small",
diff --git a/tensorflow/core/api_def/base_api/api_def_ExperimentalAssertNextDataset.pbtxt b/tensorflow/core/api_def/base_api/api_def_ExperimentalAssertNextDataset.pbtxt
new file mode 100644
index 0000000..fa8fc96
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_ExperimentalAssertNextDataset.pbtxt
@@ -0,0 +1,4 @@
+op {
+ graph_op_name: "ExperimentalAssertNextDataset"
+ visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/base_api/api_def_ExperimentalCSVDataset.pbtxt b/tensorflow/core/api_def/base_api/api_def_ExperimentalCSVDataset.pbtxt
new file mode 100644
index 0000000..5fd88e7
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_ExperimentalCSVDataset.pbtxt
@@ -0,0 +1,4 @@
+op {
+ graph_op_name: "ExperimentalCSVDataset"
+ visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/base_api/api_def_ExperimentalDirectedInterleaveDataset.pbtxt b/tensorflow/core/api_def/base_api/api_def_ExperimentalDirectedInterleaveDataset.pbtxt
new file mode 100644
index 0000000..ac1f971
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_ExperimentalDirectedInterleaveDataset.pbtxt
@@ -0,0 +1,21 @@
+op {
+ graph_op_name: "ExperimentalDirectedInterleaveDataset"
+ in_arg {
+ name: "selector_input_dataset"
+ description: <<END
+A dataset of scalar `DT_INT64` elements that determines which of the
+`N` data inputs should produce the next output element.
+END
+ }
+ in_arg {
+ name: "data_input_datasets"
+ description: <<END
+`N` datasets with the same type that will be interleaved according to
+the values of `selector_input_dataset`.
+END
+ }
+ summary: <<END
+A substitute for `InterleaveDataset` on a fixed list of `N` datasets.
+END
+ visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/base_api/api_def_ExperimentalFunctionBufferingResource.pbtxt b/tensorflow/core/api_def/base_api/api_def_ExperimentalFunctionBufferingResource.pbtxt
new file mode 100644
index 0000000..66511ef
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_ExperimentalFunctionBufferingResource.pbtxt
@@ -0,0 +1,58 @@
+op {
+ graph_op_name: "ExperimentalFunctionBufferingResource"
+ in_arg {
+ name: "string_arg"
+ description: <<END
+String argument to the function call.
+END
+ }
+ in_arg {
+ name: "target_device"
+ description: <<END
+Target device to execute the function on.
+END
+ }
+ out_arg {
+ name: "resource"
+ description: <<END
+Handle to the resource created.
+END
+ }
+ attr {
+ name: "shared_name"
+ description: <<END
+If non-empty, this resource will be shared under the given name across
+multiple sessions.
+END
+ }
+ attr {
+ name: "container"
+ description: <<END
+If non-empty, this resource is placed in the given container.
+Otherwise, a default container is used.
+END
+ }
+ attr {
+ name: "f"
+ description: <<END
+Function to be executed.
+END
+ }
+ attr {
+ name: "buffer_size"
+ description: <<END
+Size of the buffer.
+END
+ }
+ attr {
+ name: "output_types"
+ description: <<END
+The type list for the return values.
+END
+ }
+ summary: <<END
+Creates a resource that fills up a buffer by making function calls.
+END
+ visibility: HIDDEN
+}
+
diff --git a/tensorflow/core/api_def/base_api/api_def_ExperimentalFunctionBufferingResourceGetNext.pbtxt b/tensorflow/core/api_def/base_api/api_def_ExperimentalFunctionBufferingResourceGetNext.pbtxt
new file mode 100644
index 0000000..bf4b66b
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_ExperimentalFunctionBufferingResourceGetNext.pbtxt
@@ -0,0 +1,25 @@
+op {
+ graph_op_name: "ExperimentalFunctionBufferingResourceGetNext"
+ in_arg {
+ name: "function_buffer_resource"
+ description: <<END
+The FunctionBufferingResource handle.
+END
+ }
+ out_arg {
+ name: "output"
+ description: <<END
+A list of return values.
+END
+ }
+ attr {
+ name: "output_types"
+ description: <<END
+The type list for the return values.
+END
+ }
+ summary: <<END
+Gets the next element from a FunctionBufferingResource.
+END
+ visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/base_api/api_def_ExperimentalFunctionBufferingResourceReset.pbtxt b/tensorflow/core/api_def/base_api/api_def_ExperimentalFunctionBufferingResourceReset.pbtxt
new file mode 100644
index 0000000..729718d
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_ExperimentalFunctionBufferingResourceReset.pbtxt
@@ -0,0 +1,13 @@
+op {
+ graph_op_name: "ExperimentalFunctionBufferingResourceReset"
+ in_arg {
+ name: "function_buffer_resource"
+ description: <<END
+The FunctionBufferingResource handle.
+END
+ }
+ summary: <<END
+Resets the FunctionBufferingResource.
+END
+ visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/base_api/api_def_ExperimentalIdentityIndexedDataset.pbtxt b/tensorflow/core/api_def/base_api/api_def_ExperimentalIdentityIndexedDataset.pbtxt
new file mode 100644
index 0000000..fe266c1
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_ExperimentalIdentityIndexedDataset.pbtxt
@@ -0,0 +1,4 @@
+op {
+ graph_op_name: "ExperimentalIdentityIndexedDataset"
+ visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/base_api/api_def_ExperimentalIgnoreErrorsDataset.pbtxt b/tensorflow/core/api_def/base_api/api_def_ExperimentalIgnoreErrorsDataset.pbtxt
new file mode 100644
index 0000000..d425465
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_ExperimentalIgnoreErrorsDataset.pbtxt
@@ -0,0 +1,8 @@
+op {
+ graph_op_name: "ExperimentalIgnoreErrorsDataset"
+ summary: <<END
+Creates a dataset that contains the elements of `input_dataset` ignoring errors.
+END
+ visibility: HIDDEN
+}
+
diff --git a/tensorflow/core/api_def/base_api/api_def_ExperimentalIndexedDatasetGet.pbtxt b/tensorflow/core/api_def/base_api/api_def_ExperimentalIndexedDatasetGet.pbtxt
new file mode 100644
index 0000000..e285f87
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_ExperimentalIndexedDatasetGet.pbtxt
@@ -0,0 +1,4 @@
+op {
+ graph_op_name: "ExperimentalIndexedDatasetGet"
+ visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/base_api/api_def_ExperimentalIndexedDatasetMaterialize.pbtxt b/tensorflow/core/api_def/base_api/api_def_ExperimentalIndexedDatasetMaterialize.pbtxt
new file mode 100644
index 0000000..60c3247
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_ExperimentalIndexedDatasetMaterialize.pbtxt
@@ -0,0 +1,4 @@
+op {
+ graph_op_name: "ExperimentalIndexedDatasetMaterialize"
+ visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/base_api/api_def_ExperimentalIteratorGetDevice.pbtxt b/tensorflow/core/api_def/base_api/api_def_ExperimentalIteratorGetDevice.pbtxt
new file mode 100644
index 0000000..b72b229
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_ExperimentalIteratorGetDevice.pbtxt
@@ -0,0 +1,8 @@
+op {
+ graph_op_name: "ExperimentalIteratorGetDevice"
+ summary: <<END
+Returns the name of the device on which `resource` has been placed.
+END
+ visibility: HIDDEN
+}
+
diff --git a/tensorflow/core/api_def/base_api/api_def_ExperimentalLMDBDataset.pbtxt b/tensorflow/core/api_def/base_api/api_def_ExperimentalLMDBDataset.pbtxt
new file mode 100644
index 0000000..b38b23a
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_ExperimentalLMDBDataset.pbtxt
@@ -0,0 +1,4 @@
+op {
+ graph_op_name: "ExperimentalLMDBDataset"
+ visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/base_api/api_def_ExperimentalMaterializedIndexDatasetHandle.pbtxt b/tensorflow/core/api_def/base_api/api_def_ExperimentalMaterializedIndexDatasetHandle.pbtxt
new file mode 100644
index 0000000..9676b9d
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_ExperimentalMaterializedIndexDatasetHandle.pbtxt
@@ -0,0 +1,4 @@
+op {
+ graph_op_name: "ExperimentalMaterializedIndexDatasetHandle"
+ visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/base_api/api_def_ExperimentalThreadPoolDataset.pbtxt b/tensorflow/core/api_def/base_api/api_def_ExperimentalThreadPoolDataset.pbtxt
new file mode 100644
index 0000000..d73b5bf
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_ExperimentalThreadPoolDataset.pbtxt
@@ -0,0 +1,13 @@
+op {
+ graph_op_name: "ExperimentalThreadPoolDataset"
+ in_arg {
+ name: "thread_pool"
+ description: <<END
+A resource produced by the ThreadPoolHandle op.
+END
+ }
+ summary: <<END
+Creates a dataset that uses a custom thread pool to compute `input_dataset`.
+END
+ visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/base_api/api_def_ExperimentalThreadPoolHandle.pbtxt b/tensorflow/core/api_def/base_api/api_def_ExperimentalThreadPoolHandle.pbtxt
new file mode 100644
index 0000000..48bf934
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_ExperimentalThreadPoolHandle.pbtxt
@@ -0,0 +1,35 @@
+op {
+ graph_op_name: "ExperimentalThreadPoolHandle"
+ out_arg {
+ name: "handle"
+ description: <<END
+A resource that can be consumed by one or more ExperimentalThreadPoolDataset
+ops.
+END
+ }
+ attr {
+ name: "num_threads"
+ description: <<END
+The number of threads in the thread pool.
+END
+ }
+ attr {
+ name: "max_intra_op_parallelism"
+ description: <<END
+The maximum degree of parallelism to use within operations that execute on this
+threadpool.
+END
+ }
+ attr {
+ name: "display_name"
+ description: <<END
+A human-readable name for the threads that may be visible in some
+visualizations.
+threadpool.
+END
+ }
+ summary: <<END
+Creates a dataset that uses a custom thread pool to compute `input_dataset`.
+END
+ visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/base_api/api_def_ExperimentalUniqueDataset.pbtxt b/tensorflow/core/api_def/base_api/api_def_ExperimentalUniqueDataset.pbtxt
new file mode 100644
index 0000000..68ed797
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_ExperimentalUniqueDataset.pbtxt
@@ -0,0 +1,8 @@
+op {
+ graph_op_name: "ExperimentalUniqueDataset"
+ summary: <<END
+Creates a dataset that contains the unique elements of `input_dataset`.
+END
+ visibility: HIDDEN
+}
+
diff --git a/tensorflow/core/api_def/base_api/api_def_Igamma.pbtxt b/tensorflow/core/api_def/base_api/api_def_Igamma.pbtxt
index 40d7d37..7142a0e 100644
--- a/tensorflow/core/api_def/base_api/api_def_Igamma.pbtxt
+++ b/tensorflow/core/api_def/base_api/api_def_Igamma.pbtxt
@@ -9,7 +9,7 @@
where
-\\(gamma(a, x) = int_{0}^{x} t^{a-1} exp(-t) dt\\)
+\\(gamma(a, x) = \\int_{0}^{x} t^{a-1} exp(-t) dt\\)
is the lower incomplete Gamma function.
diff --git a/tensorflow/core/api_def/base_api/api_def_ReduceDataset.pbtxt b/tensorflow/core/api_def/base_api/api_def_ReduceDataset.pbtxt
new file mode 100644
index 0000000..08414b3
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_ReduceDataset.pbtxt
@@ -0,0 +1,26 @@
+op {
+ visibility: HIDDEN
+ graph_op_name: "ReduceDataset"
+ in_arg {
+ name: "input_dataset"
+ description: <<END
+A variant tensor representing the input dataset.
+END
+ }
+ in_arg {
+ name: "initial_state"
+ description: <<END
+A nested structure of tensors, representing the initial state of the
+transformation.
+END
+ }
+ attr {
+ name: "f"
+ description: <<END
+A function that maps `(old_state, input_element)` to `new_state`. It must take
+two arguments and return a nested structures of tensors. The structure of
+`new_state` must match the structure of `initial_state`.
+END
+ }
+ summary: "Reduces the input dataset to a singleton using a reduce function."
+}
diff --git a/tensorflow/core/api_def/base_api/api_def_UnicodeScript.pbtxt b/tensorflow/core/api_def/base_api/api_def_UnicodeScript.pbtxt
new file mode 100644
index 0000000..7898fe8
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_UnicodeScript.pbtxt
@@ -0,0 +1,28 @@
+op {
+ graph_op_name: "UnicodeScript"
+ endpoint {
+ name: "UnicodeScript"
+ }
+ in_arg {
+ name: "input"
+ description: <<END
+A Tensor of int32 Unicode code points.
+END
+ }
+ out_arg {
+ name: "output"
+ description: <<END
+A Tensor of int32 script codes corresponding to each input code point.
+END
+ }
+ summary: <<END
+Determine the script codes of a given tensor of Unicode integer code points.
+END
+ description: <<END
+This operation converts Unicode code points to script codes corresponding to
+each code point. Script codes correspond to International Components for
+Unicode (ICU) UScriptCode values. See http://icu-project.org/apiref/icu4c/uscript_8h.html.
+Returns -1 (USCRIPT_INVALID_CODE) for invalid codepoints. Output shape will
+match input shape.
+END
+}
diff --git a/tensorflow/core/api_def/base_api/api_def_Xdivy.pbtxt b/tensorflow/core/api_def/base_api/api_def_Xdivy.pbtxt
new file mode 100644
index 0000000..ca107ab
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_Xdivy.pbtxt
@@ -0,0 +1,4 @@
+op {
+ graph_op_name: "Xdivy"
+ summary: "Returns 0 if x == 0, and x / y otherwise, elementwise."
+}
diff --git a/tensorflow/core/api_def/base_api/api_def_Xlogy.pbtxt b/tensorflow/core/api_def/base_api/api_def_Xlogy.pbtxt
new file mode 100644
index 0000000..da625f7
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_Xlogy.pbtxt
@@ -0,0 +1,4 @@
+op {
+ graph_op_name: "Xlogy"
+ summary: "Returns 0 if x == 0, and x * log(y) otherwise, elementwise."
+}
diff --git a/tensorflow/core/api_def/python_api/api_def_BatchToSpaceND.pbtxt b/tensorflow/core/api_def/python_api/api_def_BatchToSpaceND.pbtxt
index 9552fc9..e395e33 100644
--- a/tensorflow/core/api_def/python_api/api_def_BatchToSpaceND.pbtxt
+++ b/tensorflow/core/api_def/python_api/api_def_BatchToSpaceND.pbtxt
@@ -1,10 +1,10 @@
op {
graph_op_name: "BatchToSpaceND"
endpoint {
- name: "manip.batch_to_space_nd"
+ name: "batch_to_space_nd"
}
endpoint {
- name: "batch_to_space_nd"
+ name: "manip.batch_to_space_nd"
deprecated: true
}
}
diff --git a/tensorflow/core/api_def/python_api/api_def_GatherNd.pbtxt b/tensorflow/core/api_def/python_api/api_def_GatherNd.pbtxt
index 71257c8..598f23b 100644
--- a/tensorflow/core/api_def/python_api/api_def_GatherNd.pbtxt
+++ b/tensorflow/core/api_def/python_api/api_def_GatherNd.pbtxt
@@ -1,10 +1,10 @@
op {
graph_op_name: "GatherNd"
endpoint {
- name: "manip.gather_nd"
+ name: "gather_nd"
}
endpoint {
- name: "gather_nd"
+ name: "manip.gather_nd"
deprecated: true
}
}
diff --git a/tensorflow/core/api_def/python_api/api_def_Reshape.pbtxt b/tensorflow/core/api_def/python_api/api_def_Reshape.pbtxt
index c469665..b3d596d 100644
--- a/tensorflow/core/api_def/python_api/api_def_Reshape.pbtxt
+++ b/tensorflow/core/api_def/python_api/api_def_Reshape.pbtxt
@@ -1,10 +1,10 @@
op {
graph_op_name: "Reshape"
endpoint {
- name: "manip.reshape"
+ name: "reshape"
}
endpoint {
- name: "reshape"
+ name: "manip.reshape"
deprecated: true
}
}
diff --git a/tensorflow/core/api_def/python_api/api_def_ReverseV2.pbtxt b/tensorflow/core/api_def/python_api/api_def_ReverseV2.pbtxt
index 77f5959..51478b7 100644
--- a/tensorflow/core/api_def/python_api/api_def_ReverseV2.pbtxt
+++ b/tensorflow/core/api_def/python_api/api_def_ReverseV2.pbtxt
@@ -1,10 +1,10 @@
op {
graph_op_name: "ReverseV2"
endpoint {
- name: "manip.reverse"
+ name: "reverse"
}
endpoint {
- name: "reverse"
+ name: "manip.reverse"
deprecated: true
}
endpoint {
diff --git a/tensorflow/core/api_def/python_api/api_def_ScatterNd.pbtxt b/tensorflow/core/api_def/python_api/api_def_ScatterNd.pbtxt
index a65a19b..85888da 100644
--- a/tensorflow/core/api_def/python_api/api_def_ScatterNd.pbtxt
+++ b/tensorflow/core/api_def/python_api/api_def_ScatterNd.pbtxt
@@ -1,10 +1,10 @@
op {
graph_op_name: "ScatterNd"
endpoint {
- name: "manip.scatter_nd"
+ name: "scatter_nd"
}
endpoint {
- name: "scatter_nd"
+ name: "manip.scatter_nd"
deprecated: true
}
}
diff --git a/tensorflow/core/api_def/python_api/api_def_SpaceToBatchND.pbtxt b/tensorflow/core/api_def/python_api/api_def_SpaceToBatchND.pbtxt
index af323a6..146b97f 100644
--- a/tensorflow/core/api_def/python_api/api_def_SpaceToBatchND.pbtxt
+++ b/tensorflow/core/api_def/python_api/api_def_SpaceToBatchND.pbtxt
@@ -1,10 +1,10 @@
op {
graph_op_name: "SpaceToBatchND"
endpoint {
- name: "manip.space_to_batch_nd"
+ name: "space_to_batch_nd"
}
endpoint {
- name: "space_to_batch_nd"
+ name: "manip.space_to_batch_nd"
deprecated: true
}
}
diff --git a/tensorflow/core/api_def/python_api/api_def_StatelessMultinomial.pbtxt b/tensorflow/core/api_def/python_api/api_def_StatelessMultinomial.pbtxt
new file mode 100644
index 0000000..d3c7019
--- /dev/null
+++ b/tensorflow/core/api_def/python_api/api_def_StatelessMultinomial.pbtxt
@@ -0,0 +1,4 @@
+op {
+ graph_op_name: "StatelessMultinomial"
+ visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/python_api/api_def_StatelessRandomNormal.pbtxt b/tensorflow/core/api_def/python_api/api_def_StatelessRandomNormal.pbtxt
new file mode 100644
index 0000000..e294325
--- /dev/null
+++ b/tensorflow/core/api_def/python_api/api_def_StatelessRandomNormal.pbtxt
@@ -0,0 +1,4 @@
+op {
+ graph_op_name: "StatelessRandomNormal"
+ visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/python_api/api_def_StatelessRandomUniform.pbtxt b/tensorflow/core/api_def/python_api/api_def_StatelessRandomUniform.pbtxt
new file mode 100644
index 0000000..95d414c
--- /dev/null
+++ b/tensorflow/core/api_def/python_api/api_def_StatelessRandomUniform.pbtxt
@@ -0,0 +1,4 @@
+op {
+ graph_op_name: "StatelessRandomUniform"
+ visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/python_api/api_def_StatelessTruncatedNormal.pbtxt b/tensorflow/core/api_def/python_api/api_def_StatelessTruncatedNormal.pbtxt
new file mode 100644
index 0000000..c72bdda
--- /dev/null
+++ b/tensorflow/core/api_def/python_api/api_def_StatelessTruncatedNormal.pbtxt
@@ -0,0 +1,4 @@
+op {
+ graph_op_name: "StatelessTruncatedNormal"
+ visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/python_api/api_def_Tile.pbtxt b/tensorflow/core/api_def/python_api/api_def_Tile.pbtxt
index c34061c..1d8695f 100644
--- a/tensorflow/core/api_def/python_api/api_def_Tile.pbtxt
+++ b/tensorflow/core/api_def/python_api/api_def_Tile.pbtxt
@@ -1,10 +1,10 @@
op {
graph_op_name: "Tile"
endpoint {
- name: "manip.tile"
+ name: "tile"
}
endpoint {
- name: "tile"
+ name: "manip.tile"
deprecated: true
}
}
diff --git a/tensorflow/core/api_def/python_api/api_def_UnicodeScript.pbtxt b/tensorflow/core/api_def/python_api/api_def_UnicodeScript.pbtxt
new file mode 100644
index 0000000..a884a46
--- /dev/null
+++ b/tensorflow/core/api_def/python_api/api_def_UnicodeScript.pbtxt
@@ -0,0 +1,6 @@
+op {
+ graph_op_name: "UnicodeScript"
+ endpoint {
+ name: "strings.unicode_script"
+ }
+}
diff --git a/tensorflow/core/api_def/python_api/api_def_Xdivy.pbtxt b/tensorflow/core/api_def/python_api/api_def_Xdivy.pbtxt
new file mode 100644
index 0000000..984442b
--- /dev/null
+++ b/tensorflow/core/api_def/python_api/api_def_Xdivy.pbtxt
@@ -0,0 +1,6 @@
+op {
+ graph_op_name: "Xdivy"
+ endpoint {
+ name: "math.xdivy"
+ }
+}
diff --git a/tensorflow/core/api_def/python_api/api_def_Xlogy.pbtxt b/tensorflow/core/api_def/python_api/api_def_Xlogy.pbtxt
new file mode 100644
index 0000000..b4a5299
--- /dev/null
+++ b/tensorflow/core/api_def/python_api/api_def_Xlogy.pbtxt
@@ -0,0 +1,6 @@
+op {
+ graph_op_name: "Xlogy"
+ endpoint {
+ name: "math.xlogy"
+ }
+}
diff --git a/tensorflow/core/common_runtime/constant_folding.cc b/tensorflow/core/common_runtime/constant_folding.cc
index 99cb9ac..419867f 100644
--- a/tensorflow/core/common_runtime/constant_folding.cc
+++ b/tensorflow/core/common_runtime/constant_folding.cc
@@ -470,19 +470,19 @@
const ConstantFoldNameGenerator& generate_new_name) {
// Be conservative when replacing a tensor with a constant, when not
// running on CPU.
- // 1) If the destination tensor is not an int32 tensor, and has HOST_MEMORY
+ // 1) Do not replace another constant.
+ // 2) If the destination tensor is not an int32 tensor, and has HOST_MEMORY
// constraint, do not replace it.
- // 2) If the destination tensor is an int32 tensor, but has DEVICE_MEMORY
- // constraint, do not replace it.
- // 3) If the constant op created does not have a kernel implementation
- // for the device, do not use it.
- // 4) If the size of the constant in bytes is too large (>
+ // 3) If the size of the constant in bytes is too large (>
// max_constant_in_bytes), do not replace it. This prevents the size of the
// Graph from growing too large.
+ // 4) If the constant op created does not have a kernel implementation
+ // for the device, do not use it.
// TODO(keveman): Consider adding a new constant op that has a kernel
// implementation for all types, but with HostMemory constraint on it's
// output.
- // 5) Do not replace another constant.
+ // 5) If the constant op for the device has different output memory type
+ // from the original op output memory type, do not replace it.
if (tensor.first->IsConstant()) {
return false;
}
@@ -497,8 +497,7 @@
return false;
}
bool is_int32 = tensor.first->output_type(tensor.second) == DT_INT32;
- if ((memory_type == HOST_MEMORY && !is_int32) ||
- (memory_type == DEVICE_MEMORY && is_int32)) {
+ if (memory_type == HOST_MEMORY && !is_int32) {
return false;
}
}
@@ -536,6 +535,23 @@
if (!NodeBuilder(builder).Finalize(graph, &constant_node).ok()) {
return false;
}
+ if (partition_device && device_type != DEVICE_CPU) {
+ MemoryType original_output_memory_type;
+ if (!MemoryTypeForOutput(device_type, graph, tensor.first, tensor.second,
+ &original_output_memory_type)
+ .ok()) {
+ return false;
+ }
+ MemoryType const_output_memory_type;
+ if (!MemoryTypeForOutput(device_type, graph, constant_node, 0,
+ &const_output_memory_type)
+ .ok()) {
+ return false;
+ }
+ if (original_output_memory_type != const_output_memory_type) {
+ return false;
+ }
+ }
for (auto edge : edges_to_remove) {
graph->AddEdge(constant_node, 0, edge->dst(), edge->dst_input());
graph->RemoveEdge(edge);
diff --git a/tensorflow/core/common_runtime/copy_tensor.cc b/tensorflow/core/common_runtime/copy_tensor.cc
index d800a86..6e2eb66 100644
--- a/tensorflow/core/common_runtime/copy_tensor.cc
+++ b/tensorflow/core/common_runtime/copy_tensor.cc
@@ -61,26 +61,33 @@
status_cb->Unref();
};
auto copier = std::bind(
- [dst, recv_dev_context, out_allocator, status_cb](
- StatusCallback wrapped_done_,
- // Begin unbound arguments
- const Tensor& from, Tensor* to) {
- if (!DMAHelper::CanUseDMA(&from)) {
- Status err = errors::InvalidArgument(
- "During Variant Host->Device Copy: "
- "non-DMA-copy attempted of tensor type: ",
- DataTypeString(from.dtype()));
- status_cb->UpdateStatus(err);
- return err;
- }
- if (status_cb->ok()) {
+ [dst, recv_dev_context, out_allocator, status_cb, cpu_allocator,
+ edge_name](StatusCallback wrapped_done_,
+ // Begin unbound arguments
+ const Tensor& from, Tensor* to) {
+ if (from.dtype() == DT_VARIANT) {
status_cb->Ref();
- *to = Tensor(out_allocator, from.dtype(), from.shape());
- recv_dev_context->CopyCPUTensorToDevice(&from, dst, to,
- wrapped_done_);
+ CopyHostToDevice(&from, cpu_allocator, out_allocator, edge_name,
+ dst, to, recv_dev_context, wrapped_done_);
return Status::OK();
} else {
- return status_cb->status();
+ if (!DMAHelper::CanUseDMA(&from)) {
+ Status err = errors::InvalidArgument(
+ "During Variant Host->Device Copy: "
+ "non-DMA-copy attempted of tensor type: ",
+ DataTypeString(from.dtype()));
+ status_cb->UpdateStatus(err);
+ return err;
+ }
+ if (status_cb->ok()) {
+ status_cb->Ref();
+ *to = Tensor(out_allocator, from.dtype(), from.shape());
+ recv_dev_context->CopyCPUTensorToDevice(&from, dst, to,
+ wrapped_done_);
+ return Status::OK();
+ } else {
+ return status_cb->status();
+ }
}
},
std::move(wrapped_done), std::placeholders::_1, std::placeholders::_2);
@@ -119,26 +126,33 @@
status_cb->Unref();
};
auto copier = std::bind(
- [edge_name, src, send_dev_context, out_allocator, status_cb](
- StatusCallback wrapped_done_,
- // Begin unbound arguments
- const Tensor& from, Tensor* to) {
- if (!DMAHelper::CanUseDMA(&from)) {
- Status err = errors::InvalidArgument(
- "During Variant Device->Host Copy: "
- "non-DMA-copy attempted of tensor type: ",
- DataTypeString(from.dtype()));
- status_cb->UpdateStatus(err);
- return err;
- }
- if (status_cb->ok()) {
+ [edge_name, src, send_dev_context, out_allocator, status_cb,
+ cpu_allocator](StatusCallback wrapped_done_,
+ // Begin unbound arguments
+ const Tensor& from, Tensor* to) {
+ if (from.dtype() == DT_VARIANT) {
status_cb->Ref();
- *to = Tensor(out_allocator, from.dtype(), from.shape());
- send_dev_context->CopyDeviceTensorToCPU(&from, edge_name, src, to,
- wrapped_done_);
+ CopyDeviceToHost(&from, cpu_allocator, out_allocator, edge_name,
+ src, to, send_dev_context, wrapped_done_);
return Status::OK();
} else {
- return status_cb->status();
+ if (!DMAHelper::CanUseDMA(&from)) {
+ Status err = errors::InvalidArgument(
+ "During Variant Device->Host Copy: "
+ "non-DMA-copy attempted of tensor type: ",
+ DataTypeString(from.dtype()));
+ status_cb->UpdateStatus(err);
+ return err;
+ }
+ if (status_cb->ok()) {
+ status_cb->Ref();
+ *to = Tensor(out_allocator, from.dtype(), from.shape());
+ send_dev_context->CopyDeviceTensorToCPU(&from, edge_name, src, to,
+ wrapped_done_);
+ return Status::OK();
+ } else {
+ return status_cb->status();
+ }
}
},
std::move(wrapped_done), std::placeholders::_1, std::placeholders::_2);
diff --git a/tensorflow/core/common_runtime/direct_session.cc b/tensorflow/core/common_runtime/direct_session.cc
index af5d5b1..458e133 100644
--- a/tensorflow/core/common_runtime/direct_session.cc
+++ b/tensorflow/core/common_runtime/direct_session.cc
@@ -40,6 +40,7 @@
#include "tensorflow/core/framework/graph_def_util.h"
#include "tensorflow/core/framework/log_memory.h"
#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/run_handler.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/versions.pb.h"
#include "tensorflow/core/graph/algorithm.h"
@@ -244,6 +245,21 @@
#endif // __ANDROID__
}
+static RunHandlerPool* GetOrCreateRunHandlerPool(
+ const SessionOptions& options) {
+ static RunHandlerPool* pool =
+ new RunHandlerPool(NumInterOpThreadsFromSessionOptions(options));
+ return pool;
+}
+
+bool DirectSession::ShouldUseRunHandlerPool() const {
+ if (options_.config.session_inter_op_thread_pool_size() > 0 ||
+ options_.config.use_per_session_threads()) {
+ return false;
+ }
+ return true;
+}
+
DirectSession::DirectSession(const SessionOptions& options,
const DeviceMgr* device_mgr,
DirectSessionFactory* const factory)
@@ -363,7 +379,7 @@
Status DirectSession::Create(const GraphDef& graph) {
TF_RETURN_IF_ERROR(init_error_);
if (graph.node_size() > 0) {
- mutex_lock l(graph_def_lock_);
+ mutex_lock l(graph_state_lock_);
if (graph_created_) {
return errors::AlreadyExists(
"A Graph has already been created for this session.");
@@ -375,7 +391,7 @@
Status DirectSession::Extend(const GraphDef& graph) {
TF_RETURN_IF_ERROR(CheckNotClosed());
- mutex_lock l(graph_def_lock_);
+ mutex_lock l(graph_state_lock_);
return ExtendLocked(graph);
}
@@ -582,16 +598,37 @@
}
}
- Executor::Args::Runner default_runner = [this,
- pool](Executor::Args::Closure c) {
- SchedClosure(pool, std::move(c));
- };
+ std::unique_ptr<RunHandler> handler;
+ if (ShouldUseRunHandlerPool() &&
+ run_options.experimental().use_run_handler_pool()) {
+ // Non-null only when a global inter-op pool is used.
+ VLOG(1) << "Using RunHandler to scheduler inter-op closures.";
+ handler = GetOrCreateRunHandlerPool(options_)->Get();
+ }
+ auto* handler_ptr = handler.get();
+
+ Executor::Args::Runner default_runner = nullptr;
+
+ if (pool == nullptr) {
+ default_runner = [](Executor::Args::Closure c) { c(); };
+ } else if (handler_ptr != nullptr) {
+ default_runner = [handler_ptr](Executor::Args::Closure c) {
+ handler_ptr->ScheduleInterOpClosure(std::move(c));
+ };
+ } else {
+ default_runner = [this, pool](Executor::Args::Closure c) {
+ SchedClosure(pool, std::move(c));
+ };
+ }
+
for (const auto& item : executors_and_keys->items) {
- // TODO(zhengxq): support partial run.
- // TODO(zhengxq): if the device picks its own threadpool, we need to assign
+ // TODO(azaks): support partial run.
+ // TODO(azaks): if the device picks its own threadpool, we need to assign
// less threads to the main compute pool by default.
thread::ThreadPool* device_thread_pool =
item.device->tensorflow_device_thread_pool();
+ // TODO(crk): Investigate usage of RunHandlerPool when using device specific
+ // thread pool(s).
if (!device_thread_pool) {
args.runner = default_runner;
} else {
@@ -1172,7 +1209,7 @@
int graph_def_version;
{
- mutex_lock l(graph_def_lock_);
+ mutex_lock l(graph_state_lock_);
graph_def_version =
execution_state_->original_graph_def().versions().producer();
}
@@ -1400,7 +1437,7 @@
std::unique_ptr<FunctionLibraryDefinition>* flib_def,
RunStateArgs* run_state_args, DataTypeVector* input_types,
DataTypeVector* output_types, int64* collective_graph_key) {
- mutex_lock l(graph_def_lock_);
+ mutex_lock l(graph_state_lock_);
std::unique_ptr<ClientGraph> client_graph;
std::unique_ptr<GraphExecutionState> temp_exec_state_holder;
diff --git a/tensorflow/core/common_runtime/direct_session.h b/tensorflow/core/common_runtime/direct_session.h
index c2cf3c7..3a168bb 100644
--- a/tensorflow/core/common_runtime/direct_session.h
+++ b/tensorflow/core/common_runtime/direct_session.h
@@ -215,7 +215,7 @@
// if not already initialized.
Status MaybeInitializeExecutionState(const GraphDef& graph,
bool* out_already_initialized)
- EXCLUSIVE_LOCKS_REQUIRED(graph_def_lock_);
+ EXCLUSIVE_LOCKS_REQUIRED(graph_state_lock_);
// Retrieves an already existing set of executors to run 'inputs' and
// 'outputs', or creates and caches them for future use.
@@ -247,8 +247,11 @@
ExecutorsAndKeys* executors_and_keys,
RunMetadata* run_metadata);
+ // Returns whether inter-op execution uses a global pool.
+ bool ShouldUseRunHandlerPool() const;
+
::tensorflow::Status ExtendLocked(const GraphDef& graph)
- EXCLUSIVE_LOCKS_REQUIRED(graph_def_lock_);
+ EXCLUSIVE_LOCKS_REQUIRED(graph_state_lock_);
::tensorflow::Status ResourceHandleToInputTensor(
const Tensor& resource_tensor, Tensor* retrieved_tensor);
@@ -289,7 +292,7 @@
}
::tensorflow::Status CheckGraphCreated(const char* method) {
- mutex_lock l(graph_def_lock_);
+ mutex_lock l(graph_state_lock_);
if (!graph_created_) {
return errors::InvalidArgument(
"Session was not created with a graph before ", method, "!");
@@ -313,10 +316,8 @@
DeviceSet device_set_;
string session_handle_;
- bool graph_created_ GUARDED_BY(graph_def_lock_) = false;
-
- mutex graph_def_lock_;
- GraphDef graph_def_ GUARDED_BY(graph_def_lock_);
+ mutex graph_state_lock_;
+ bool graph_created_ GUARDED_BY(graph_state_lock_) = false;
// The thread-pools to use for running ops, with a bool indicating if the pool
// is owned.
@@ -367,11 +368,11 @@
// nodes can not be moved to a different device. Maps node names to
// device names.
std::unordered_map<string, string> stateful_placements_
- GUARDED_BY(graph_def_lock_);
+ GUARDED_BY(graph_state_lock_);
// Execution_state; used when placing the entire graph.
std::unique_ptr<GraphExecutionState> execution_state_
- GUARDED_BY(graph_def_lock_);
+ GUARDED_BY(graph_state_lock_);
// The function library, before any rewrites or optimizations have been
// performed. In particular, CreateGraphs() may need to modify the function
@@ -386,7 +387,7 @@
std::atomic<int64> edge_name_counter_ = {0};
std::atomic<int64> handle_name_counter_ = {0};
- // For generating step ids that are unique across all sessions.
+ // For generating step ids that are unique across this sessions.
static std::atomic_int_fast64_t step_id_counter_;
// Global timeout for all blocking operations in this session.
@@ -395,8 +396,6 @@
// Manages all the cost models for the graphs executed in this session.
CostModelManager cost_model_manager_;
- Executor::Args::NodeOutputsCallback node_outputs_callback_ = nullptr;
-
// For testing collective graph key generation.
mutex collective_graph_key_lock_;
int64 collective_graph_key_ GUARDED_BY(collective_graph_key_lock_) = -1;
diff --git a/tensorflow/core/common_runtime/direct_session_test.cc b/tensorflow/core/common_runtime/direct_session_test.cc
index 65e816c..e3e431f 100644
--- a/tensorflow/core/common_runtime/direct_session_test.cc
+++ b/tensorflow/core/common_runtime/direct_session_test.cc
@@ -625,6 +625,34 @@
EXPECT_EQ(run_metadata.step_stats().dev_stats_size(), 2);
}
+TEST_F(DirectSessionMinusAXTest, UseRunHandlerPool) {
+ Initialize({3, 2, -1, 0});
+ auto session = CreateSession();
+ ASSERT_TRUE(session != nullptr);
+ TF_ASSERT_OK(session->Create(def_));
+ std::vector<std::pair<string, Tensor>> inputs;
+
+ // Request two targets: one fetch output and one non-fetched output.
+ std::vector<string> output_names = {y_ + ":0"};
+ std::vector<string> target_nodes = {y_neg_};
+ std::vector<Tensor> outputs;
+
+ // Prepares RunOptions and RunMetadata
+ RunOptions run_options;
+ run_options.mutable_experimental()->set_use_run_handler_pool(true);
+
+ Status s = session->Run(run_options, inputs, output_names, target_nodes,
+ &outputs, nullptr);
+ TF_ASSERT_OK(s);
+
+ ASSERT_EQ(1, outputs.size());
+ // The first output should be initialized and have the correct
+ // output.
+ auto mat = outputs[0].matrix<float>();
+ ASSERT_TRUE(outputs[0].IsInitialized());
+ EXPECT_FLOAT_EQ(5.0, mat(0, 0));
+}
+
TEST(DirectSessionTest, KeepsStateAcrossRunsOfSession) {
GraphDef def;
Graph g(OpRegistry::Global());
diff --git a/tensorflow/core/common_runtime/direct_session_with_tracking_alloc_test.cc b/tensorflow/core/common_runtime/direct_session_with_tracking_alloc_test.cc
index efd6185..2c63b87 100644
--- a/tensorflow/core/common_runtime/direct_session_with_tracking_alloc_test.cc
+++ b/tensorflow/core/common_runtime/direct_session_with_tracking_alloc_test.cc
@@ -117,15 +117,15 @@
// which increments the value of AllocationId.
// Thus AllocationId becomes more than TF if MKL
// is used. Now IDs for MKL are 8 more than TF.
- EXPECT_EQ(29, cm->AllocationId(node, 0));
-#else
EXPECT_EQ(21, cm->AllocationId(node, 0));
+#else
+ EXPECT_EQ(13, cm->AllocationId(node, 0));
#endif // INTEL_MKL && ENABLE_MKL
} else {
#if defined(INTEL_MKL) && defined(ENABLE_MKL)
- EXPECT_EQ(30, cm->AllocationId(node, 0));
-#else
EXPECT_EQ(22, cm->AllocationId(node, 0));
+#else
+ EXPECT_EQ(14, cm->AllocationId(node, 0));
#endif // INTEL_MKL && ENABLE_MKL
}
}
diff --git a/tensorflow/core/common_runtime/eager/BUILD b/tensorflow/core/common_runtime/eager/BUILD
index be5f3ba..7b74c67 100644
--- a/tensorflow/core/common_runtime/eager/BUILD
+++ b/tensorflow/core/common_runtime/eager/BUILD
@@ -147,10 +147,11 @@
"kernel_and_device.h",
],
visibility = ["//tensorflow:internal"],
- deps = select({
+ deps = [
+ "@farmhash_archive//:farmhash",
+ ] + select({
"//tensorflow:android": [
"//tensorflow/core:android_tensorflow_lib_lite",
- "//util/hash:farmhash_fingerprint",
],
"//conditions:default": [
"//tensorflow/core:core_cpu_lib",
@@ -219,13 +220,13 @@
visibility = ["//tensorflow:internal"],
deps = [
":kernel_and_device",
+ "@farmhash_archive//:farmhash",
# Only the TF_AttrType enum is required, so pull in just the C headers.
# TODO(b/113535673): Break this dependency and avoid the C header completely.
"//tensorflow/c:c_api_headers",
] + select({
"//tensorflow:android": [
"//tensorflow/core:android_tensorflow_lib_lite",
- "//util/hash:farmhash_fingerprint",
],
"//conditions:default": [
"//tensorflow/core:core_cpu",
diff --git a/tensorflow/core/common_runtime/executor.h b/tensorflow/core/common_runtime/executor.h
index 6cd4fd2..34bf739 100644
--- a/tensorflow/core/common_runtime/executor.h
+++ b/tensorflow/core/common_runtime/executor.h
@@ -97,12 +97,6 @@
typedef std::function<void()> Closure;
typedef std::function<void(Closure)> Runner;
Runner runner = nullptr;
-
- // A callback that is invoked each time a node has finished executing.
- typedef std::function<Status(const string& node_name, const int output_slot,
- const Tensor* tensor, const bool is_ref,
- OpKernelContext* ctx)>
- NodeOutputsCallback;
};
typedef std::function<void(const Status&)> DoneCallback;
virtual void RunAsync(const Args& args, DoneCallback done) = 0;
diff --git a/tensorflow/core/common_runtime/graph_optimizer.cc b/tensorflow/core/common_runtime/graph_optimizer.cc
index 96ecfb4..37a979a 100644
--- a/tensorflow/core/common_runtime/graph_optimizer.cc
+++ b/tensorflow/core/common_runtime/graph_optimizer.cc
@@ -38,7 +38,8 @@
std::unique_ptr<Graph>* graph,
const std::unordered_map<string, std::vector<PartialTensorShape>>*
shape_map,
- const std::function<bool(const Node*)>& cse_consider_fn) {
+ const std::function<bool(const Node*)>& cse_consider_fn,
+ const std::function<bool(const Node*)>& cf_consider_fn) {
Graph* g = graph->get();
DumpGraph("Initial", g);
@@ -62,6 +63,7 @@
if (opts_.do_constant_folding()) {
ConstantFoldingOptions cf_opts;
cf_opts.shape_map = shape_map;
+ cf_opts.consider = cf_consider_fn;
if (opts_.max_folded_constant_in_bytes() > 0) {
cf_opts.max_constant_size_in_bytes =
opts_.max_folded_constant_in_bytes();
diff --git a/tensorflow/core/common_runtime/graph_optimizer.h b/tensorflow/core/common_runtime/graph_optimizer.h
index 8024628..789cc56 100644
--- a/tensorflow/core/common_runtime/graph_optimizer.h
+++ b/tensorflow/core/common_runtime/graph_optimizer.h
@@ -45,12 +45,15 @@
//
// If cse_consider_fn is not null then only nodes for which cse_consider_fn
// returns true will be considered for CSE.
+ // If cf_consider_fn is not null then only nodes for which cf_consider_fn
+ // returns true will be considered for CF.
void Optimize(
FunctionLibraryRuntime* runtime, Env* env, Device* device,
std::unique_ptr<Graph>* graph,
const std::unordered_map<string, std::vector<PartialTensorShape>>*
shape_map,
- const std::function<bool(const Node*)>& cse_consider_fn = nullptr);
+ const std::function<bool(const Node*)>& cse_consider_fn = nullptr,
+ const std::function<bool(const Node*)>& cf_consider_fn = nullptr);
const OptimizerOptions& options() { return opts_; }
diff --git a/tensorflow/core/common_runtime/lower_if_op.cc b/tensorflow/core/common_runtime/lower_if_op.cc
index dfce7c2..a02084f 100644
--- a/tensorflow/core/common_runtime/lower_if_op.cc
+++ b/tensorflow/core/common_runtime/lower_if_op.cc
@@ -38,11 +38,12 @@
public:
enum Branch { kElseBranch = 0, kThenBranch = 1 };
- // Create a CondBuilder to create the lowering of If op. that has then and
+ // Create a CondBuilder to create the lowered form of `if_op` with then and
// else functions named `then_fn_name` and `else_fn_name` respectively in the
- // given graph.
+ // `graph`. The functions should be available in `flib`.
CondBuilder(Node* if_op, const string& then_fn_name,
- const string& else_fn_name, Graph* graph);
+ const string& else_fn_name, const FunctionLibraryDefinition& flib,
+ Graph* graph);
// Constructs the basic conditional control flow using switch and merge nodes.
Status CreatePivotNodes();
@@ -89,6 +90,7 @@
Node* then_call_node_;
Node* else_call_node_;
Graph* graph_;
+ const FunctionLibraryDefinition& flib_;
string name_;
NodeBuilder then_call_builder_;
@@ -96,9 +98,11 @@
};
CondBuilder::CondBuilder(Node* if_op, const string& then_fn_name,
- const string& else_fn_name, Graph* graph)
+ const string& else_fn_name,
+ const FunctionLibraryDefinition& flib, Graph* graph)
: if_op_(if_op),
graph_(graph),
+ flib_(flib),
name_(if_op->name()),
then_call_builder_(NewName("then"), then_fn_name, graph->op_registry()),
else_call_builder_(NewName("else"), else_fn_name, graph->op_registry()) {
@@ -193,15 +197,15 @@
return Status::OK();
}
-Status InlineCallInGraph(Node* n, Graph* g) {
- const auto& lib = g->flib_def();
- const FunctionDef* fdef = lib.Find(n->type_string());
+Status InlineCallInGraph(Node* n, const FunctionLibraryDefinition& flib,
+ Graph* g) {
+ const FunctionDef* fdef = flib.Find(n->type_string());
CHECK(fdef != nullptr);
FunctionBody* fbody;
TF_RETURN_IF_ERROR(
- FunctionDefToBodyHelper(*fdef, n->attrs(), &lib,
- [&lib](const string& op, const OpDef** sig) {
- return lib.LookUpOpDef(op, sig);
+ FunctionDefToBodyHelper(*fdef, n->attrs(), &flib,
+ [&flib](const string& op, const OpDef** sig) {
+ return flib.LookUpOpDef(op, sig);
},
&fbody));
// TODO(jpienaar): Improve this interface to make the need to delete it
@@ -219,8 +223,8 @@
}
Status CondBuilder::InlineCallNodes() {
- TF_RETURN_IF_ERROR(InlineCallInGraph(then_call_node_, graph_));
- TF_RETURN_IF_ERROR(InlineCallInGraph(else_call_node_, graph_));
+ TF_RETURN_IF_ERROR(InlineCallInGraph(then_call_node_, flib_, graph_));
+ TF_RETURN_IF_ERROR(InlineCallInGraph(else_call_node_, flib_, graph_));
return Status::OK();
}
@@ -240,6 +244,12 @@
return errors::Internal("Lowering If op requires a graph to be available.");
}
+ FunctionLibraryDefinition* flib = options.flib_def;
+ if (flib == nullptr) {
+ return errors::Internal(
+ "Lowering If op requires a FunctionLibraryDefinition to be available.");
+ }
+
// Match all the nodes that need to be rewritten.
gtl::InlinedVector<Node*, 2> matches;
for (Node* n : g->op_nodes()) {
@@ -251,12 +261,14 @@
}
}
for (Node* n : matches) {
- TF_RETURN_IF_ERROR(RewriteNode(n, g));
+ TF_RETURN_IF_ERROR(RewriteNode(n, *flib, g));
}
return Status::OK();
}
-Status LowerIfOpPass::RewriteNode(Node* n, Graph* g) {
+Status LowerIfOpPass::RewriteNode(Node* n,
+ const FunctionLibraryDefinition& flib,
+ Graph* g) {
const AttrValue* then_attr = n->attrs().Find("then_branch");
if (then_attr == nullptr) {
return errors::InvalidArgument("Then branch function missing");
@@ -266,7 +278,8 @@
return errors::InvalidArgument("Else branch function missing");
}
- CondBuilder cb(n, then_attr->func().name(), else_attr->func().name(), g);
+ CondBuilder cb(n, then_attr->func().name(), else_attr->func().name(), flib,
+ g);
TF_RETURN_IF_ERROR(cb.CreatePivotNodes());
TF_RETURN_IF_ERROR(cb.AddInputs());
TF_RETURN_IF_ERROR(cb.AddOutputs());
diff --git a/tensorflow/core/common_runtime/lower_if_op.h b/tensorflow/core/common_runtime/lower_if_op.h
index a9ef39a..5ab1123 100644
--- a/tensorflow/core/common_runtime/lower_if_op.h
+++ b/tensorflow/core/common_runtime/lower_if_op.h
@@ -29,8 +29,9 @@
Status Run(const GraphOptimizationPassOptions& options) override;
private:
- // Rewrite the given If node `n` in graph `g` to use the switch-merge form.
- Status RewriteNode(Node* n, Graph* g);
+ // Rewrite the given If node `n` in graph `g` to use the switch-merge
+ // form. `flib` should contain the branch functions referenced by `n`.
+ Status RewriteNode(Node* n, const FunctionLibraryDefinition& flib, Graph* g);
};
} // namespace tensorflow
diff --git a/tensorflow/core/common_runtime/lower_if_op_test.cc b/tensorflow/core/common_runtime/lower_if_op_test.cc
index 319a617..044a355 100644
--- a/tensorflow/core/common_runtime/lower_if_op_test.cc
+++ b/tensorflow/core/common_runtime/lower_if_op_test.cc
@@ -36,9 +36,7 @@
namespace {
Status Rewrite(std::unique_ptr<Graph>* graph) {
- FunctionDefLibrary flib;
- FunctionLibraryDefinition flib_def((*graph)->op_registry(), flib);
-
+ FunctionLibraryDefinition flib_def((*graph)->flib_def());
GraphOptimizationPassOptions opt_options;
opt_options.graph = graph;
opt_options.flib_def = &flib_def;
diff --git a/tensorflow/core/common_runtime/mkl_cpu_allocator.h b/tensorflow/core/common_runtime/mkl_cpu_allocator.h
index 538a706..429b195 100644
--- a/tensorflow/core/common_runtime/mkl_cpu_allocator.h
+++ b/tensorflow/core/common_runtime/mkl_cpu_allocator.h
@@ -251,6 +251,7 @@
// max_alloc_size from large_size_allocator would be the maximum
// size allocated by MklCPUAllocator.
stats->max_alloc_size = l_stats.max_alloc_size;
+ stats->bytes_limit = std::max(s_stats.bytes_limit, l_stats.bytes_limit);
}
void ClearStats() override {
diff --git a/tensorflow/core/common_runtime/ring_reducer.cc b/tensorflow/core/common_runtime/ring_reducer.cc
index a81f865..b1fe928 100644
--- a/tensorflow/core/common_runtime/ring_reducer.cc
+++ b/tensorflow/core/common_runtime/ring_reducer.cc
@@ -41,6 +41,16 @@
// Set true for greater intelligibility of debug mode log messages.
#define READABLE_KEYS false
+// RingReduce algorithm exchanges chunks of tensor between devices. The chunk
+// size depends on the number of subdivisions specified in the algorithm. If
+// the user does not specify the number of subdivisions, we infer the number
+// dynamically so that the resulting chunk size does not exceed
+// kMaxChunkSizeBytes, empirically set at 4 MiB.
+constexpr size_t kMaxChunkSizeBytes = (4 * 1024 * 1024);
+// kMaxSubdivsPerDev is used to give an upper bound on the number of
+// subdivisions dynamically generated. A reasonable value would be a small
+// multiple of the number of NICs adjacent to each device.
+constexpr int kMaxSubdivsPerDevice = 2;
namespace tensorflow {
namespace {
@@ -92,7 +102,62 @@
RingReducer::~RingReducer() { group_size_tensor_ready_.WaitForNotification(); }
+Status GenerateSubdivsInCollectiveParams(CollectiveParams* col_params) {
+ if (col_params->instance.shape.num_elements() == 0) {
+ return errors::Internal("shape in CollectiveParams should be non-empty");
+ }
+ const int kAvgDevPerTask =
+ col_params->group.group_size / col_params->group.num_tasks;
+ const int kMaxNumSubdivs = kMaxSubdivsPerDevice * kAvgDevPerTask;
+ if (kMaxNumSubdivs <= 0) {
+ return errors::Internal("Unexpected kMaxNumSubdivs ", kMaxNumSubdivs,
+ " in RingReducer");
+ }
+ // NOTE(ayushd): If no subdiv_offsets have been specified, dynamically add
+ // as many offsets as needed so that the size of tensor chunks <=
+ // kMaxChunkSizeBytes. Empirically, chunks that are too small or too large
+ // lead to worse performance.
+ int num_subdivs = 0;
+ const size_t tensor_size = col_params->instance.shape.num_elements() *
+ DataTypeSize(col_params->instance.data_type);
+ size_t chunk_size;
+ do {
+ ++num_subdivs;
+ int num_chunks = col_params->group.group_size * num_subdivs;
+ chunk_size = tensor_size / num_chunks;
+ VLOG(2) << "num_subdivs " << num_subdivs << " num_chunks " << num_chunks
+ << " chunk_size " << chunk_size;
+ } while (chunk_size > kMaxChunkSizeBytes && num_subdivs < kMaxNumSubdivs);
+ if (num_subdivs <= 0) {
+ return errors::Internal("Unexpected num_subdivs ", num_subdivs,
+ " in RingReducer");
+ }
+
+ int subdiv_stride = kAvgDevPerTask / num_subdivs;
+ if (subdiv_stride == 0) subdiv_stride = 1;
+ col_params->instance.impl_details.subdiv_offsets.reserve(num_subdivs);
+ for (int sdi = 0; sdi < num_subdivs; ++sdi) {
+ int subdiv_offset = subdiv_stride * sdi;
+ if (sdi % 2 == 1) subdiv_offset *= -1;
+ col_params->instance.impl_details.subdiv_offsets.push_back(subdiv_offset);
+ }
+
+ if (VLOG_IS_ON(2)) {
+ string subdiv_buf;
+ for (const int subdiv_offset :
+ col_params->instance.impl_details.subdiv_offsets) {
+ strings::StrAppend(&subdiv_buf, " ", subdiv_offset);
+ }
+ VLOG(2) << "Dynamically generated " << num_subdivs
+ << " subdiv_offsets:" << subdiv_buf << " tensor_size "
+ << tensor_size << " chunk_size " << chunk_size;
+ }
+
+ return Status::OK();
+}
+
Status RingReducer::InitializeCollectiveParams(CollectiveParams* col_params) {
+ // TODO(b/113171733): change CHECKs to return errors.
CHECK_EQ(col_params->instance.type, REDUCTION_COLLECTIVE);
CHECK_EQ(col_params->instance.impl_details.collective_name, "RingReduce");
const string& device_name =
@@ -123,12 +188,11 @@
dev_per_task.push_back(dev_count);
CHECK_EQ(col_params->group.num_tasks, dev_per_task.size());
- // Generate a ring permutation for each requested offset.
if (col_params->instance.impl_details.subdiv_offsets.empty()) {
- return errors::Internal(
- "Subdiv offsets should be non-empty for ring reducer, size=",
- col_params->instance.impl_details.subdiv_offsets.size());
+ TF_RETURN_IF_ERROR(GenerateSubdivsInCollectiveParams(col_params));
}
+
+ // Generate a ring permutation for requested offset.
VLOG(2) << "Setting up perms for col_params " << col_params
<< " subdiv_permutations "
<< &col_params->instance.impl_details.subdiv_permutations;
@@ -646,7 +710,8 @@
case RF_SEND:
--send_pending_count;
break;
- default: {} // Ignore any other actions
+ default: {
+ } // Ignore any other actions
}
}
}
diff --git a/tensorflow/core/common_runtime/ring_reducer_test.cc b/tensorflow/core/common_runtime/ring_reducer_test.cc
index 28df853..75aba43 100644
--- a/tensorflow/core/common_runtime/ring_reducer_test.cc
+++ b/tensorflow/core/common_runtime/ring_reducer_test.cc
@@ -549,37 +549,38 @@
int32 reduce_counter_ GUARDED_BY(mu_) = 0;
};
-TEST_F(RingReducerTest, InitializeParams) {
- static const int kNumDevsPerTask = 8;
- static const int kNumTasks = 3;
- static const int kNumDevs = kNumDevsPerTask * kNumTasks;
+CollectiveParams SetUpCollectiveParams(const int num_devs_per_task,
+ const int num_tasks) {
CollectiveParams cp;
- std::vector<string> device_names;
- std::vector<string> task_names;
+ const int kNumDevs = num_devs_per_task * num_tasks;
cp.group.group_key = 1;
cp.group.group_size = kNumDevs;
cp.group.device_type = DeviceType("GPU");
- cp.group.num_tasks = kNumTasks;
+ cp.group.num_tasks = num_tasks;
cp.instance.instance_key = 3;
cp.instance.type = REDUCTION_COLLECTIVE;
cp.instance.data_type = DataType(DT_FLOAT);
- cp.instance.shape = TensorShape({5});
+ cp.instance.shape = TensorShape({kNumDevs});
cp.instance.impl_details.collective_name = "RingReduce";
cp.instance.impl_details.subdiv_offsets.push_back(0);
cp.is_source = false;
for (int i = 0; i < kNumDevs; ++i) {
- int task_id = i / kNumDevsPerTask;
- int dev_id = i % kNumDevsPerTask;
+ int task_id = i / num_devs_per_task;
+ int dev_id = i % num_devs_per_task;
string task_name = strings::StrCat("/job:worker/replica:0/task:", task_id);
- task_names.push_back(task_name);
string device_name = strings::StrCat(task_name, "/device:GPU:", dev_id);
- device_names.push_back(device_name);
cp.instance.task_names.push_back(task_name);
cp.instance.device_names.push_back(device_name);
}
+ return cp;
+}
- int test_rank = 0;
- cp.default_rank = test_rank;
+TEST_F(RingReducerTest, InitializeParams) {
+ const int kNumDevsPerTask = 8;
+ const int kNumTasks = 3;
+ CollectiveParams cp = SetUpCollectiveParams(kNumDevsPerTask, kNumTasks);
+
+ cp.default_rank = 0;
cp.instance.impl_details.subdiv_offsets = {0, 4};
RunSubdivPermsTest(&cp,
{{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11,
@@ -588,8 +589,15 @@
8, 9, 10, 11, 20, 21, 22, 23, 16, 17, 18, 19}},
{0, 4});
- test_rank = 3;
- cp.default_rank = test_rank;
+ cp.instance.impl_details.subdiv_offsets = {0, -4};
+ RunSubdivPermsTest(&cp,
+ {{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11,
+ 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23},
+ {3, 2, 1, 0, 7, 6, 5, 4, 11, 10, 9, 8,
+ 15, 14, 13, 12, 19, 18, 17, 16, 23, 22, 21, 20}},
+ {0, 3});
+
+ cp.default_rank = 3;
cp.instance.impl_details.subdiv_offsets = {3, -3};
RunSubdivPermsTest(&cp,
{{3, 4, 5, 6, 7, 0, 1, 2, 11, 12, 13, 14,
@@ -599,6 +607,49 @@
{0, 1});
}
+TEST_F(RingReducerTest, AutomaticSubdivs) {
+ const int kNumDevsPerTask = 8;
+ const int kNumTasks = 3;
+ const int kNumDevs = kNumDevsPerTask * kNumTasks;
+ CollectiveParams cp = SetUpCollectiveParams(kNumDevsPerTask, kNumTasks);
+
+ // Test automatic generation of subdiv offsets.
+ cp.default_rank = 0;
+ cp.instance.impl_details.subdiv_offsets.clear();
+ RunSubdivPermsTest(&cp, {{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11,
+ 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23}},
+ {0});
+
+ // Set shape so that with 2 subdivs chunk_size is 3 MiB. This should cause 2
+ // offsets, {0, -4}, to be generated.
+ {
+ int num_subdivs = 2;
+ int num_chunks = kNumDevs * num_subdivs;
+ size_t chunk_size = 3 * 1048576; // 3 MB
+ size_t tensor_size = chunk_size * num_chunks;
+ cp.instance.shape =
+ TensorShape({static_cast<int64>(tensor_size / DataTypeSize(DT_FLOAT))});
+ }
+ cp.instance.impl_details.subdiv_offsets.clear();
+ RunSubdivPermsTest(&cp,
+ {{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11,
+ 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23},
+ {3, 2, 1, 0, 7, 6, 5, 4, 11, 10, 9, 8,
+ 15, 14, 13, 12, 19, 18, 17, 16, 23, 22, 21, 20}},
+ {0, 3});
+}
+
+TEST_F(RingReducerTest, AutomaticSubdivUpperBound) {
+ const int kNumDevsPerTask = 1;
+ const int kNumTasks = 4;
+ CollectiveParams cp = SetUpCollectiveParams(kNumDevsPerTask, kNumTasks);
+
+ cp.default_rank = 0;
+ cp.instance.impl_details.subdiv_offsets.clear();
+ cp.instance.shape = TensorShape({104857600 / DataTypeSize(DT_FLOAT)});
+ RunSubdivPermsTest(&cp, {{0, 1, 2, 3}, {0, 1, 2, 3}}, {0, 0});
+}
+
// TODO(b/113171733): change to use TEST_P.
#define DEF_TEST(B, T, W, D, S, L, A) \
TEST_F(RingReducerTest, \
diff --git a/tensorflow/core/framework/common_shape_fns.cc b/tensorflow/core/framework/common_shape_fns.cc
index 20a07d8..50403b4 100644
--- a/tensorflow/core/framework/common_shape_fns.cc
+++ b/tensorflow/core/framework/common_shape_fns.cc
@@ -1306,6 +1306,113 @@
return Status::OK();
}
+namespace {
+
+// This SliceHelper processes the output shape of the `slice`
+// when the tensor of `sizes` is available.
+template <typename T>
+Status SliceHelper(InferenceContext* c, ShapeHandle begin_value,
+ const Tensor* sizes_value,
+ std::vector<DimensionHandle>* dims) {
+ auto sizes_vec = sizes_value->vec<T>();
+ for (int i = 0; i < sizes_value->NumElements(); ++i) {
+ DimensionHandle dim = c->Dim(c->input(0), i);
+ if (sizes_vec(i) != -1) {
+ auto dim_val = c->Value(dim);
+ if (sizes_vec(i) < 0) {
+ return errors::InvalidArgument(
+ "Out of bounds slicing on dimension ", i, " of length ", dim_val,
+ ": sizes vector cannot be < -1, but was ", sizes_vec(i));
+ }
+
+ dims->emplace_back(c->MakeDim(sizes_vec(i)));
+ } else {
+ DimensionHandle result;
+ TF_RETURN_IF_ERROR(c->Subtract(dim, c->Dim(begin_value, i), &result));
+ dims->emplace_back(result);
+ }
+ }
+
+ return Status::OK();
+}
+} // namespace
+
+Status SliceShape(InferenceContext* c) {
+ ShapeHandle input = c->input(0);
+ ShapeHandle begin_shape;
+ TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 1, &begin_shape));
+ ShapeHandle sizes_shape;
+ TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 1, &sizes_shape));
+
+ // Merge to check compatibility of begin and sizes tensors.
+ TF_RETURN_IF_ERROR(c->Merge(begin_shape, sizes_shape, &begin_shape));
+
+ DimensionHandle ndims = c->Dim(begin_shape, 0);
+ if (c->ValueKnown(ndims)) {
+ TF_RETURN_IF_ERROR(c->WithRank(input, c->Value(ndims), &input));
+ }
+
+ // NOTE(mrry): Use MakeShapeFromShapeTensor to handle partially-known
+ // values, even though the `begin` value does not represent a shape.
+ ShapeHandle begin_value;
+ TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensor(1, &begin_value));
+
+ // We check the tensor value here and will only use
+ // `MakeShapeFromShapeTensor` when `sizes_value` is null.
+ // The reason is that `sizes` might contain -1, which can't
+ // be represented (-1 in the ShapeHandle would mean "unknown").
+ const Tensor* sizes_value = c->input_tensor(2);
+
+ if (sizes_value != nullptr) {
+ TF_RETURN_IF_ERROR(
+ c->WithRank(begin_value, sizes_value->NumElements(), &begin_value));
+ std::vector<DimensionHandle> dims;
+ // If the begin and sizes tensors are available, then
+ // we can be precise about the shape of the output.
+ if (sizes_value->dtype() == DT_INT64) {
+ TF_RETURN_IF_ERROR(
+ SliceHelper<int64>(c, begin_value, sizes_value, &dims));
+ } else {
+ TF_RETURN_IF_ERROR(
+ SliceHelper<int32>(c, begin_value, sizes_value, &dims));
+ }
+ c->set_output(0, c->MakeShape(dims));
+ return Status::OK();
+ } else {
+ // In case `sizes` is not available (`sizes_value` is null),
+ // we could try to use `MakeShapeFromShapeTensor` here.
+ // If sizes contain -1, we will simply consider it as `Unknown`.
+ // This is less than ideal but still an improvement of shape inference.
+ // The following is an example that returns [None, 1, None] with this
+ // code path:
+ // z = tf.zeros((1, 2, 3))
+ // m = tf.slice(z, [0, 0, 0], [tf.constant(1) + 0, 1, -1])
+ // m.get_shape().as_list()
+ ShapeHandle sizes_value;
+ TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensor(2, &sizes_value));
+ if (c->RankKnown(sizes_value)) {
+ TF_RETURN_IF_ERROR(
+ c->WithRank(begin_value, c->Rank(sizes_value), &begin_value));
+ std::vector<DimensionHandle> dims;
+ dims.reserve(c->Rank(sizes_value));
+ for (int i = 0; i < c->Rank(sizes_value); ++i) {
+ dims.emplace_back(c->Dim(sizes_value, i));
+ }
+ c->set_output(0, c->MakeShape(dims));
+ return Status::OK();
+ }
+ // We might know the rank of the input.
+ if (c->RankKnown(input)) {
+ c->set_output(0, c->UnknownShapeOfRank(c->Rank(input)));
+ return Status::OK();
+ } else {
+ return shape_inference::UnknownShape(c);
+ }
+ }
+
+ return Status::OK();
+}
+
Status ValidateSparseTensor(InferenceContext* c, ShapeHandle indices_shape,
ShapeHandle values_shape, ShapeHandle shape_shape) {
// Validate ranks.
diff --git a/tensorflow/core/framework/common_shape_fns.h b/tensorflow/core/framework/common_shape_fns.h
index e6f9f93..3a496e0 100644
--- a/tensorflow/core/framework/common_shape_fns.h
+++ b/tensorflow/core/framework/common_shape_fns.h
@@ -293,6 +293,9 @@
// Shape function for random operations.
Status RandomShape(shape_inference::InferenceContext* c);
+// Shape function for Slice opertaions.
+Status SliceShape(shape_inference::InferenceContext* c);
+
// Validates the 3 component tensors of a sparse tensor have the proper
// shapes. This mimics SparseTensor.__init__ in python/framework/ops.py.
Status ValidateSparseTensor(InferenceContext* c, ShapeHandle indices_shape,
diff --git a/tensorflow/core/framework/dataset.h b/tensorflow/core/framework/dataset.h
index 697e060..8c1151c 100644
--- a/tensorflow/core/framework/dataset.h
+++ b/tensorflow/core/framework/dataset.h
@@ -657,15 +657,15 @@
// When performance modeling is enabled, this method adds a tunable parameter
// to the model node corresponding to this iterator.
//
- // The performance modeling logic may use `value` to set the value of the
+ // The performance modeling logic may use `state` to set the value of the
// tunable parameter at any point during the lifetime of this iterator. When
- // it does, it notifies `cond_var`.
+ // it does, it acquires `state->mu` and notifies `state->cond_var`.
void AddTunableParameter(IteratorContext* ctx, const string& name,
- std::atomic<int64>* value, int64 min, int64 max,
- condition_variable* cond_var) {
+ std::shared_ptr<model::SharedState> state, int64 min,
+ int64 max) {
if (ctx->model()) {
- ctx->model()->AddTunableParameter(prefix(), name, value, min, max,
- cond_var);
+ ctx->model()->AddTunableParameter(prefix(), name, std::move(state), min,
+ max);
}
}
diff --git a/tensorflow/core/framework/function.cc b/tensorflow/core/framework/function.cc
index a17959a..20f9571 100644
--- a/tensorflow/core/framework/function.cc
+++ b/tensorflow/core/framework/function.cc
@@ -1101,6 +1101,14 @@
return Status::OK();
}
+Status FunctionLibraryDefinition::ReplaceGradient(const GradientDef& grad) {
+ mutex_lock l(mu_);
+ bool added;
+ TF_RETURN_IF_ERROR(RemoveGradient(grad.function_name()));
+ TF_RETURN_IF_ERROR(AddGradientDefHelper(grad, &added));
+ return Status::OK();
+}
+
Status FunctionLibraryDefinition::RemoveFunction(const string& func) {
const auto& i = function_defs_.find(func);
if (i == function_defs_.end()) {
diff --git a/tensorflow/core/framework/function.h b/tensorflow/core/framework/function.h
index e01eb75..4d6d68e 100644
--- a/tensorflow/core/framework/function.h
+++ b/tensorflow/core/framework/function.h
@@ -331,6 +331,11 @@
// a non-OK status if "func" was not found in the library, OK otherwise.
Status ReplaceFunction(const string& func, const FunctionDef& fdef);
+ // Replaces the gradient corresponding to `grad.function_name()`. Returns
+ // a non-OK status if "grad.function_name()" was not found in the library, OK
+ // otherwise.
+ Status ReplaceGradient(const GradientDef& grad);
+
// Adds the functions and gradients in 'other' to this function library.
// Duplicate functions and gradients are ignored.
// This operation is atomic.
diff --git a/tensorflow/core/framework/function_testlib.cc b/tensorflow/core/framework/function_testlib.cc
index d5c203d..0445c24 100644
--- a/tensorflow/core/framework/function_testlib.cc
+++ b/tensorflow/core/framework/function_testlib.cc
@@ -93,7 +93,6 @@
FunctionDef RandomUniform() {
const Tensor kZero = test::AsScalar<int64>(0);
- const Tensor kTen = test::AsScalar<int64>(10);
return FDH::Define(
// Name
@@ -108,19 +107,11 @@
"Const",
{},
{{"value", kZero}, {"dtype", DT_INT64}}},
- {{"random_uniform/min"},
- "Const",
- {},
- {{"value", kZero}, {"dtype", DT_INT64}}},
- {{"random_uniform/max"},
- "Const",
- {},
- {{"value", kTen}, {"dtype", DT_INT64}}},
{{"random_uniform"},
- "RandomUniformInt",
- {},
- {{"T", DT_INT64},
- {"Tout", DT_INT64},
+ "RandomUniform",
+ {"random_uniform/shape"},
+ {{"T", DT_INT32},
+ {"Tout", DT_FLOAT},
{"seed", 87654321},
{"seed2", 42}}}});
}
diff --git a/tensorflow/core/framework/model.cc b/tensorflow/core/framework/model.cc
index b0330ec..bfdb3a6 100644
--- a/tensorflow/core/framework/model.cc
+++ b/tensorflow/core/framework/model.cc
@@ -296,12 +296,12 @@
void Model::AddTunableParameter(const string& node_name,
const string& parameter_name,
- std::atomic<int64>* value, int64 min, int64 max,
- condition_variable* cond_var) {
+ std::shared_ptr<SharedState> state, int64 min,
+ int64 max) {
tf_shared_lock l(mu_);
auto node = *gtl::FindOrNull(lookup_table_, node_name);
DCHECK(node);
- node->add_tunable_param(parameter_name, value, min, max, cond_var);
+ node->add_tunable_param(parameter_name, std::move(state), min, max);
}
// The optimization algorithm starts by setting all tunable parallelism
@@ -311,54 +311,55 @@
// is less than or equal to the processing time needed to produce an element
// divided by CPU budget.
void Model::Optimize(int64 cpu_budget) {
- tf_shared_lock lock(mu_);
std::vector<std::shared_ptr<Model::Node::Tunable>> tunables;
- const int64 processing_time = ProcessingTime();
- tunables = CollectTunables();
- for (auto tunable : tunables) {
- tunable->value = 1;
- }
- while (true) {
- const int64 output_time = OutputTime();
- bool all_tunables = true;
- for (auto& tunable : tunables) {
- if (tunable->value < tunable->max) {
- all_tunables = false;
+ {
+ tf_shared_lock lock(mu_);
+ const int64 processing_time = ProcessingTime();
+ tunables = CollectTunables();
+ for (auto tunable : tunables) {
+ tunable->value = 1;
+ }
+ while (true) {
+ const int64 output_time = OutputTime();
+ bool all_tunables = true;
+ for (auto& tunable : tunables) {
+ if (tunable->value < tunable->max) {
+ all_tunables = false;
+ break;
+ }
+ }
+ if (output_time < processing_time / cpu_budget || all_tunables) {
break;
}
- }
- if (output_time < processing_time / cpu_budget || all_tunables) {
- break;
- }
- int64 best_delta = -1;
- Model::Node::Tunable* best_tunable = nullptr;
- for (auto& tunable : tunables) {
- if (tunable->value == tunable->max) {
- continue;
+ int64 best_delta = -1;
+ Model::Node::Tunable* best_tunable = nullptr;
+ for (auto& tunable : tunables) {
+ if (tunable->value == tunable->max) {
+ continue;
+ }
+ tunable->value++;
+ int64 delta = output_time - OutputTime();
+ if (delta > best_delta) {
+ best_delta = delta;
+ best_tunable = tunable.get();
+ }
+ tunable->value--;
}
- tunable->value++;
- int64 delta = output_time - OutputTime();
- if (delta > best_delta) {
- best_delta = delta;
- best_tunable = tunable.get();
+ if (!best_tunable) {
+ // NOTE: This can happen because we are performing the optimization
+ // while the model data is changing. If this becomes an issue, we should
+ // look into performing the optimization using a model snapshot.
+ break;
}
- tunable->value--;
+ best_tunable->value++;
}
- if (!best_tunable) {
- // NOTE: This can happen because we are performing the optimization
- // while the model data is changing. If this becomes an issue, we should
- // look into performing the optimization using a model snapshot.
- break;
- }
- best_tunable->value++;
}
VLOG(2) << "Number of knobs: " << tunables.size();
for (auto& tunable : tunables) {
VLOG(2) << "Setting tunable parameter: " << tunable->value;
- tunable->value_ptr->store(tunable->value);
- if (tunable->cond_var) {
- tunable->cond_var->notify_all();
- }
+ mutex_lock l(*tunable->state->mu);
+ tunable->state->value = tunable->value;
+ tunable->state->cond_var->notify_all();
}
}
diff --git a/tensorflow/core/framework/model.h b/tensorflow/core/framework/model.h
index 26402f5..eae0fa7 100644
--- a/tensorflow/core/framework/model.h
+++ b/tensorflow/core/framework/model.h
@@ -33,6 +33,19 @@
namespace data {
namespace model {
+// Represents thread-safe state that can be shared between an input pipeline and
+// the performance model.
+struct SharedState {
+ public:
+ explicit 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)) {}
+
+ std::shared_ptr<mutex> mu;
+ std::shared_ptr<condition_variable> cond_var;
+ int64 value;
+};
+
// Abstract representation of a TensorFlow input pipeline that can be used
// for collecting runtime information and optimizing performance. It collects
// runtime information about execution of the input pipeline that is used to
@@ -62,8 +75,8 @@
// Adds a tunable parameter for the given node.
void AddTunableParameter(const string& node_name,
const string& parameter_name,
- std::atomic<int64>* value, int64 min, int64 max,
- condition_variable* cond_var) LOCKS_EXCLUDED(mu_);
+ std::shared_ptr<SharedState> value, int64 min,
+ int64 max) LOCKS_EXCLUDED(mu_);
// Runs optimization.
void Optimize(int64 cpu_budget) LOCKS_EXCLUDED(mu_);
@@ -109,13 +122,8 @@
public:
// Represents a tunable parameter.
struct Tunable {
- Tunable(std::atomic<int64>* value, int64 min, int64 max,
- condition_variable* cond_var)
- : value(*value),
- min(min),
- max(max),
- value_ptr(value),
- cond_var(cond_var) {}
+ Tunable(std::shared_ptr<SharedState> state, int64 min, int64 max)
+ : value(state->value), min(min), max(max), state(std::move(state)) {}
// Identifies the model value of the parameter. This can be different from
// the actual value (e.g. during optimization search).
@@ -127,12 +135,8 @@
// Identifies the maximum value of the parameter.
int64 max;
- // Points to the actual value of the parameter. Not owned.
- std::atomic<int64>* value_ptr;
-
- // If non-null, this condition variable is notified when the model updates
- // the actual value of the parameter (via `value_ptr`). Not owned.
- condition_variable* cond_var;
+ // Shared state of the parameter.
+ std::shared_ptr<SharedState> state;
};
Node(int64 id, const string& name, std::shared_ptr<Node> output)
@@ -158,12 +162,12 @@
}
// Adds a tunable parameter.
- void add_tunable_param(const string& name, std::atomic<int64>* value,
- int64 min, int64 max, condition_variable* cond_var)
- LOCKS_EXCLUDED(mu_) {
+ void add_tunable_param(const string& name,
+ std::shared_ptr<SharedState> state, int64 min,
+ int64 max) LOCKS_EXCLUDED(mu_) {
mutex_lock l(mu_);
tunable_params_[name] =
- std::make_shared<Tunable>(value, min, max, cond_var);
+ std::make_shared<Tunable>(std::move(state), min, max);
}
// Returns the unique node ID.
diff --git a/tensorflow/core/framework/node_def_util.h b/tensorflow/core/framework/node_def_util.h
index 187bfa2..0ff6755 100644
--- a/tensorflow/core/framework/node_def_util.h
+++ b/tensorflow/core/framework/node_def_util.h
@@ -17,7 +17,6 @@
#define TENSORFLOW_CORE_FRAMEWORK_NODE_DEF_UTIL_H_
#include <string>
-#include <unordered_map>
#include <vector>
#include "tensorflow/core/framework/attr_value_util.h"
diff --git a/tensorflow/core/framework/op.h b/tensorflow/core/framework/op.h
index 25f8de8..81ed5f9 100644
--- a/tensorflow/core/framework/op.h
+++ b/tensorflow/core/framework/op.h
@@ -209,16 +209,16 @@
class OpDefBuilderWrapper<true> {
public:
OpDefBuilderWrapper(const char name[]) : builder_(name) {}
- OpDefBuilderWrapper<true>& Attr(StringPiece spec) {
- builder_.Attr(spec);
+ OpDefBuilderWrapper<true>& Attr(string spec) {
+ builder_.Attr(std::move(spec));
return *this;
}
- OpDefBuilderWrapper<true>& Input(StringPiece spec) {
- builder_.Input(spec);
+ OpDefBuilderWrapper<true>& Input(string spec) {
+ builder_.Input(std::move(spec));
return *this;
}
- OpDefBuilderWrapper<true>& Output(StringPiece spec) {
- builder_.Output(spec);
+ OpDefBuilderWrapper<true>& Output(string spec) {
+ builder_.Output(std::move(spec));
return *this;
}
OpDefBuilderWrapper<true>& SetIsCommutative() {
@@ -237,12 +237,12 @@
builder_.SetAllowsUninitializedInput();
return *this;
}
- OpDefBuilderWrapper<true>& Deprecated(int version, StringPiece explanation) {
- builder_.Deprecated(version, explanation);
+ OpDefBuilderWrapper<true>& Deprecated(int version, string explanation) {
+ builder_.Deprecated(version, std::move(explanation));
return *this;
}
- OpDefBuilderWrapper<true>& Doc(StringPiece text) {
- builder_.Doc(text);
+ OpDefBuilderWrapper<true>& Doc(string text) {
+ builder_.Doc(std::move(text));
return *this;
}
OpDefBuilderWrapper<true>& SetShapeFn(
diff --git a/tensorflow/core/framework/op_def_builder.cc b/tensorflow/core/framework/op_def_builder.cc
index 34a7a43..8a9bb63 100644
--- a/tensorflow/core/framework/op_def_builder.cc
+++ b/tensorflow/core/framework/op_def_builder.cc
@@ -526,32 +526,32 @@
} // namespace
-OpDefBuilder::OpDefBuilder(StringPiece op_name) {
- op_def()->set_name(string(op_name)); // NOLINT
+OpDefBuilder::OpDefBuilder(string op_name) {
+ op_def()->set_name(std::move(op_name));
}
-OpDefBuilder& OpDefBuilder::Attr(StringPiece spec) {
- attrs_.emplace_back(spec.data(), spec.size());
+OpDefBuilder& OpDefBuilder::Attr(string spec) {
+ attrs_.push_back(std::move(spec));
return *this;
}
-OpDefBuilder& OpDefBuilder::Input(StringPiece spec) {
- inputs_.emplace_back(spec.data(), spec.size());
+OpDefBuilder& OpDefBuilder::Input(string spec) {
+ inputs_.push_back(std::move(spec));
return *this;
}
-OpDefBuilder& OpDefBuilder::Output(StringPiece spec) {
- outputs_.emplace_back(spec.data(), spec.size());
+OpDefBuilder& OpDefBuilder::Output(string spec) {
+ outputs_.push_back(std::move(spec));
return *this;
}
#ifndef TF_LEAN_BINARY
-OpDefBuilder& OpDefBuilder::Doc(StringPiece text) {
+OpDefBuilder& OpDefBuilder::Doc(string text) {
if (!doc_.empty()) {
errors_.push_back(
strings::StrCat("Extra call to Doc() for Op ", op_def()->name()));
} else {
- doc_.assign(text.data(), text.size());
+ doc_ = std::move(text);
}
return *this;
}
@@ -577,14 +577,14 @@
return *this;
}
-OpDefBuilder& OpDefBuilder::Deprecated(int version, StringPiece explanation) {
+OpDefBuilder& OpDefBuilder::Deprecated(int version, string explanation) {
if (op_def()->has_deprecation()) {
errors_.push_back(
strings::StrCat("Deprecated called twice for Op ", op_def()->name()));
} else {
OpDeprecation* deprecation = op_def()->mutable_deprecation();
deprecation->set_version(version);
- deprecation->set_explanation(string(explanation));
+ deprecation->set_explanation(std::move(explanation));
}
return *this;
}
diff --git a/tensorflow/core/framework/op_def_builder.h b/tensorflow/core/framework/op_def_builder.h
index 0b39d6e..8077b20 100644
--- a/tensorflow/core/framework/op_def_builder.h
+++ b/tensorflow/core/framework/op_def_builder.h
@@ -51,7 +51,7 @@
class OpDefBuilder {
public:
// Constructs an OpDef with just the name field set.
- explicit OpDefBuilder(StringPiece op_name);
+ explicit OpDefBuilder(string op_name);
// Adds an attr to this OpDefBuilder (and returns *this). The spec has
// format "<name>:<type>" or "<name>:<type>=<default>"
@@ -84,7 +84,7 @@
// * Ability to restrict the type of the tensor like the existing
// restrictions for type attrs.
// Perhaps by linking the type of the tensor to a type attr?
- OpDefBuilder& Attr(StringPiece spec);
+ OpDefBuilder& Attr(string spec);
// Adds an input or output to this OpDefBuilder (and returns *this).
// The spec has form "<name>:<type-expr>" or "<name>:Ref(<type-expr>)"
@@ -101,8 +101,8 @@
// in the spec?
// TODO(josh11b): SparseInput() and SparseOutput() matching the Python
// handling?
- OpDefBuilder& Input(StringPiece spec);
- OpDefBuilder& Output(StringPiece spec);
+ OpDefBuilder& Input(string spec);
+ OpDefBuilder& Output(string spec);
// Turns on the indicated boolean flag in this OpDefBuilder (and
// returns *this).
@@ -112,7 +112,7 @@
OpDefBuilder& SetAllowsUninitializedInput();
// Deprecate the op at a certain GraphDef version.
- OpDefBuilder& Deprecated(int version, StringPiece explanation);
+ OpDefBuilder& Deprecated(int version, string explanation);
// Adds docs to this OpDefBuilder (and returns *this).
// Docs have the format:
@@ -128,9 +128,9 @@
// to suppress the automatically-generated type documentation in
// generated output.
#ifndef TF_LEAN_BINARY
- OpDefBuilder& Doc(StringPiece text);
+ OpDefBuilder& Doc(string text);
#else
- OpDefBuilder& Doc(StringPiece text) { return *this; }
+ OpDefBuilder& Doc(string text) { return *this; }
#endif
// Sets the shape function to be used for shape inference.
diff --git a/tensorflow/core/framework/resource_mgr.cc b/tensorflow/core/framework/resource_mgr.cc
index ebdaaec..508a8d3 100644
--- a/tensorflow/core/framework/resource_mgr.cc
+++ b/tensorflow/core/framework/resource_mgr.cc
@@ -288,4 +288,13 @@
return ctx->resource_manager()->Delete(p);
}
+Status ResourceHandlesShape(shape_inference::InferenceContext* c) {
+ int n;
+ TF_RETURN_IF_ERROR(c->GetAttr("N", &n));
+ for (int i = 0; i < n; ++i) {
+ c->set_output(i, c->Scalar());
+ }
+ return Status::OK();
+}
+
} // end namespace tensorflow
diff --git a/tensorflow/core/framework/resource_mgr.h b/tensorflow/core/framework/resource_mgr.h
index d58deaa..4a53164 100644
--- a/tensorflow/core/framework/resource_mgr.h
+++ b/tensorflow/core/framework/resource_mgr.h
@@ -16,6 +16,7 @@
#ifndef TENSORFLOW_CORE_FRAMEWORK_RESOURCE_MGR_H_
#define TENSORFLOW_CORE_FRAMEWORK_RESOURCE_MGR_H_
+#include <memory>
#include <string>
#include <typeindex>
#include <typeinfo>
@@ -127,6 +128,14 @@
Status Lookup(const string& container, const string& name,
T** resource) const TF_MUST_USE_RESULT;
+ // Similar to Lookup, but looks up multiple resources at once, with only a
+ // single lock acquisition.
+ template <typename T>
+ Status LookupMany(absl::Span<std::pair<const string*, const string*> const>
+ containers_and_names,
+ std::vector<std::unique_ptr<T, core::RefCountDeleter>>*
+ resource) const TF_MUST_USE_RESULT;
+
// If "container" has a resource "name", returns it in
// "*resource". Otherwise, invokes creator() to create the resource.
// The caller takes the ownership of one ref on "*resource".
@@ -239,14 +248,31 @@
ResourceHandle* handle);
// Create a resource pointed by a given resource handle.
+//
+// If successful, the caller transfers the ownership of one ref on `resource` to
+// `ctx->resource_mgr()`.
template <typename T>
Status CreateResource(OpKernelContext* ctx, const ResourceHandle& p, T* value);
// Looks up a resource pointed by a given resource handle.
+//
+// If the lookup is successful, the caller takes the ownership of one ref on
+// `*value`, and must call its `Unref()` method when it has finished using it.
template <typename T>
Status LookupResource(OpKernelContext* ctx, const ResourceHandle& p, T** value);
+// Looks up multiple resources pointed by a sequence of resource handles.
+template <typename T>
+Status LookupResources(
+ OpKernelContext* ctx, absl::Span<ResourceHandle const> p,
+ std::vector<std::unique_ptr<T, core::RefCountDeleter>>* values);
+
// Looks up or creates a resource.
+//
+// If successful, the caller takes the ownership of one ref on `*value`, and
+// must call its `Unref()` method when it has finished using it. If the
+// `creator` is invoked, its reference on the created resource is transferred
+// to `ctx->resource_mgr()`.
template <typename T>
Status LookupOrCreateResource(OpKernelContext* ctx, const ResourceHandle& p,
T** value, std::function<Status(T**)> creator);
@@ -358,6 +384,26 @@
std::atomic<bool> initialized_{false};
};
+// Utility op kernel to produce a handle to a resource of type T.
+template <typename T>
+class ResourceHandlesOp : public OpKernel {
+ public:
+ explicit ResourceHandlesOp(OpKernelConstruction* context);
+
+ void Compute(OpKernelContext* ctx) override;
+
+ bool IsExpensive() override { return false; }
+
+ private:
+ std::vector<string> containers_;
+ std::vector<string> names_;
+ mutex mutex_;
+ std::vector<Tensor> resources_;
+ std::atomic<bool> initialized_{false};
+};
+
+Status ResourceHandlesShape(shape_inference::InferenceContext* c);
+
// Registers a kernel for an op which produces a handle to a resource of the
// specified type.
#define REGISTER_RESOURCE_HANDLE_KERNEL(Type) \
@@ -390,6 +436,24 @@
}
template <typename T>
+Status ResourceMgr::LookupMany(
+ absl::Span<std::pair<const string*, const string*> const>
+ containers_and_names,
+ std::vector<std::unique_ptr<T, core::RefCountDeleter>>* resources) const {
+ CheckDeriveFromResourceBase<T>();
+ tf_shared_lock l(mu_);
+ resources->resize(containers_and_names.size());
+ for (size_t i = 0; i < containers_and_names.size(); ++i) {
+ T* resource;
+ TF_RETURN_IF_ERROR(LookupInternal(*containers_and_names[i].first,
+ *containers_and_names[i].second,
+ &resource));
+ (*resources)[i].reset(resource);
+ }
+ return Status::OK();
+}
+
+template <typename T>
Status ResourceMgr::LookupInternal(const string& container, const string& name,
T** resource) const {
ResourceBase* found = nullptr;
@@ -499,6 +563,19 @@
}
template <typename T>
+Status LookupResources(
+ OpKernelContext* ctx, absl::Span<ResourceHandle const* const> p,
+ std::vector<std::unique_ptr<T, core::RefCountDeleter>>* values) {
+ std::vector<std::pair<const string*, const string*>> containers_and_names(
+ p.size());
+ for (size_t i = 0; i < p.size(); ++i) {
+ TF_RETURN_IF_ERROR(internal::ValidateDeviceAndType<T>(ctx, *p[i]));
+ containers_and_names[i] = {&p[i]->container(), &p[i]->name()};
+ }
+ return ctx->resource_manager()->LookupMany(containers_and_names, values);
+}
+
+template <typename T>
Status LookupOrCreateResource(OpKernelContext* ctx, const ResourceHandle& p,
T** value, std::function<Status(T**)> creator) {
TF_RETURN_IF_ERROR(internal::ValidateDeviceAndType<T>(ctx, p));
@@ -555,6 +632,46 @@
ctx->set_output(0, resource_);
}
+template <typename T>
+ResourceHandlesOp<T>::ResourceHandlesOp(OpKernelConstruction* context)
+ : OpKernel(context) {
+ int n;
+ OP_REQUIRES_OK(context, context->GetAttr("N", &n));
+ OP_REQUIRES_OK(context, context->GetAttr("containers", &containers_));
+ OP_REQUIRES_OK(context, context->GetAttr("shared_names", &names_));
+ OP_REQUIRES(
+ context, containers_.size() == n,
+ errors::InvalidArgument("Number of containers (", containers_.size(),
+ ") must be equal to N (", n, ")"));
+ OP_REQUIRES(context, names_.size() == n,
+ errors::InvalidArgument("Number of names (", containers_.size(),
+ ") must be equal to N (", n, ")"));
+ resources_.resize(n);
+}
+
+template <typename T>
+void ResourceHandlesOp<T>::Compute(OpKernelContext* ctx) {
+ if (!initialized_.load()) {
+ mutex_lock ml(mutex_);
+ // Checking again to see if another thread has initialized the resource.
+ if (!initialized_.load()) {
+ AllocatorAttributes attr;
+ attr.set_on_host(true);
+ for (size_t i = 0; i < resources_.size(); ++i) {
+ OP_REQUIRES_OK(ctx, ctx->allocate_temp(DT_RESOURCE, TensorShape({}),
+ &resources_[i], attr));
+ ResourceHandle h =
+ MakeResourceHandle<T>(ctx, containers_[i], names_[i]);
+ resources_[i].template scalar<ResourceHandle>()() = h;
+ }
+ initialized_.store(true);
+ }
+ }
+ for (size_t i = 0; i < resources_.size(); ++i) {
+ ctx->set_output(i, resources_[i]);
+ }
+}
+
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_FRAMEWORK_RESOURCE_MGR_H_
diff --git a/tensorflow/core/framework/run_handler.cc b/tensorflow/core/framework/run_handler.cc
new file mode 100644
index 0000000..0c4007e
--- /dev/null
+++ b/tensorflow/core/framework/run_handler.cc
@@ -0,0 +1,249 @@
+/* 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.
+==============================================================================*/
+
+#define EIGEN_USE_THREADS
+
+#include "tensorflow/core/framework/run_handler.h"
+
+#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
+#include "tensorflow/core/framework/run_handler_util.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/util/ptr_util.h"
+
+namespace tensorflow {
+
+// Contains the concrete implementation of the RunHandler.
+// Externally visible RunHandler class simply forwards the work to this one.
+class RunHandler::Impl {
+ public:
+ explicit Impl(RunHandlerPool::Impl* pool_impl) : pool_impl_(pool_impl) {
+ Reset();
+ }
+
+ ~Impl() {}
+
+ void set_inter_op_scheduling_range(std::uint_fast32_t start,
+ std::uint_fast32_t limit) {
+ inter_op_scheduling_range_.store(EncodePartition(start, limit),
+ std::memory_order_release);
+ }
+
+ std::uint_fast32_t inter_op_scheduling_range() const {
+ return inter_op_scheduling_range_.load(std::memory_order_acquire);
+ }
+
+ // Stores now time (in microseconds) since unix epoch when the handler is
+ // requested via RunHandlerPool::Get().
+ uint64 start_time_us() const { return start_time_us_; }
+
+ void ScheduleInterOpClosure(std::function<void()> fn);
+
+ void Reset();
+
+ RunHandlerPool::Impl* pool_impl() { return pool_impl_; }
+
+ private:
+ // Encoding/decoding logic for storing [start, limit) into a single
+ // uint_fast32_t int. We assume that pool_num_threads < (1 << 16).
+ const int kMaxPartitionBits = 16;
+ const int kMaxThreads = 1 << kMaxPartitionBits;
+
+ std::uint_fast32_t EncodePartition(std::uint_fast32_t start,
+ std::uint_fast32_t limit) {
+ return (start << kMaxPartitionBits) | limit;
+ }
+
+ void DecodePartition(std::uint_fast32_t val, std::uint_fast32_t* start,
+ std::uint_fast32_t* limit) {
+ *limit = val & (kMaxThreads - 1);
+ val >>= kMaxPartitionBits;
+ *start = val;
+ }
+
+ std::atomic_uint_fast32_t inter_op_scheduling_range_;
+ RunHandlerPool::Impl* pool_impl_; // NOT OWNED.
+ uint64 start_time_us_;
+};
+
+// Contains shared state across all run handlers present in the pool. Also
+// responsible for pool management decisions.
+// This class is thread safe.
+class RunHandlerPool::Impl {
+ public:
+ explicit Impl(int num_inter_op_threads)
+ : max_handlers_(128),
+ inter_op_thread_pool_(new thread::ThreadPool(
+ Env::Default(), ThreadOptions(), "inter_op", num_inter_op_threads)),
+ iterations_(0) {
+ VLOG(1) << "Creating a RunHandlerPool with max handlers: " << max_handlers_;
+ for (int i = 0; i < max_handlers_; ++i) {
+ handlers_.emplace_back(new RunHandler::Impl(this));
+ free_handlers_.push_back(handlers_.back().get());
+ }
+ }
+
+ ~Impl() {
+ // Sanity check that all handlers have been returned back to the pool before
+ // destruction.
+ DCHECK_EQ(handlers_.size(), max_handlers_);
+ DCHECK_EQ(free_handlers_.size(), handlers_.size());
+ DCHECK_EQ(sorted_active_handlers_.size(), 0);
+ }
+
+ thread::ThreadPool* inter_op_thread_pool() const {
+ return inter_op_thread_pool_.get();
+ }
+
+ std::unique_ptr<RunHandler> Get() LOCKS_EXCLUDED(mu_) {
+ mutex_lock l(mu_);
+ while (free_handlers_.empty()) {
+ one_handler_free_.wait(l);
+ }
+ // Remove the last entry from free_handlers_ and add to the end of
+ // sorted_active_handlers_.
+ auto* handler_impl = free_handlers_.back();
+ handler_impl->Reset();
+ // Sortedness isn't violated if we simply add at the end of the list, since
+ // handlers are expected to be obtained in increasing order of time.
+ sorted_active_handlers_.push_back(handler_impl);
+ DCHECK_LE(sorted_active_handlers_.size(), max_handlers_);
+ free_handlers_.pop_back();
+
+ RecomputePoolStatsLocked();
+ return WrapUnique<RunHandler>(new RunHandler(handler_impl));
+ }
+
+ void ReleaseHandler(RunHandler::Impl* handler) LOCKS_EXCLUDED(mu_) {
+ {
+ mutex_lock l(mu_);
+ DCHECK_GT(sorted_active_handlers_.size(), 0);
+
+ uint64 now = tensorflow::Env::Default()->NowMicros();
+ double elapsed = (now - handler->start_time_us()) / 1000.0;
+ time_hist_.Add(elapsed);
+
+ // Erase from and update sorted_active_handlers_. Add it to the end of
+ // free_handlers_.
+ auto iter = std::find(sorted_active_handlers_.begin(),
+ sorted_active_handlers_.end(), handler);
+ DCHECK(iter != sorted_active_handlers_.end())
+ << "Unexpected handler: " << handler
+ << " is being requested for release";
+
+ // Remove this handler from this list and add it to the list of free
+ // handlers.
+ sorted_active_handlers_.erase(iter);
+ free_handlers_.push_back(handler);
+ DCHECK_LE(free_handlers_.size(), max_handlers_);
+
+ RecomputePoolStatsLocked();
+ }
+ one_handler_free_.notify_one();
+ }
+
+ private:
+ void RecomputePoolStatsLocked() EXCLUSIVE_LOCKS_REQUIRED(mu_);
+
+ // Maximum number of handlers pre-created during pool construction time. The
+ // number has been chosen expecting each handler might at least want 1
+ // inter-op thread for execution (during compute intensive workloads like
+ // inference).
+ const int max_handlers_;
+
+ // Thread safe part.
+ const std::unique_ptr<thread::ThreadPool> inter_op_thread_pool_;
+
+ // Thread compatible part used only by lock under RunHandlerPool.
+ // Handlers are sorted by start time.
+ std::vector<RunHandler::Impl*> sorted_active_handlers_ GUARDED_BY(mu_);
+ std::vector<RunHandler::Impl*> free_handlers_ GUARDED_BY(mu_);
+ std::vector<std::unique_ptr<RunHandler::Impl>> handlers_ GUARDED_BY(mu_);
+ // Histogram of elapsed runtime of every handler (in ms).
+ histogram::Histogram time_hist_ GUARDED_BY(mu_);
+ std::vector<std::uint_fast32_t> inter_op_start_ GUARDED_BY(mu_);
+ std::vector<std::uint_fast32_t> inter_op_limit_ GUARDED_BY(mu_);
+ int64 iterations_ GUARDED_BY(mu_);
+ condition_variable one_handler_free_;
+ mutex mu_;
+};
+
+void RunHandlerPool::Impl::RecomputePoolStatsLocked() {
+ int num_active_requests = sorted_active_handlers_.size();
+ if (num_active_requests == 0) return;
+
+ int num_threads = inter_op_thread_pool_->NumThreads();
+
+ inter_op_start_.resize(num_active_requests);
+ inter_op_limit_.resize(num_active_requests);
+
+ const int kMinThreadsPerRequest = 3;
+ ComputeInterOpSchedulingRanges(num_active_requests, num_threads,
+ kMinThreadsPerRequest, &inter_op_start_,
+ &inter_op_limit_);
+
+ for (int i = 0; i < num_active_requests; ++i) {
+ sorted_active_handlers_[i]->set_inter_op_scheduling_range(
+ inter_op_start_[i], inter_op_limit_[i]);
+ }
+
+ if (iterations_++ % 5000 == 0 && VLOG_IS_ON(1)) {
+ VLOG(1) << "Printing time histogram: " << time_hist_.ToString();
+ VLOG(1) << "Active session runs: " << num_active_requests;
+ uint64 now = tensorflow::Env::Default()->NowMicros();
+ string ranges_str = "";
+ string times_str = "";
+ for (int i = 0; i < num_active_requests; ++i) {
+ if (i > 0) {
+ times_str += " ";
+ ranges_str += " ";
+ }
+
+ times_str += strings::StrCat(
+ (now - sorted_active_handlers_[i]->start_time_us()) / 1000.0, " ms.");
+ ranges_str += strings::StrCat("[", inter_op_start_[i], ", ",
+ inter_op_limit_[i], ")");
+ }
+ VLOG(1) << "Elapsed times are: " << times_str;
+ VLOG(1) << "Ranges are: " << ranges_str;
+ }
+}
+
+void RunHandler::Impl::ScheduleInterOpClosure(std::function<void()> fn) {
+ std::uint_fast32_t start = 0, limit = 0;
+ DecodePartition(inter_op_scheduling_range(), &start, &limit);
+ pool_impl_->inter_op_thread_pool()->Schedule(std::move(fn));
+}
+
+void RunHandler::Impl::Reset() {
+ set_inter_op_scheduling_range(
+ 0, pool_impl_->inter_op_thread_pool()->NumThreads());
+ start_time_us_ = tensorflow::Env::Default()->NowMicros();
+}
+
+RunHandlerPool::RunHandlerPool(int num_inter_op_threads)
+ : impl_(new Impl(num_inter_op_threads)) {}
+
+RunHandlerPool::~RunHandlerPool() {}
+
+std::unique_ptr<RunHandler> RunHandlerPool::Get() { return impl_->Get(); }
+
+RunHandler::RunHandler(Impl* impl) : impl_(impl) {}
+
+void RunHandler::ScheduleInterOpClosure(std::function<void()> fn) {
+ impl_->ScheduleInterOpClosure(std::move(fn));
+}
+
+RunHandler::~RunHandler() { impl_->pool_impl()->ReleaseHandler(impl_); }
+} // namespace tensorflow
diff --git a/tensorflow/core/framework/run_handler.h b/tensorflow/core/framework/run_handler.h
new file mode 100644
index 0000000..72fa630
--- /dev/null
+++ b/tensorflow/core/framework/run_handler.h
@@ -0,0 +1,95 @@
+/* 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_FRAMEWORK_RUN_HANDLER_H_
+#define TENSORFLOW_CORE_FRAMEWORK_RUN_HANDLER_H_
+
+#include "tensorflow/core/lib/core/threadpool.h"
+#include "tensorflow/core/lib/histogram/histogram.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/thread_annotations.h"
+#include "tensorflow/core/protobuf/config.pb.h"
+
+namespace tensorflow {
+
+class RunHandler;
+
+// RunHandlerPool is a fixed size pool of pre-allocated RunHandlers
+// that can be used for tracking inter-op work for a given Session::Run().
+// RunHandler(s) in the pool are initially 'inactive'. A RunHandler becomes
+// 'active' when its unique_ptr is returned by Get() and is being used by a
+// client. It becomes 'inactive' once more when its unique_ptr gets destroyed.
+//
+// Expected usage:
+//
+// * Create a single RunHandlerPool (say run_handler_pool_).
+//
+// * When a Session::Run() is invoked, obtain a handler by:
+// auto handler = run_handler_pool_->Get();
+//
+// * Use handler for scheduling all inter-op work by:
+// handler->ScheduleInterOpClosure(closure);
+//
+// This class is thread safe.
+class RunHandlerPool {
+ public:
+ explicit RunHandlerPool(int num_inter_op_threads);
+ ~RunHandlerPool();
+
+ // Returns an inactive RunHandler from the pool.
+ //
+ // RunHandlers in RunHandlerPool are initially 'inactive'.
+ // A RunHandler becomes 'active' when its unique_ptr its returned by Get()
+ // and is being used by a client. It becomes 'inactive' once more when the
+ // unique_ptr is destroyed.
+ //
+ // Will block unless there is an inactive handler.
+ std::unique_ptr<RunHandler> Get();
+
+ private:
+ class Impl;
+ friend class RunHandler;
+
+ std::unique_ptr<Impl> impl_;
+};
+
+// RunHandler can be used to schedule inter-op closures to run on a global pool
+// shared across all Session::Run(s).
+//
+// It can only be created via RunHandlerPool::Get().
+//
+// This class can be used instead of directly scheduling closures on a global
+// pool since it maintains a global view across all sessions and optimizes pool
+// scheduling to improve (median and tail) latency.
+//
+// This class is thread safe.
+class RunHandler {
+ public:
+ void ScheduleInterOpClosure(std::function<void()> fn);
+
+ ~RunHandler();
+
+ private:
+ class Impl;
+ friend class RunHandlerPool::Impl;
+
+ explicit RunHandler(Impl* impl);
+
+ Impl* impl_; // NOT OWNED.
+};
+
+} // end namespace tensorflow.
+
+#endif // TENSORFLOW_CORE_FRAMEWORK_RUN_HANDLER_H_
diff --git a/tensorflow/core/framework/run_handler_util.cc b/tensorflow/core/framework/run_handler_util.cc
new file mode 100644
index 0000000..3087998
--- /dev/null
+++ b/tensorflow/core/framework/run_handler_util.cc
@@ -0,0 +1,57 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/core/framework/run_handler_util.h"
+
+#include <algorithm>
+#include <cmath>
+#include "tensorflow/core/platform/logging.h"
+
+namespace tensorflow {
+
+void ComputeInterOpSchedulingRanges(int num_active_requests, int num_threads,
+ int min_threads_per_request,
+ std::vector<std::uint_fast32_t>* start_vec,
+ std::vector<std::uint_fast32_t>* end_vec) {
+ // Each request is expected to have weight W[i] = num_active_requests - i.
+ // Therefore, total_weight = sum of all request weights.
+ float total_weight = 0.5f * num_active_requests * (num_active_requests + 1);
+ float demand_factor = static_cast<float>(num_threads) / total_weight;
+ float last_cumulative_weight = 0.0;
+ min_threads_per_request = std::max(1, min_threads_per_request);
+ for (int i = 0; i != num_active_requests; i++) {
+ float cumulative_weight =
+ static_cast<float>(i + 1) *
+ (num_active_requests - static_cast<float>(i) * 0.5f);
+ float weight = cumulative_weight - last_cumulative_weight;
+ // Quantize thread_demand by rounding up, and also satisfying
+ // `min_threads_per_request` constraint.
+ // Note: We subtract a small epsilon (0.00001) to prevent ceil(..) from
+ // rounding weights like 4.0 to 5.
+ int demand =
+ std::max(min_threads_per_request,
+ static_cast<int>(ceil(weight * demand_factor - 0.00001f)));
+ // For the quantized range [start, end); compute the floor of real start,
+ // and expand downwards from there with length `demand` and adjust for
+ // boundary conditions.
+ int start = last_cumulative_weight * demand_factor;
+ int end = std::min(num_threads, start + demand);
+ start = std::max(0, std::min(start, end - demand));
+ start_vec->at(i) = start;
+ end_vec->at(i) = end;
+ last_cumulative_weight = cumulative_weight;
+ }
+}
+} // namespace tensorflow
diff --git a/tensorflow/core/framework/run_handler_util.h b/tensorflow/core/framework/run_handler_util.h
new file mode 100644
index 0000000..c0c36ae
--- /dev/null
+++ b/tensorflow/core/framework/run_handler_util.h
@@ -0,0 +1,43 @@
+/* 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_FRAMEWORK_RUN_HANDLER_UTIL_H_
+#define TENSORFLOW_CORE_FRAMEWORK_RUN_HANDLER_UTIL_H_
+
+#include <cstdint>
+#include <vector>
+
+namespace tensorflow {
+
+// Assign thread ranges to requests.
+// Requests are numbered 0...num_active_requests-1, and
+// threads are numbered 0...num_threads-1.
+// On return, the range start_vec->at(i)...end_vec->at(i)-1
+// indicates the subrange of the threads available to request i.
+// The ranges given to different requests may overlap.
+// Lower numbered requests will tend to be assigned more threads.
+// Thus, a client might associate older requests with lower
+// array indices so they receive access to more threads.
+// However, the routine ensures that each request is given access
+// to at least min(min_threads_per_request, num_threads) threads.
+// Every thread will be assigned to at least one request range,
+// assuming there is at least one request.
+void ComputeInterOpSchedulingRanges(int num_active_requests, int num_threads,
+ int min_threads_per_request,
+ std::vector<std::uint_fast32_t>* start_vec,
+ std::vector<std::uint_fast32_t>* end_vec);
+
+} // end namespace tensorflow
+#endif // TENSORFLOW_CORE_FRAMEWORK_RUN_HANDLER_UTIL_H_
diff --git a/tensorflow/core/framework/run_handler_util_test.cc b/tensorflow/core/framework/run_handler_util_test.cc
new file mode 100644
index 0000000..a1928c1
--- /dev/null
+++ b/tensorflow/core/framework/run_handler_util_test.cc
@@ -0,0 +1,93 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/core/framework/run_handler_util.h"
+
+#include <vector>
+#include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/test.h"
+namespace tensorflow {
+namespace {
+
+void VerifyFunction(int num_active_requests, int num_threads,
+ int min_threads_per_request, bool print_stats = false) {
+ if (print_stats) {
+ LOG(INFO) << "Test case# num_active_requests: " << num_active_requests
+ << " num_threads: " << num_threads
+ << " min_threads: " << min_threads_per_request;
+ }
+ std::vector<std::uint_fast32_t> start(num_active_requests);
+ std::vector<std::uint_fast32_t> end(num_active_requests);
+
+ ComputeInterOpSchedulingRanges(num_active_requests, num_threads,
+ min_threads_per_request, &start, &end);
+ string range_str = "";
+ for (int i = 0; i < num_active_requests; ++i) {
+ if (i > 0) range_str += " ";
+ range_str += strings::StrCat("[", start[i], ", ", end[i], ")");
+
+ ASSERT_GE(start[i], 0) << range_str;
+ ASSERT_LE(end[i], num_threads) << range_str;
+ if (i > 0) {
+ // Due to linearly decreasing demand, #threads(i - 1) >= #threads(i)
+ ASSERT_GE(end[i - 1] - start[i - 1], end[i] - start[i]) << range_str;
+ // No missing threads.
+ ASSERT_GE(end[i - 1], start[i]) << range_str;
+ }
+ // Each interval is at least of size 'min_threads_per_request'.
+ ASSERT_GE((end[i] - start[i]), min_threads_per_request) << range_str;
+ // Verify that assigned (quantized) threads is not overly estimated
+ // from real demand, when the demand is high (>=
+ // min_threads_per_request).
+ float entry_weight = num_active_requests - i;
+ float total_weight = 0.5f * num_active_requests * (num_active_requests + 1);
+ float thread_demand = (entry_weight * num_threads) / total_weight;
+ if (thread_demand > min_threads_per_request) {
+ // We expect some over-estimation of threads due to quantization,
+ // but we hope it's not more than 1 extra thread.
+ ASSERT_NEAR(end[i] - start[i], thread_demand, 1.0)
+ << "Ranges: " << range_str << " thread_demand: " << thread_demand
+ << " i: " << i;
+ }
+ }
+ ASSERT_EQ(end[num_active_requests - 1], num_threads);
+ ASSERT_EQ(start[0], 0);
+ if (print_stats) {
+ LOG(INFO) << "Assigned ranges: " << range_str;
+ }
+}
+
+TEST(RunHandlerUtilTest, TestComputeInterOpSchedulingRanges) {
+ const int kMinThreadsPerRequestBound = 12;
+ const int kMaxActiveRequests = 128;
+ const int kMaxThreads = 128;
+
+ for (int min_threads_per_request = 1;
+ min_threads_per_request <= kMinThreadsPerRequestBound;
+ ++min_threads_per_request) {
+ for (int num_active_requests = 1; num_active_requests <= kMaxActiveRequests;
+ ++num_active_requests) {
+ for (int num_threads = min_threads_per_request;
+ num_threads <= kMaxThreads; ++num_threads) {
+ VerifyFunction(num_active_requests, num_threads,
+ min_threads_per_request);
+ }
+ }
+ }
+}
+
+} // namespace
+} // namespace tensorflow
diff --git a/tensorflow/core/graph/mkl_layout_pass.cc b/tensorflow/core/graph/mkl_layout_pass.cc
index 37b88f1..06d3fef 100644
--- a/tensorflow/core/graph/mkl_layout_pass.cc
+++ b/tensorflow/core/graph/mkl_layout_pass.cc
@@ -2450,6 +2450,7 @@
csinfo_.tanh = "Tanh";
csinfo_.tanh_grad = "TanhGrad";
csinfo_.reshape = "Reshape";
+ csinfo_.slice = "Slice";
csinfo_.softmax = "Softmax";
csinfo_.split = "Split";
// Element-wise ops. Ensure you also add any new ops to IsOpElementWise
@@ -2557,6 +2558,9 @@
rinfo_.push_back({csinfo_.reshape,
mkl_op_registry::GetMklOpName(csinfo_.reshape),
CopyAttrsReshape, AlwaysRewrite});
+ rinfo_.push_back({csinfo_.slice,
+ mkl_op_registry::GetMklOpName(csinfo_.slice),
+ CopyAttrsSlice, AlwaysRewrite});
rinfo_.push_back({csinfo_.softmax,
mkl_op_registry::GetMklOpName(csinfo_.softmax),
CopyAttrsDataType, AlwaysRewrite});
@@ -2676,6 +2680,7 @@
string tanh;
string tanh_grad;
string reshape;
+ string slice;
string softmax;
string split;
string squared_difference;
@@ -3134,6 +3139,7 @@
static void CopyAttrsLRN(const Node* orig_node, NodeBuilder* nb);
static void CopyAttrsPooling(const Node* orig_node, NodeBuilder* nb);
static void CopyAttrsReshape(const Node* orig_node, NodeBuilder* nb);
+ static void CopyAttrsSlice(const Node* orig_node, NodeBuilder* nb);
static void CopyAttrsSplit(const Node* orig_node, NodeBuilder* nb);
// Generate a graph node in graph 'g' representing a dummy Mkl tensor node,
@@ -3739,6 +3745,19 @@
nb->Attr("Tshape", Tshape);
}
+void MklLayoutRewritePass::CopyAttrsSlice(const Node* orig_node,
+ NodeBuilder* nb) {
+ DataType T;
+ DataType Index;
+
+ // Get all attributes from old node.
+ TF_CHECK_OK(GetNodeAttr(orig_node->def(), "T", &T));
+ TF_CHECK_OK(GetNodeAttr(orig_node->def(), "Index", &Index));
+ // Add attributes to new node.
+ nb->Attr("T", T);
+ nb->Attr("Index", Index);
+}
+
void MklLayoutRewritePass::CopyAttrsSplit(const Node* orig_node,
NodeBuilder* nb) {
DataType T;
diff --git a/tensorflow/core/graph/mkl_layout_pass_test.cc b/tensorflow/core/graph/mkl_layout_pass_test.cc
index f42a4ee..77640e2 100644
--- a/tensorflow/core/graph/mkl_layout_pass_test.cc
+++ b/tensorflow/core/graph/mkl_layout_pass_test.cc
@@ -3510,6 +3510,26 @@
"B->C:1;C->E;D->E:1;E->Z;M->C:2;N->C:3;Y->Z:1");
}
+TEST_F(MklLayoutPassTest, NodeRewrite_Slice_DeviceTest) {
+ InitGraph(
+ "node { name: 'A' op: 'Input'}"
+ "node { name: 'B' op: 'Int32Input'}"
+ "node { name: 'C' op: 'Int32Input'}"
+ "node { name: 'D' op: 'Slice'"
+ " attr { key: 'T' value { type: DT_FLOAT } }"
+ " attr { key: 'Index' value { type: DT_INT32 } }"
+ " input: ['A', 'B', 'C'] }"
+ "node { name: 'E' op: 'Zeta' attr { key: 'T' value { type: DT_FLOAT } }"
+ " input: ['A', 'D'] }");
+ EXPECT_EQ(DoMklLayoutOptimizationPass(),
+ "A(Input);B(Int32Input);C(Int32Input);"
+ "D(_MklSlice);DMT/_0(Const);DMT/_1(Const);DMT/"
+ "_2(Const);E(Zeta)|A->D;A->E;"
+ "A:control->DMT/_0:control;A:control->DMT/"
+ "_1:control;A:control->DMT/_2:control;"
+ "B->D:1;C->D:2;D->E:1;DMT/_0->D:3;DMT/_1->D:4;DMT/_2->D:5");
+}
+
/////////////////////////////////////////////////////////////////////
// Post-rewrite fixup pass test
diff --git a/tensorflow/core/graph/node_builder.cc b/tensorflow/core/graph/node_builder.cc
index a446e0d..d928749 100644
--- a/tensorflow/core/graph/node_builder.cc
+++ b/tensorflow/core/graph/node_builder.cc
@@ -99,6 +99,11 @@
return *this;
}
+NodeBuilder& NodeBuilder::AssignedDevice(StringPiece device) {
+ assigned_device_ = string(device);
+ return *this;
+}
+
Status NodeBuilder::Finalize(Graph* graph, Node** created_node) const {
// In case of error, set *created_node to nullptr.
if (created_node != nullptr) *created_node = nullptr;
@@ -115,6 +120,8 @@
Node* node = graph->AddNode(node_def, &status);
if (!status.ok()) return status;
+ node->set_assigned_device_name(assigned_device_);
+
for (size_t i = 0; i < inputs_.size(); ++i) {
if (inputs_[i].node != nullptr) { // Skip back edges.
graph->AddEdge(inputs_[i].node, inputs_[i].index, node, i);
diff --git a/tensorflow/core/graph/node_builder.h b/tensorflow/core/graph/node_builder.h
index 4727ee7..d576985 100644
--- a/tensorflow/core/graph/node_builder.h
+++ b/tensorflow/core/graph/node_builder.h
@@ -100,6 +100,9 @@
// "assigned device" in the Node).
NodeBuilder& Device(StringPiece device_spec);
+ // Sets the device name in the "assigned device" field in tensorflow::Node.
+ NodeBuilder& AssignedDevice(StringPiece device);
+
// Set the value of an attr. attr_name must match the name of one of
// attrs defined by the Op, and value must have the corresponding type
// (see SetAttrValue() in ../framework/attr_value_util.h for legal
@@ -141,6 +144,7 @@
std::vector<NodeOut> inputs_;
std::vector<Node*> control_inputs_;
std::vector<string> errors_;
+ string assigned_device_;
};
// IMPLEMENTATION -------------------------------------------------------------
diff --git a/tensorflow/core/grappler/graph_view.cc b/tensorflow/core/grappler/graph_view.cc
index 2619a9a..0b8cb5e 100644
--- a/tensorflow/core/grappler/graph_view.cc
+++ b/tensorflow/core/grappler/graph_view.cc
@@ -72,7 +72,7 @@
void GraphView::AddFanouts(NodeDef* node) {
for (int i = 0; i < node->input_size(); ++i) {
OutputPort fanin;
- string fanin_name = ParseNodeName(node->input(i), &fanin.port_id);
+ const string fanin_name = ParseNodeName(node->input(i), &fanin.port_id);
fanin.node = nodes_[fanin_name];
InputPort input;
diff --git a/tensorflow/core/grappler/grappler_item.cc b/tensorflow/core/grappler/grappler_item.cc
index bbc0fed..2c490f3 100644
--- a/tensorflow/core/grappler/grappler_item.cc
+++ b/tensorflow/core/grappler/grappler_item.cc
@@ -38,6 +38,7 @@
restore_op = other.restore_op;
save_restore_loc_tensor = other.save_restore_loc_tensor;
queue_runners = other.queue_runners;
+ allowed_optimizations = other.allowed_optimizations;
graph.Swap(graph_def);
}
diff --git a/tensorflow/core/grappler/grappler_item.h b/tensorflow/core/grappler/grappler_item.h
index 939e5fa..a0748ab 100644
--- a/tensorflow/core/grappler/grappler_item.h
+++ b/tensorflow/core/grappler/grappler_item.h
@@ -77,6 +77,15 @@
// Return a set of node names that must be preserved. This includes feed and
// fetch nodes, keep_ops, init_ops.
std::unordered_set<string> NodesToPreserve() const;
+
+ // Restrict types of optimizations that are allowed for this GrapplerItem.
+ struct AllowedOptimizations {
+ // Is it allowed to add nodes to the graph that do not have registered
+ // gradient function.
+ bool non_differentiable_rewrites = true;
+ };
+
+ AllowedOptimizations allowed_optimizations;
};
// Return the transitive fanin of a set of terminal nodes.
diff --git a/tensorflow/core/grappler/grappler_item_builder.cc b/tensorflow/core/grappler/grappler_item_builder.cc
index 029515a..3690466 100644
--- a/tensorflow/core/grappler/grappler_item_builder.cc
+++ b/tensorflow/core/grappler/grappler_item_builder.cc
@@ -192,9 +192,13 @@
const string feed_name = NodeName(feed_node);
new_item->feed.emplace_back(feed_name, Tensor());
}
+ for (const auto& fetch_node : cfg.fetch_nodes) {
+ new_item->fetch.emplace_back(NodeName(fetch_node));
+ }
- // Attempt to detect the fetch node(s).
- if (meta_graph.collection_def().count("train_op") > 0) {
+ // Attempt to detect the fetch node(s) if they were not set explicitly.
+ if (new_item->fetch.empty() &&
+ meta_graph.collection_def().count("train_op") > 0) {
const CollectionDef& nodes = meta_graph.collection_def().at("train_op");
if (nodes.has_node_list()) {
for (const auto& node : nodes.node_list().value()) {
diff --git a/tensorflow/core/grappler/grappler_item_builder.h b/tensorflow/core/grappler/grappler_item_builder.h
index aafd2fd..1698587 100644
--- a/tensorflow/core/grappler/grappler_item_builder.h
+++ b/tensorflow/core/grappler/grappler_item_builder.h
@@ -49,6 +49,8 @@
bool prune_graph = false;
// Override feed nodes list.
std::set<string> feed_nodes;
+ // Override fetch nodes list.
+ std::set<string> fetch_nodes;
};
// Factory method for creating a GrapplerItem from a MetaGraphDef.
diff --git a/tensorflow/core/grappler/grappler_item_builder_test.cc b/tensorflow/core/grappler/grappler_item_builder_test.cc
index 4b90bf3..d00981f 100644
--- a/tensorflow/core/grappler/grappler_item_builder_test.cc
+++ b/tensorflow/core/grappler/grappler_item_builder_test.cc
@@ -313,6 +313,29 @@
EXPECT_EQ(item2->feed[0].second.NumElements(), 1);
}
+TEST_F(GrapplerItemBuilderTest, ExplicitFeedAndFetch) {
+ tensorflow::Scope s = tensorflow::Scope::NewRootScope();
+ auto x = ops::Const(s.WithOpName("x"), 0);
+ auto y = ops::Const(s.WithOpName("y"), 1);
+ auto z = ops::Add(s.WithOpName("z"), x, y);
+
+ MetaGraphDef meta_graph;
+ TF_CHECK_OK(s.ToGraphDef(meta_graph.mutable_graph_def()));
+
+ ItemConfig config;
+ config.feed_nodes.insert("x");
+ config.fetch_nodes.insert("z");
+
+ std::unique_ptr<GrapplerItem> item =
+ GrapplerItemFromMetaGraphDef("0", meta_graph, config);
+ ASSERT_TRUE(item != nullptr);
+
+ EXPECT_EQ(item->feed.size(), 1);
+ EXPECT_EQ(item->fetch.size(), 1);
+ EXPECT_EQ(item->feed[0].first, "x");
+ EXPECT_EQ(item->fetch[0], "z");
+}
+
} // namespace
} // namespace grappler
} // namespace tensorflow
diff --git a/tensorflow/core/grappler/op_types.cc b/tensorflow/core/grappler/op_types.cc
index 3521669..9f0d9db 100644
--- a/tensorflow/core/grappler/op_types.cc
+++ b/tensorflow/core/grappler/op_types.cc
@@ -425,6 +425,10 @@
return op == "Switch" || op == "RefSwitch";
}
+bool IsSymbolicGradient(const NodeDef& node) {
+ return node.op() == "SymbolicGradient";
+}
+
bool IsTanhGrad(const NodeDef& node) { return node.op() == "TanhGrad"; }
bool IsTile(const NodeDef& node) { return node.op() == "Tile"; }
diff --git a/tensorflow/core/grappler/op_types.h b/tensorflow/core/grappler/op_types.h
index 25ab6b6..7f86a5f 100644
--- a/tensorflow/core/grappler/op_types.h
+++ b/tensorflow/core/grappler/op_types.h
@@ -149,6 +149,7 @@
bool IsSub(const NodeDef& node);
bool IsSum(const NodeDef& node);
bool IsSwitch(const NodeDef& node);
+bool IsSymbolicGradient(const NodeDef& node);
bool IsTanhGrad(const NodeDef& node);
bool IsTile(const NodeDef& node);
bool IsTranspose(const NodeDef& node);
diff --git a/tensorflow/core/grappler/optimizers/BUILD b/tensorflow/core/grappler/optimizers/BUILD
index 960d1ad..c708f84 100644
--- a/tensorflow/core/grappler/optimizers/BUILD
+++ b/tensorflow/core/grappler/optimizers/BUILD
@@ -525,6 +525,7 @@
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
+ "//tensorflow/core:lib_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/grappler:grappler_item",
"//tensorflow/core/grappler/utils:colocation",
@@ -541,6 +542,7 @@
":custom_graph_optimizer_registry",
":meta_optimizer",
"//tensorflow/cc:cc_ops",
+ "//tensorflow/core:lib_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:tensorflow",
"//tensorflow/core:test",
diff --git a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
index 75ed126..7d5014e 100644
--- a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
@@ -276,7 +276,7 @@
for (const NodeDef* output : ctx().node_map->GetOutputs(node.name())) {
for (int i = 0; i < output->input_size(); ++i) {
auto input = output->input(i);
- string name = ParseNodeName(input, &position);
+ StringPiece name = ParseNodeNameAsStringPiece(input, &position);
if (name == node.name() && /*control input*/ position < 0) {
return true;
}
@@ -1568,7 +1568,8 @@
for (NodeDef* output : outputs) {
if (IsControlInput(output->input(0))) continue;
int port;
- const string node_name = ParseNodeName(output->input(0), &port);
+ const StringPiece node_name =
+ ParseNodeNameAsStringPiece(output->input(0), &port);
if (node_name == node.name()) {
tails->insert(ChainLink(output, port));
} else {
@@ -1618,7 +1619,8 @@
} else {
for (NodeDef* new_tail : ctx().node_map->GetOutputs(tail->name())) {
int port;
- const string node_name = ParseNodeName(new_tail->input(0), &port);
+ const StringPiece node_name =
+ ParseNodeNameAsStringPiece(new_tail->input(0), &port);
if (node_name != tail->name()) {
return Status::OK();
}
@@ -2929,8 +2931,8 @@
for (const auto& input : node.input()) {
int pos;
- string node_name = ParseNodeName(input, &pos);
- h = Hash64CombineUnordered(Hash64(node_name), h);
+ const StringPiece node_name = ParseNodeNameAsStringPiece(input, &pos);
+ h = Hash64CombineUnordered(Hash64(node_name.data(), node_name.size()), h);
h = Hash64CombineUnordered(std::hash<int>()(pos), h);
}
for (const auto& attr : node.attr()) {
@@ -3247,6 +3249,10 @@
optimized_graph_ = &optimized_item.graph;
node_map_.reset(new NodeMap(optimized_graph_));
+ // Disable restricted graph rewrites.
+ options_.unary_ops_composition &=
+ item.allowed_optimizations.non_differentiable_rewrites;
+
if (options_.dedup_computations) {
DedupComputations();
}
diff --git a/tensorflow/core/grappler/optimizers/data/BUILD b/tensorflow/core/grappler/optimizers/data/BUILD
index cf305ce..5a3abbb 100644
--- a/tensorflow/core/grappler/optimizers/data/BUILD
+++ b/tensorflow/core/grappler/optimizers/data/BUILD
@@ -22,6 +22,7 @@
"//tensorflow/core/grappler/utils:topological_sort",
"//tensorflow/core/grappler/optimizers:custom_graph_optimizer",
"//tensorflow/core/grappler/optimizers:custom_graph_optimizer_registry",
+ "//tensorflow/core:lib_internal",
] + tf_protos_all(),
)
@@ -31,6 +32,7 @@
visibility = ["//visibility:public"],
deps = [
":filter_fusion",
+ ":graph_test_utils",
":graph_utils",
"//tensorflow/core:framework",
"//tensorflow/core:test",
@@ -87,11 +89,12 @@
],
visibility = ["//visibility:public"],
deps = [
+ ":graph_utils",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core/grappler:mutable_graph_view",
- "//tensorflow/core/grappler:grappler_item",
"//tensorflow/core/grappler:utils",
+ "//tensorflow/core:lib_internal",
] + tf_protos_all(),
)
@@ -121,10 +124,10 @@
],
visibility = ["//visibility:public"],
deps = [
+ "//tensorflow/core:core_cpu",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core/grappler:mutable_graph_view",
- "//tensorflow/core/grappler:grappler_item",
"//tensorflow/core/grappler:utils",
] + tf_protos_all(),
)
@@ -146,6 +149,62 @@
)
cc_library(
+ name = "graph_test_utils",
+ testonly = 1,
+ srcs = ["graph_test_utils.cc"],
+ hdrs = [
+ "graph_test_utils.h",
+ ],
+ visibility = ["//visibility:public"],
+ deps = [
+ "//tensorflow/core:framework",
+ "//tensorflow/core:lib",
+ "//tensorflow/core/grappler:mutable_graph_view",
+ "//tensorflow/core/grappler:grappler_item",
+ "//tensorflow/core/grappler:utils",
+ "//tensorflow/core:testlib",
+ ] + tf_protos_all(),
+)
+
+cc_library(
+ name = "hoist_random_uniform",
+ srcs = ["hoist_random_uniform.cc"],
+ hdrs = [
+ "hoist_random_uniform.h",
+ ],
+ visibility = ["//visibility:public"],
+ deps = [
+ ":function_utils",
+ ":graph_utils",
+ "//tensorflow/core:lib",
+ "//tensorflow/core/grappler:mutable_graph_view",
+ "//tensorflow/core/grappler:grappler_item",
+ "//tensorflow/core/grappler:op_types",
+ "//tensorflow/core/grappler:utils",
+ "//tensorflow/core/grappler/clusters:cluster",
+ "//tensorflow/core/grappler/optimizers:custom_graph_optimizer",
+ "//tensorflow/core/grappler/optimizers:custom_graph_optimizer_registry",
+ "//tensorflow/core:lib_internal",
+ ] + tf_protos_all(),
+)
+
+tf_cc_test(
+ name = "hoist_random_uniform_test",
+ srcs = ["hoist_random_uniform_test.cc"],
+ visibility = ["//visibility:public"],
+ deps = [
+ ":graph_test_utils",
+ ":graph_utils",
+ ":hoist_random_uniform",
+ "//tensorflow/core:framework",
+ "//tensorflow/core:test",
+ "//tensorflow/core:test_main",
+ "//tensorflow/core:testlib",
+ "//tensorflow/core/grappler:grappler_item",
+ ] + tf_protos_all(),
+)
+
+cc_library(
name = "latency_all_edges",
srcs = ["latency_all_edges.cc"],
hdrs = [
@@ -256,7 +315,7 @@
"//tensorflow/core/grappler/optimizers:custom_graph_optimizer",
"//tensorflow/core/grappler/utils:topological_sort",
"//tensorflow/core/grappler/optimizers:custom_graph_optimizer_registry",
- "//tensorflow/core:ptr_util",
+ "//tensorflow/core:lib_internal",
] + tf_protos_all(),
)
@@ -265,6 +324,7 @@
srcs = ["map_and_filter_fusion_test.cc"],
visibility = ["//visibility:public"],
deps = [
+ ":graph_test_utils",
":graph_utils",
":map_and_filter_fusion",
"//tensorflow/core:framework",
@@ -294,6 +354,7 @@
"//tensorflow/core/grappler/utils:topological_sort",
"//tensorflow/core/grappler/optimizers:custom_graph_optimizer",
"//tensorflow/core/grappler/optimizers:custom_graph_optimizer_registry",
+ "//tensorflow/core:lib_internal",
] + tf_protos_all(),
)
@@ -302,6 +363,7 @@
srcs = ["map_fusion_test.cc"],
visibility = ["//visibility:public"],
deps = [
+ ":graph_test_utils",
":graph_utils",
":map_fusion",
"//tensorflow/core:framework",
@@ -339,6 +401,7 @@
srcs = ["map_parallelization_test.cc"],
visibility = ["//visibility:public"],
deps = [
+ ":graph_test_utils",
":graph_utils",
":map_parallelization",
"//tensorflow/core:framework",
@@ -422,6 +485,7 @@
visibility = ["//visibility:public"],
deps = [
":filter_fusion",
+ ":hoist_random_uniform",
":latency_all_edges",
":map_and_batch_fusion",
":map_and_filter_fusion",
@@ -459,6 +523,7 @@
":function_utils",
":graph_utils",
"@com_google_absl//absl/strings",
+ "//tensorflow/core:core_cpu",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
@@ -474,6 +539,7 @@
srcs = ["vectorization_utils_test.cc"],
visibility = ["//visibility:public"],
deps = [
+ ":graph_utils",
":function_utils",
":vectorization_utils",
"//tensorflow/core:framework",
@@ -483,7 +549,10 @@
"//tensorflow/core/grappler:grappler_item",
"//tensorflow/core/grappler/optimizers:custom_graph_optimizer",
"//tensorflow/core/grappler/optimizers:custom_graph_optimizer_registry",
+ # For ops we need registered
+ "//tensorflow/core/kernels/data:dataset_ops",
"//tensorflow/core/kernels:cast_op",
+ "//tensorflow/core/kernels:logging_ops",
"//tensorflow/tools/graph_transforms:transform_utils",
] + tf_protos_all(),
)
diff --git a/tensorflow/core/grappler/optimizers/data/filter_fusion.cc b/tensorflow/core/grappler/optimizers/data/filter_fusion.cc
index c71aa6e..1ad495b 100644
--- a/tensorflow/core/grappler/optimizers/data/filter_fusion.cc
+++ b/tensorflow/core/grappler/optimizers/data/filter_fusion.cc
@@ -43,19 +43,14 @@
fused_node.set_op("FilterDataset");
fused_node.add_input(first_filter_node.input(0));
- auto copy_attribute = [](const string& attribute_name, const NodeDef& from,
- NodeDef* to) {
- (*to->mutable_attr())[attribute_name] = from.attr().at(attribute_name);
- };
-
auto attr = first_filter_node.attr().at("predicate");
*attr.mutable_func()->mutable_name() = fused_function.signature().name();
(*fused_node.mutable_attr())["predicate"] = std::move(attr);
- copy_attribute("Targuments", first_filter_node, &fused_node);
+ graph_utils::CopyAttribute("Targuments", first_filter_node, &fused_node);
for (auto key : {"output_shapes", "output_types"})
- copy_attribute(key, second_filter_node, &fused_node);
+ graph_utils::CopyAttribute(key, second_filter_node, &fused_node);
return fused_node;
}
@@ -120,8 +115,8 @@
// functions, or make sure that optimization passes run after filter
// fusion.
TF_RETURN_IF_ERROR(function_library.AddFunctionDef(*fused_predicate));
- // TODO(prazek): we could also remove map functions from library if they
- // are not used anymore.
+ // TODO(b/116285210): we could also remove map functions from library if
+ // they are not used anymore.
nodes_to_delete.insert(first_filter_node->name());
nodes_to_delete.insert(second_filter_node->name());
}
diff --git a/tensorflow/core/grappler/optimizers/data/filter_fusion_test.cc b/tensorflow/core/grappler/optimizers/data/filter_fusion_test.cc
index 12b1924..c8becc5 100644
--- a/tensorflow/core/grappler/optimizers/data/filter_fusion_test.cc
+++ b/tensorflow/core/grappler/optimizers/data/filter_fusion_test.cc
@@ -19,8 +19,8 @@
#include "tensorflow/core/framework/function_testlib.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/grappler/grappler_item.h"
+#include "tensorflow/core/grappler/optimizers/data/graph_test_utils.h"
#include "tensorflow/core/grappler/optimizers/data/graph_utils.h"
-
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
@@ -28,14 +28,7 @@
namespace grappler {
namespace {
-NodeDef MakeFilterNode(StringPiece name, StringPiece input_node_name) {
- return test::function::NDef(
- name, "FilterDataset", {string(input_node_name)},
- {{"predicate", FunctionDefHelper::FunctionRef("IsZero")},
- {"Targuments", {}},
- {"output_shapes", {}},
- {"output_types", {}}});
-}
+using graph_tests_utils::MakeFilterNode;
TEST(FilterFusionTest, FuseTwoFilterIntoOne) {
using test::function::NDef;
diff --git a/tensorflow/core/grappler/optimizers/data/function_utils.cc b/tensorflow/core/grappler/optimizers/data/function_utils.cc
index e95ea1a..311df15 100644
--- a/tensorflow/core/grappler/optimizers/data/function_utils.cc
+++ b/tensorflow/core/grappler/optimizers/data/function_utils.cc
@@ -14,31 +14,16 @@
==============================================================================*/
#include "tensorflow/core/grappler/optimizers/data/function_utils.h"
+#include "tensorflow/core/grappler/optimizers/data/graph_utils.h"
#include "tensorflow/core/framework/device_base.h"
#include "tensorflow/core/framework/op_def.pb.h"
+#include "tensorflow/core/lib/strings/scanner.h"
#include "tensorflow/core/util/ptr_util.h"
namespace tensorflow {
namespace grappler {
namespace function_utils {
-namespace {
-
-template <typename Predicate, typename Collection>
-std::vector<int> GetElementIndicesWithPredicate(const Predicate& predicate,
- const Collection& collection) {
- std::vector<int> indices = {};
- unsigned idx = 0;
- for (auto&& element : collection) {
- if (predicate(element)) {
- indices.push_back(idx);
- }
- idx++;
- }
- return indices;
-}
-
-} // namespace
FunctionDefTensorDesc::FunctionDefTensorDesc(const string& node_name,
const string& output, int position)
@@ -152,32 +137,27 @@
}
int FindFunctionInputWithName(StringPiece name, const FunctionDef& function) {
- std::vector<int> indices = GetElementIndicesWithPredicate(
+ return graph_utils::GetFirstElementIndexWithPredicate(
[&name](const OpDef_ArgDef& arg) { return arg.name() == name; },
function.signature().input_arg());
- return indices.empty() ? -1 : indices.front();
}
int FindFunctionOutputWithName(StringPiece name, const FunctionDef& function) {
- std::vector<int> indices = GetElementIndicesWithPredicate(
+ return graph_utils::GetFirstElementIndexWithPredicate(
[&name](const OpDef_ArgDef& arg) { return arg.name() == name; },
function.signature().output_arg());
- return indices.empty() ? -1 : indices.front();
}
int FindFunctionNodeWithName(StringPiece name, const FunctionDef& function) {
- std::vector<int> indices = GetElementIndicesWithPredicate(
+ return graph_utils::GetFirstElementIndexWithPredicate(
[&name](const NodeDef& node) { return node.name() == name; },
function.node_def());
- return indices.empty() ? -1 : indices.front();
}
int FindFunctionNodeWithOp(StringPiece op, const FunctionDef& function) {
- std::vector<int> indices = GetElementIndicesWithPredicate(
+ return graph_utils::GetFirstElementIndexWithPredicate(
[&op](const NodeDef& node) { return node.op() == op; },
function.node_def());
-
- return indices.empty() ? -1 : indices.front();
}
void SetUniqueFunctionNodeName(StringPiece prefix, FunctionDef* function,
diff --git a/tensorflow/core/grappler/optimizers/data/graph_test_utils.cc b/tensorflow/core/grappler/optimizers/data/graph_test_utils.cc
new file mode 100644
index 0000000..b2eec72
--- /dev/null
+++ b/tensorflow/core/grappler/optimizers/data/graph_test_utils.cc
@@ -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.
+==============================================================================*/
+
+#include "tensorflow/core/grappler/optimizers/data/graph_test_utils.h"
+
+#include "tensorflow/core/framework/function_testlib.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/grappler/grappler_item.h"
+
+namespace tensorflow {
+namespace grappler {
+namespace graph_tests_utils {
+
+NodeDef MakeMapNode(StringPiece name, StringPiece input_node_name,
+ StringPiece function_name) {
+ return test::function::NDef(
+ name, "MapDataset", {string(input_node_name)},
+ {{"f", FunctionDefHelper::FunctionRef(string(function_name))},
+ {"Targuments", {}},
+ {"output_shapes", gtl::ArraySlice<TensorShape>{}},
+ {"output_types", gtl::ArraySlice<DataType>{}}});
+}
+
+NodeDef MakeFilterNode(StringPiece name, StringPiece input_node_name,
+ StringPiece function_name) {
+ return test::function::NDef(
+ name, "FilterDataset", {string(input_node_name)},
+ {{"predicate", FunctionDefHelper::FunctionRef(string(function_name))},
+ {"Targuments", {}},
+ {"output_shapes", gtl::ArraySlice<TensorShape>{}},
+ {"output_types", gtl::ArraySlice<TensorShape>{}}});
+}
+
+} // end namespace graph_tests_utils
+} // end namespace grappler
+} // end namespace tensorflow
diff --git a/tensorflow/core/grappler/optimizers/data/graph_test_utils.h b/tensorflow/core/grappler/optimizers/data/graph_test_utils.h
new file mode 100644
index 0000000..ca0fde9
--- /dev/null
+++ b/tensorflow/core/grappler/optimizers/data/graph_test_utils.h
@@ -0,0 +1,36 @@
+/* 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_CORE_GRAPPLER_OPTIMIZERS_DATA_GRAPH_TEST_UTILS_H_
+#define TENSORFLOW_CORE_GRAPPLER_OPTIMIZERS_DATA_GRAPH_TEST_UTILS_H_
+
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/lib/core/stringpiece.h"
+
+namespace tensorflow {
+namespace grappler {
+namespace graph_tests_utils {
+
+NodeDef MakeMapNode(StringPiece name, StringPiece input_node_name,
+ StringPiece function_name = "XTimesTwo");
+
+NodeDef MakeFilterNode(StringPiece name, StringPiece input_node_name,
+ StringPiece function_name = "IsZero");
+
+} // end namespace graph_tests_utils
+} // end namespace grappler
+} // end namespace tensorflow
+
+#endif // TENSORFLOW_CORE_GRAPPLER_OPTIMIZERS_DATA_GRAPH_TEST_UTILS_H_
diff --git a/tensorflow/core/grappler/optimizers/data/graph_utils.cc b/tensorflow/core/grappler/optimizers/data/graph_utils.cc
index 2dd9ee8..3eaaf8f 100644
--- a/tensorflow/core/grappler/optimizers/data/graph_utils.cc
+++ b/tensorflow/core/grappler/optimizers/data/graph_utils.cc
@@ -201,25 +201,22 @@
int FindGraphFunctionWithName(StringPiece name,
const FunctionDefLibrary& library) {
- std::vector<int> indices = GetElementIndicesWithPredicate(
+ return GetFirstElementIndexWithPredicate(
[&name](const FunctionDef& function) {
return function.signature().name() == name;
},
library.function());
- return indices.empty() ? -1 : indices.front();
}
int FindGraphNodeWithName(StringPiece name, const GraphDef& graph) {
- std::vector<int> indices = GetElementIndicesWithPredicate(
+ return GetFirstElementIndexWithPredicate(
[&name](const NodeDef& node) { return node.name() == name; },
graph.node());
- return indices.empty() ? -1 : indices.front();
}
int FindGraphNodeWithOp(StringPiece op, const GraphDef& graph) {
- std::vector<int> indices = GetElementIndicesWithPredicate(
+ return GetFirstElementIndexWithPredicate(
[&op](const NodeDef& node) { return node.op() == op; }, graph.node());
- return indices.empty() ? -1 : indices.front();
}
std::vector<int> FindAllGraphNodesWithOp(const string& op,
@@ -260,6 +257,21 @@
}
function->mutable_signature()->set_name(std::move(name));
}
+
+void CopyAttribute(const string& attribute_name, const NodeDef& from,
+ NodeDef* to_node) {
+ (*to_node->mutable_attr())[attribute_name] = from.attr().at(attribute_name);
+}
+
+void ConcatAttributeList(const string& attribute_name, const NodeDef& first,
+ const NodeDef& second, NodeDef* to_node) {
+ CopyAttribute(attribute_name, first, to_node);
+ (*to_node->mutable_attr())
+ .at(attribute_name)
+ .mutable_list()
+ ->MergeFrom(second.attr().at(attribute_name).list());
+}
+
} // end namespace graph_utils
} // end namespace grappler
} // end namespace tensorflow
diff --git a/tensorflow/core/grappler/optimizers/data/graph_utils.h b/tensorflow/core/grappler/optimizers/data/graph_utils.h
index b117482..3af34f6 100644
--- a/tensorflow/core/grappler/optimizers/data/graph_utils.h
+++ b/tensorflow/core/grappler/optimizers/data/graph_utils.h
@@ -31,6 +31,21 @@
namespace grappler {
namespace graph_utils {
+// Returns the index of the first element in collection that fulfills predicate.
+// If no such element exists, returns -1.
+template <typename Predicate, typename Collection>
+int GetFirstElementIndexWithPredicate(const Predicate& predicate,
+ const Collection& collection) {
+ unsigned idx = 0;
+ for (auto&& element : collection) {
+ if (predicate(element)) {
+ return idx;
+ }
+ idx++;
+ }
+ return -1;
+}
+
// Adds a node to the graph.
NodeDef* AddNode(StringPiece name, StringPiece op,
const std::vector<string>& inputs,
@@ -101,11 +116,21 @@
// is unique across the graph.
void SetUniqueGraphNodeName(StringPiece prefix, GraphDef* graph, NodeDef* node);
-// Sets the node name using the `prefix` name as a prefix while guaranteeing the
-// name is unique across the graph.
+// Sets the function name using the `prefix` name as a prefix while guaranteeing
+// the name is unique across the function library.
void SetUniqueGraphFunctionName(StringPiece prefix, FunctionDefLibrary* library,
FunctionDef* function);
+// Copies attribute having name `attribute_name` from node `from` to node
+// `to_node`.
+void CopyAttribute(const string& attribute_name, const NodeDef& from,
+ NodeDef* to_node);
+
+// Concatenates list attribute having name `attribute_name` from `first` and
+// `second` node, setting it to `to_node`.
+void ConcatAttributeList(const string& attribute_name, const NodeDef& first,
+ const NodeDef& second, NodeDef* to_node);
+
} // end namespace graph_utils
} // end namespace grappler
} // end namespace tensorflow
diff --git a/tensorflow/core/grappler/optimizers/data/graph_utils_test.cc b/tensorflow/core/grappler/optimizers/data/graph_utils_test.cc
index 6877c20..db98654 100644
--- a/tensorflow/core/grappler/optimizers/data/graph_utils_test.cc
+++ b/tensorflow/core/grappler/optimizers/data/graph_utils_test.cc
@@ -24,6 +24,18 @@
namespace graph_utils {
namespace {
+TEST(GraphUtilsTest, GetFirstElementIndexWithPredicate) {
+ std::vector<int> vec({1, 2, 3, 4, 5, 6});
+ auto result = GetFirstElementIndexWithPredicate(
+ [](int elem) { return elem % 3 == 0; }, vec);
+
+ EXPECT_EQ(result, 2);
+
+ result = GetFirstElementIndexWithPredicate(
+ [](int elem) { return elem % 7 == 0; }, vec);
+ EXPECT_EQ(result, -1);
+}
+
TEST(GraphUtilsTest, AddScalarConstNodeBool) {
GraphDef graph_def;
MutableGraphView graph(&graph_def);
diff --git a/tensorflow/core/grappler/optimizers/data/hoist_random_uniform.cc b/tensorflow/core/grappler/optimizers/data/hoist_random_uniform.cc
new file mode 100644
index 0000000..ce0b2db
--- /dev/null
+++ b/tensorflow/core/grappler/optimizers/data/hoist_random_uniform.cc
@@ -0,0 +1,289 @@
+/* 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/grappler/optimizers/data/hoist_random_uniform.h"
+
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/op_def.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/grappler/clusters/cluster.h"
+#include "tensorflow/core/grappler/grappler_item.h"
+#include "tensorflow/core/grappler/mutable_graph_view.h"
+#include "tensorflow/core/grappler/op_types.h"
+#include "tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.h"
+#include "tensorflow/core/grappler/optimizers/data/function_utils.h"
+#include "tensorflow/core/grappler/optimizers/data/graph_utils.h"
+#include "tensorflow/core/grappler/utils.h"
+#include "tensorflow/core/lib/gtl/map_util.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/platform/protobuf.h"
+
+namespace tensorflow {
+namespace grappler {
+namespace {
+
+NodeDef MakeStatelessMap(const NodeDef& map_node, const NodeDef& zip_node,
+ const FunctionDef& stateless_function,
+ MutableGraphView* graph) {
+ NodeDef stateless_map;
+ graph_utils::SetUniqueGraphNodeName("stateless_map", graph->GetGraph(),
+ &stateless_map);
+
+ stateless_map.set_op("MapDataset");
+ stateless_map.add_input(zip_node.name());
+ // Add placeholders.
+ for (int i = 1; i < map_node.input_size(); i++)
+ stateless_map.add_input(map_node.input(i));
+
+ auto attr = map_node.attr().at("f");
+ *attr.mutable_func()->mutable_name() = stateless_function.signature().name();
+ *attr.mutable_func()->mutable_attr() = stateless_function.attr();
+ (*stateless_map.mutable_attr())["f"] = std::move(attr);
+
+ graph_utils::CopyAttribute("Targuments", map_node, &stateless_map);
+ for (auto key : {"output_shapes", "output_types"})
+ graph_utils::CopyAttribute(key, map_node, &stateless_map);
+
+ if (const auto* attr =
+ gtl::FindOrNull(map_node.attr(), "use_inter_op_parallelism"))
+ (*stateless_map.mutable_attr())["use_inter_op_parallelism"] = *attr;
+
+ return stateless_map;
+}
+
+NodeDef MakeRandomDataset(const NodeDef& random_uniform_node,
+ MutableGraphView* graph) {
+ NodeDef random_dataset;
+ random_dataset.set_op("RandomDataset");
+ graph_utils::SetUniqueGraphNodeName("RandomDataset", graph->GetGraph(),
+ &random_dataset);
+
+ const auto* seed = graph_utils::AddScalarConstNode<int64>(
+ random_uniform_node.attr().at("seed").i(), graph);
+ const auto* seed2 = graph_utils::AddScalarConstNode<int64>(
+ random_uniform_node.attr().at("seed2").i(), graph);
+
+ random_dataset.add_input(seed->name());
+ random_dataset.add_input(seed2->name());
+
+ (*random_dataset.mutable_attr())["output_shapes"].mutable_list()->add_shape();
+ (*random_dataset.mutable_attr())["output_types"].mutable_list()->add_type(
+ DT_INT64);
+
+ return random_dataset;
+}
+
+NodeDef MakeBatchTwo(const NodeDef& random_dataset, MutableGraphView* graph) {
+ NodeDef batch_dataset;
+ batch_dataset.set_op("BatchDatasetV2");
+ graph_utils::SetUniqueGraphNodeName("pair_of_random", graph->GetGraph(),
+ &batch_dataset);
+ const auto* batch_size = graph_utils::AddScalarConstNode<int64>(2, graph);
+ const auto* drop_reminder = graph_utils::AddScalarConstNode(false, graph);
+ batch_dataset.add_input(random_dataset.name());
+ batch_dataset.add_input(batch_size->name());
+ batch_dataset.add_input(drop_reminder->name());
+
+ (*batch_dataset.mutable_attr())["output_shapes"]
+ .mutable_list()
+ ->add_shape()
+ ->mutable_dim()
+ ->Add()
+ ->set_size(-1);
+ (*batch_dataset.mutable_attr())["output_types"].mutable_list()->add_type(
+ DT_INT64);
+
+ return batch_dataset;
+}
+
+NodeDef MakeZipNode(const NodeDef& first_node, const NodeDef& second_node,
+ MutableGraphView* graph) {
+ NodeDef zip_node;
+ graph_utils::SetUniqueGraphNodeName("zip_with_random", graph->GetGraph(),
+ &zip_node);
+
+ zip_node.set_op("ZipDataset");
+ zip_node.add_input(first_node.name());
+ zip_node.add_input(second_node.name());
+
+ for (auto key : {"output_shapes", "output_types"})
+ graph_utils::ConcatAttributeList(key, first_node, second_node, &zip_node);
+
+ (*zip_node.mutable_attr())["N"].set_i(2);
+
+ return zip_node;
+}
+
+// We need to insert our argument before the placeholders, which are the last
+// arguments.
+OpDef_ArgDef* InsertSeedArgument(OpDef* signature, int num_placeholders) {
+ int new_argument_idx = signature->input_arg_size() - num_placeholders;
+ signature->add_input_arg();
+ for (int i = signature->input_arg_size() - 1; i > new_argument_idx; i--) {
+ signature->mutable_input_arg()->SwapElements(i - 1, i);
+ }
+ auto* seed_arg = signature->mutable_input_arg(new_argument_idx);
+ seed_arg->set_name(strings::StrCat("seed_arg", new_argument_idx));
+ seed_arg->set_type(DT_INT64);
+
+ return seed_arg;
+}
+
+// Make function that uses `StatelessRandomUniform` instead of `RandomUniform`
+// to make it less statefull. The function can still be stateful, but in when
+// other stateful ops are e.g. `Assert`, then it will be parallelizable.
+const FunctionDef* MakeLessStatefulFunction(const FunctionDef& map_function,
+ bool is_stateful,
+ int num_placeholders,
+ FunctionDefLibrary* library) {
+ FunctionDef* stateless_function = library->add_function();
+ *stateless_function = map_function;
+ if (is_stateful)
+ stateless_function->mutable_signature()->set_is_stateful(is_stateful);
+ graph_utils::SetUniqueGraphFunctionName("stateless_function", library,
+ stateless_function);
+
+ auto* seed_arg = InsertSeedArgument(stateless_function->mutable_signature(),
+ num_placeholders);
+
+ auto* const random_uniform = stateless_function->mutable_node_def(
+ function_utils::FindFunctionNodeWithOp("RandomUniform",
+ *stateless_function));
+
+ // Replace RandomUniform node with StatelessRandomUniform.
+ random_uniform->set_op("StatelessRandomUniform");
+ random_uniform->add_input(seed_arg->name());
+ (*random_uniform->mutable_attr())["Tseed"].set_type(DT_INT64);
+ random_uniform->mutable_attr()->erase("seed");
+ random_uniform->mutable_attr()->erase("seed2");
+
+ return stateless_function;
+}
+// This function returns true if function is stateful and has single
+// RandomUniform op and no other stateful ops except Assert.
+// `is_stateful_after_hoisting` is set to true if RandomUniform is the only
+// stateful op and hoisting can be performed.
+bool CanHoistRandomUniform(const FunctionDef& map_function,
+ const FunctionLibraryDefinition& library,
+ bool* is_stateful_after_hoisting,
+ const NodeDef** random_uniform_op) {
+ if (!map_function.signature().is_stateful()) return false;
+ *is_stateful_after_hoisting = true;
+
+ bool have_other_stateful_ops = false;
+
+ for (const auto& node : map_function.node_def()) {
+ const OpDef* op_def;
+ TF_CHECK_OK(library.LookUpOpDef(node.op(), &op_def));
+ // Skip stateless nodes and assert, as it does not actually have a state.
+ if (!op_def->is_stateful()) continue;
+
+ if (op_def->name() == "Assert") {
+ have_other_stateful_ops = true;
+ continue;
+ }
+
+ // TODO(prazek): For now we only handle RandomUniform, we should handle
+ // RandomUniformInt as well.
+ if (op_def->name() != "RandomUniform") return false;
+
+ // TODO(prazek): For now we can only hoist single RandomUniform.
+ if (*random_uniform_op != nullptr) return false;
+
+ *random_uniform_op = &node;
+ }
+
+ if (!have_other_stateful_ops) *is_stateful_after_hoisting = false;
+
+ // Have we found single RandomUniform?
+ return *random_uniform_op != nullptr;
+}
+
+int NumberOfPlaceholders(const NodeDef& map_node) {
+ // First input of MapDataset is the argument to the function. Rest of the
+ // inputs are placeholders.
+ return map_node.input_size() - 1;
+}
+
+} // namespace
+
+Status HoistRandomUniform::Optimize(Cluster* cluster, const GrapplerItem& item,
+ GraphDef* output) {
+ *output = item.graph;
+
+ MutableGraphView graph(output);
+ std::set<string> nodes_to_delete;
+ FunctionLibraryDefinition function_library(OpRegistry::Global(),
+ item.graph.library());
+
+ auto get_map_node = [](const NodeDef& node) -> const NodeDef* {
+ // TODO(prazek): we could also handle ParallelMapDataset and
+ // MapAndBatchDataset.
+ if (node.op() == "MapDataset") return &node;
+ return nullptr;
+ };
+
+ for (const NodeDef& node : item.graph.node()) {
+ const NodeDef* map_node = get_map_node(node);
+ if (!map_node) continue;
+
+ const auto& fun = map_node->attr().at("f");
+ const FunctionDef* func = function_library.Find(fun.func().name());
+
+ const NodeDef* random_uniform_op = nullptr;
+ bool is_stateful_after_hoisting = true;
+ if (!CanHoistRandomUniform(*func, function_library,
+ &is_stateful_after_hoisting, &random_uniform_op))
+ continue;
+ const auto* random_seed_dataset =
+ graph.AddNode(MakeRandomDataset(*random_uniform_op, &graph));
+
+ const auto* batch_dataset =
+ graph.AddNode(MakeBatchTwo(*random_seed_dataset, &graph));
+
+ const NodeDef& parent_node = *graph_utils::GetInputNode(*map_node, graph);
+
+ const auto* zip_node =
+ graph.AddNode(MakeZipNode(parent_node, *batch_dataset, &graph));
+
+ const auto* stateless_func = MakeLessStatefulFunction(
+ *func, is_stateful_after_hoisting, NumberOfPlaceholders(*map_node),
+ output->mutable_library());
+
+ const auto* stateless_map = graph.AddNode(
+ MakeStatelessMap(*map_node, *zip_node, *stateless_func, &graph));
+
+ graph.ReplaceInput(*map_node, *stateless_map);
+
+ // TODO(b/116285210): we could also remove map functions from library if
+ // they are not used anymore.
+ nodes_to_delete.insert(map_node->name());
+ }
+
+ graph.DeleteNodes(nodes_to_delete);
+ return Status::OK();
+}
+
+void HoistRandomUniform::Feedback(Cluster* cluster, const GrapplerItem& item,
+ const GraphDef& optimize_output,
+ double result) {
+ // no-op
+}
+
+REGISTER_GRAPH_OPTIMIZER_AS(HoistRandomUniform, "hoist_random_uniform");
+
+} // end namespace grappler
+} // end namespace tensorflow
diff --git a/tensorflow/core/grappler/optimizers/data/hoist_random_uniform.h b/tensorflow/core/grappler/optimizers/data/hoist_random_uniform.h
new file mode 100644
index 0000000..d1bcf67
--- /dev/null
+++ b/tensorflow/core/grappler/optimizers/data/hoist_random_uniform.h
@@ -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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_GRAPPLER_OPTIMIZERS_DATA_HOIST_RANDOM_UNIFORM_H_
+#define TENSORFLOW_CORE_GRAPPLER_OPTIMIZERS_DATA_HOIST_RANDOM_UNIFORM_H_
+
+#include "tensorflow/core/grappler/optimizers/custom_graph_optimizer.h"
+
+namespace tensorflow {
+namespace grappler {
+
+// This optimization hoists instances of `random_uniform` out of a function
+// with the aim of making it stateless. It creates a new function that takes a
+// random seed as an extra argument and uses `stateless_random_uniform` instead
+// of `random_uniform` to make it stateless.
+// It also creates RandomDataset(seed).batch(2), which is zipped with old input
+// to the map. The batching in RandomDataset is because we need 2 seeds for
+// `stateless_random_uniform`.
+// TODO(prazek): for now only `RandomUniform` is handled, but we could handle
+// `RandomUniformInt` similarly.
+class HoistRandomUniform : public CustomGraphOptimizer {
+ public:
+ HoistRandomUniform() = default;
+ ~HoistRandomUniform() override = default;
+
+ string name() const override { return "hoist_random_uniform"; };
+
+ Status Init(
+ const tensorflow::RewriterConfig_CustomGraphOptimizer* config) override {
+ return Status::OK();
+ }
+
+ Status Optimize(Cluster* cluster, const GrapplerItem& item,
+ GraphDef* output) override;
+
+ void Feedback(Cluster* cluster, const GrapplerItem& item,
+ const GraphDef& optimize_output, double result) override;
+};
+
+} // end namespace grappler
+} // end namespace tensorflow
+
+#endif // TENSORFLOW_CORE_GRAPPLER_OPTIMIZERS_DATA_HOIST_RANDOM_UNIFORM_H_
diff --git a/tensorflow/core/grappler/optimizers/data/hoist_random_uniform_test.cc b/tensorflow/core/grappler/optimizers/data/hoist_random_uniform_test.cc
new file mode 100644
index 0000000..455459e
--- /dev/null
+++ b/tensorflow/core/grappler/optimizers/data/hoist_random_uniform_test.cc
@@ -0,0 +1,84 @@
+/* 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/grappler/optimizers/data/hoist_random_uniform.h"
+
+#include "tensorflow/core/framework/function_testlib.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/grappler/grappler_item.h"
+#include "tensorflow/core/grappler/optimizers/data/graph_test_utils.h"
+#include "tensorflow/core/grappler/optimizers/data/graph_utils.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace grappler {
+namespace {
+
+TEST(HoistRandomUniform, SimpleHoisting) {
+ using test::function::NDef;
+ GrapplerItem item;
+ item.graph = test::function::GDef(
+ {NDef("start", "Const", {}, {{"value", 0}, {"dtype", DT_INT32}}),
+ NDef("stop", "Const", {}, {{"value", 10}, {"dtype", DT_INT32}}),
+ NDef("step", "Const", {}, {{"value", 1}, {"dtype", DT_INT32}}),
+ NDef("filename", "Const", {}, {{"value", ""}, {"dtype", DT_STRING}}),
+ NDef("range", "RangeDataset", {"start", "stop", "step"},
+ {{"output_shapes", gtl::ArraySlice<TensorShape>{}},
+ {"output_types", gtl::ArraySlice<DataType>{}}}),
+ graph_tests_utils::MakeMapNode("map1", "range", "RandomUniform"),
+ NDef("cache", "CacheDataset", {"map1", "filename"}, {})},
+ // FunctionLib
+ {
+ test::function::RandomUniform(),
+ });
+
+ HoistRandomUniform optimizer;
+ GraphDef output;
+ TF_ASSERT_OK(optimizer.Optimize(nullptr, item, &output));
+
+ EXPECT_FALSE(graph_utils::ContainsGraphNodeWithName("map1", output));
+ const int new_map_id = graph_utils::FindGraphNodeWithOp("MapDataset", output);
+ const int zip_dataset_id =
+ graph_utils::FindGraphNodeWithOp("ZipDataset", output);
+ const int random_dataset_id =
+ graph_utils::FindGraphNodeWithOp("RandomDataset", output);
+ const int batch_random_id =
+ graph_utils::FindGraphNodeWithOp("BatchDatasetV2", output);
+ ASSERT_NE(random_dataset_id, -1);
+ ASSERT_NE(zip_dataset_id, -1);
+ ASSERT_NE(new_map_id, -1);
+ ASSERT_NE(batch_random_id, -1);
+
+ const auto& new_map = output.node(new_map_id);
+ const auto& zip = output.node(zip_dataset_id);
+ const auto& random = output.node(random_dataset_id);
+ const auto& batch = output.node(batch_random_id);
+
+ ASSERT_EQ(new_map.input_size(), 1);
+ EXPECT_EQ(new_map.input(0), zip.name());
+
+ ASSERT_EQ(zip.input_size(), 2);
+ EXPECT_EQ(zip.input(0), "range");
+ EXPECT_EQ(zip.input(1), batch.name());
+
+ ASSERT_EQ(batch.input_size(), 3);
+ EXPECT_EQ(batch.input(0), random.name());
+}
+
+} // namespace
+} // namespace grappler
+} // namespace tensorflow
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 63945b8..e66766e 100644
--- a/tensorflow/core/grappler/optimizers/data/map_and_batch_fusion.cc
+++ b/tensorflow/core/grappler/optimizers/data/map_and_batch_fusion.cc
@@ -80,11 +80,12 @@
// Set `f` and `Targuments` attributes.
for (auto key : {"f", "Targuments"}) {
- (*new_node.mutable_attr())[key] = map_node.attr().at(key);
+ graph_utils::CopyAttribute(key, map_node, &new_node);
}
+
// Set `output_types` and `output_shapes` attributes.
for (auto key : {"output_shapes", "output_types"}) {
- (*new_node.mutable_attr())[key] = batch_node.attr().at(key);
+ graph_utils::CopyAttribute(key, batch_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 f1844a1..c4868ea 100644
--- a/tensorflow/core/grappler/optimizers/data/map_and_filter_fusion.cc
+++ b/tensorflow/core/grappler/optimizers/data/map_and_filter_fusion.cc
@@ -26,6 +26,7 @@
#include "tensorflow/core/grappler/optimizers/data/graph_utils.h"
#include "tensorflow/core/grappler/utils.h"
#include "tensorflow/core/grappler/utils/topological_sort.h"
+#include "tensorflow/core/lib/gtl/map_util.h"
#include "tensorflow/core/platform/protobuf.h"
namespace tensorflow {
@@ -41,19 +42,18 @@
fused_node.set_op("MapDataset");
fused_node.add_input(map_node.input(0));
- auto copy_attribute = [](const string& attribute_name, const NodeDef& from,
- NodeDef* to) {
- (*to->mutable_attr())[attribute_name] = from.attr().at(attribute_name);
- };
-
auto attr = map_node.attr().at("f");
attr.mutable_func()->set_name(fused_function.signature().name());
(*fused_node.mutable_attr())["f"] = std::move(attr);
- copy_attribute("Targuments", map_node, &fused_node);
+ graph_utils::CopyAttribute("Targuments", map_node, &fused_node);
for (auto key : {"output_shapes", "output_types"})
- copy_attribute(key, map_node, &fused_node);
+ graph_utils::CopyAttribute(key, map_node, &fused_node);
+
+ if (const auto* attr =
+ gtl::FindOrNull(map_node.attr(), "use_inter_op_parallelism"))
+ (*fused_node.mutable_attr())["use_inter_op_parallelism"] = *attr;
// Add the predicate output attributes.
(*fused_node.mutable_attr())["output_types"]
diff --git a/tensorflow/core/grappler/optimizers/data/map_and_filter_fusion_test.cc b/tensorflow/core/grappler/optimizers/data/map_and_filter_fusion_test.cc
index f029a09..6e6da37 100644
--- a/tensorflow/core/grappler/optimizers/data/map_and_filter_fusion_test.cc
+++ b/tensorflow/core/grappler/optimizers/data/map_and_filter_fusion_test.cc
@@ -19,6 +19,7 @@
#include "tensorflow/core/framework/function_testlib.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/grappler/grappler_item.h"
+#include "tensorflow/core/grappler/optimizers/data/graph_test_utils.h"
#include "tensorflow/core/grappler/optimizers/data/graph_utils.h"
#include "tensorflow/core/lib/core/status_test_util.h"
@@ -27,24 +28,8 @@
namespace tensorflow {
namespace grappler {
namespace {
-
-NodeDef MakeMapNode(StringPiece name, StringPiece input_node_name) {
- return test::function::NDef(
- name, "MapDataset", {string(input_node_name)},
- {{"f", FunctionDefHelper::FunctionRef("XTimesTwo")},
- {"Targuments", {}},
- {"output_shapes", {}},
- {"output_types", {}}});
-}
-
-NodeDef MakeFilterNode(StringPiece name, StringPiece input_node_name) {
- return test::function::NDef(
- name, "FilterDataset", {string(input_node_name)},
- {{"predicate", FunctionDefHelper::FunctionRef("IsZero")},
- {"Targuments", {}},
- {"output_shapes", {}},
- {"output_types", {}}});
-}
+using graph_tests_utils::MakeFilterNode;
+using graph_tests_utils::MakeMapNode;
TEST(MapAndFilterFusionTest, FuseMapAndFilter) {
using test::function::NDef;
diff --git a/tensorflow/core/grappler/optimizers/data/map_fusion.cc b/tensorflow/core/grappler/optimizers/data/map_fusion.cc
index a78ecb0..bd94334 100644
--- a/tensorflow/core/grappler/optimizers/data/map_fusion.cc
+++ b/tensorflow/core/grappler/optimizers/data/map_fusion.cc
@@ -26,6 +26,7 @@
#include "tensorflow/core/grappler/optimizers/data/graph_utils.h"
#include "tensorflow/core/grappler/utils.h"
#include "tensorflow/core/grappler/utils/topological_sort.h"
+#include "tensorflow/core/lib/gtl/map_util.h"
#include "tensorflow/core/platform/protobuf.h"
namespace tensorflow {
@@ -40,24 +41,31 @@
NodeDef fused_node;
graph_utils::SetUniqueGraphNodeName("fused_map", graph->GetGraph(),
&fused_node);
-
fused_node.set_op("MapDataset");
fused_node.add_input(parent_map_node.input(0));
- auto copy_attribute = [](const string& attribute_name, const NodeDef& from,
- NodeDef* to) {
- (*to->mutable_attr())[attribute_name] = from.attr().at(attribute_name);
- };
-
auto attr = parent_map_node.attr().at("f");
*attr.mutable_func()->mutable_name() = fused_function.signature().name();
(*fused_node.mutable_attr())["f"] = std::move(attr);
- copy_attribute("Targuments", parent_map_node, &fused_node);
-
+ graph_utils::CopyAttribute("Targuments", parent_map_node, &fused_node);
for (auto key : {"output_shapes", "output_types"})
- copy_attribute(key, map_node, &fused_node);
+ graph_utils::CopyAttribute(key, map_node, &fused_node);
+ auto value_or_false = [](const AttrValue* attr) {
+ if (!attr) return false;
+ return attr->b();
+ };
+
+ const auto* first_parallelism =
+ gtl::FindOrNull(parent_map_node.attr(), "use_inter_op_parallelism");
+ const auto* second_parallelism =
+ 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);
+ }
return fused_node;
}
@@ -123,8 +131,8 @@
// fusion.
TF_RETURN_IF_ERROR(function_library.AddFunctionDef(*fused_function));
- // TODO(prazek): we could also remove map functions from library if they
- // are not used anymore.
+ // TODO(b/116285210): we could also remove map functions from library if
+ // they are not used anymore.
nodes_to_delete.insert(parent_map_node->name());
nodes_to_delete.insert(map_node->name());
}
diff --git a/tensorflow/core/grappler/optimizers/data/map_fusion_test.cc b/tensorflow/core/grappler/optimizers/data/map_fusion_test.cc
index b25dfbd..8889f9d 100644
--- a/tensorflow/core/grappler/optimizers/data/map_fusion_test.cc
+++ b/tensorflow/core/grappler/optimizers/data/map_fusion_test.cc
@@ -19,6 +19,7 @@
#include "tensorflow/core/framework/function_testlib.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/grappler/grappler_item.h"
+#include "tensorflow/core/grappler/optimizers/data/graph_test_utils.h"
#include "tensorflow/core/grappler/optimizers/data/graph_utils.h"
#include "tensorflow/core/lib/core/status_test_util.h"
@@ -28,14 +29,7 @@
namespace grappler {
namespace {
-NodeDef MakeMapNode(StringPiece name, StringPiece input_node_name) {
- return test::function::NDef(
- name, "MapDataset", {string(input_node_name)},
- {{"f", FunctionDefHelper::FunctionRef("XTimesTwo")},
- {"Targuments", {}},
- {"output_shapes", {}},
- {"output_types", {}}});
-}
+using graph_tests_utils::MakeMapNode;
TEST(MapFusionTest, FuseTwoMapNodesIntoOne) {
using test::function::NDef;
diff --git a/tensorflow/core/grappler/optimizers/data/map_parallelization.cc b/tensorflow/core/grappler/optimizers/data/map_parallelization.cc
index 305325e..782c9f4 100644
--- a/tensorflow/core/grappler/optimizers/data/map_parallelization.cc
+++ b/tensorflow/core/grappler/optimizers/data/map_parallelization.cc
@@ -84,9 +84,6 @@
auto* parallel_map = graph.AddNode(MakeParallelMap(*map_node, &graph));
graph.ReplaceInput(*map_node, *parallel_map);
-
- // TODO(prazek): we could also remove map functions from library if they
- // are not used anymore.
nodes_to_delete.insert(map_node->name());
}
diff --git a/tensorflow/core/grappler/optimizers/data/map_parallelization_test.cc b/tensorflow/core/grappler/optimizers/data/map_parallelization_test.cc
index b2a5d9b..9fdfe8a 100644
--- a/tensorflow/core/grappler/optimizers/data/map_parallelization_test.cc
+++ b/tensorflow/core/grappler/optimizers/data/map_parallelization_test.cc
@@ -19,8 +19,8 @@
#include "tensorflow/core/framework/function_testlib.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/grappler/grappler_item.h"
+#include "tensorflow/core/grappler/optimizers/data/graph_test_utils.h"
#include "tensorflow/core/grappler/optimizers/data/graph_utils.h"
-
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
@@ -28,16 +28,7 @@
namespace grappler {
namespace {
-NodeDef MakeMapNode(StringPiece name, StringPiece input_node_name,
- StringPiece function_name) {
- return test::function::NDef(
- name, "MapDataset", {string(input_node_name)},
- {{"f", FunctionDefHelper::FunctionRef(string(function_name))},
- {"Targuments", {}},
- {"output_shapes", {}},
- {"output_types", {}}});
-}
-
+using graph_tests_utils::MakeMapNode;
const char stateless_fun_name[] = "XTimesTwo";
const char stateful_fun_name[] = "RandomUniform";
diff --git a/tensorflow/core/grappler/optimizers/data/map_vectorization.cc b/tensorflow/core/grappler/optimizers/data/map_vectorization.cc
index 7a2f191..9328a7c 100644
--- a/tensorflow/core/grappler/optimizers/data/map_vectorization.cc
+++ b/tensorflow/core/grappler/optimizers/data/map_vectorization.cc
@@ -35,10 +35,6 @@
namespace grappler {
namespace {
-void CopyAttribute(const string& attr_name, const NodeDef& from, NodeDef* to) {
- (*to->mutable_attr())[attr_name] = from.attr().at(attr_name);
-}
-
// Returns a FunctionDef containing a MapDefun op that wraps the original
// function.
FunctionDef* CreateMapDefunWrapper(const NodeDef& map_node,
@@ -61,7 +57,7 @@
for (const string& k : {"f", "output_types", "output_shapes"}) {
// Function, output types and (unbatched) shapes are the same as the
// original map node.
- CopyAttribute(k, map_node, map_defun_node);
+ graph_utils::CopyAttribute(k, map_node, map_defun_node);
}
// Get types of input arguments from original map function
@@ -90,21 +86,19 @@
// efficient vectorization with VectorizeMapDefun.
FunctionDef* vectorized_func =
CreateMapDefunWrapper(map_node, orig_func, library);
- NodeDef* map_defun_node = vectorized_func->mutable_node_def()->Mutable(0);
- DCHECK_EQ(map_defun_node->op(), "MapDefun");
+ const NodeDef& map_defun_node = vectorized_func->node_def(0);
+ DCHECK_EQ(map_defun_node.op(), "MapDefun");
- // Create a copy of the original function so that we can mutate it, and
- // attach that to the map defun node.
- FunctionDef* map_defun_fn = library->add_function();
- *map_defun_fn = orig_func;
- graph_utils::SetUniqueGraphFunctionName(orig_func.signature().name(), library,
- map_defun_fn);
- (*map_defun_node->mutable_attr())["f"].mutable_func()->set_name(
- map_defun_fn->signature().name());
+ // TODO(b/116285210): Unreferenced functions should get cleaned up later
+ FunctionDef* result;
+ Status s = vectorization_utils::VectorizeMapDefun(
+ *vectorized_func, map_defun_node, library, &result);
- vectorization_utils::VectorizeMapDefun(vectorized_func, map_defun_fn,
- map_defun_node);
- return vectorized_func;
+ if (!s.ok()) {
+ LOG(ERROR) << "VectorizeMapDefun failed: " << s;
+ return vectorized_func;
+ }
+ return result;
}
bool IsOutputShapesFullyDefined(const NodeDef& node) {
@@ -195,13 +189,16 @@
}
// Set attrs
- CopyAttribute("Targuments", old_map_node, &map_node);
+ graph_utils::CopyAttribute("Targuments", old_map_node, &map_node);
auto& func_attr = (*map_node.mutable_attr())["f"];
func_attr.mutable_func()->set_name(vectorized_func.signature().name());
for (auto key : {"output_shapes", "output_types"}) {
- CopyAttribute(key, old_batch_node, &map_node);
+ graph_utils::CopyAttribute(key, old_batch_node, &map_node);
}
+
+ (*map_node.mutable_attr())["use_inter_op_parallelism"].set_b(true);
+
return map_node;
}
diff --git a/tensorflow/core/grappler/optimizers/data/map_vectorization_test.cc b/tensorflow/core/grappler/optimizers/data/map_vectorization_test.cc
index ed1bd6b..f4faf41 100644
--- a/tensorflow/core/grappler/optimizers/data/map_vectorization_test.cc
+++ b/tensorflow/core/grappler/optimizers/data/map_vectorization_test.cc
@@ -30,72 +30,51 @@
using test::function::GDef;
using test::function::NDef;
-void MakeTensorShapeProtoHelper(const gtl::ArraySlice<int> dims,
- TensorShapeProto* t) {
- for (size_t i = 0; i < dims.size(); ++i) {
- auto* d = t->add_dim();
- d->set_size(dims[i]);
- }
-}
-
-AttrValue MakeShapeListAttr(
- const gtl::ArraySlice<const gtl::ArraySlice<int>>& shapes) {
- AttrValue shapes_attr;
- for (size_t i = 0; i < shapes.size(); ++i) {
- MakeTensorShapeProtoHelper(shapes[i],
- shapes_attr.mutable_list()->add_shape());
- }
-
- return shapes_attr;
-}
-
-NodeDef MakeMapNodeHelper(
- StringPiece name, StringPiece input_node_name, StringPiece function_name,
- StringPiece map_op_name,
- const gtl::ArraySlice<const gtl::ArraySlice<int>>& output_shapes,
- const gtl::ArraySlice<DataType>& output_types) {
+NodeDef MakeMapNodeHelper(StringPiece name, StringPiece input_node_name,
+ StringPiece function_name, StringPiece map_op_name,
+ gtl::ArraySlice<PartialTensorShape> output_shapes,
+ gtl::ArraySlice<DataType> output_types) {
return test::function::NDef(
name, map_op_name, {string(input_node_name)},
{{"f", FunctionDefHelper::FunctionRef(string(function_name))},
{"Targuments", {}},
- {"output_shapes", MakeShapeListAttr(output_shapes)},
+ {"output_shapes", output_shapes},
{"output_types", output_types}});
}
-NodeDef MakeMapNode(
- StringPiece name, StringPiece input_node_name, StringPiece function_name,
- const gtl::ArraySlice<const gtl::ArraySlice<int>>& output_shapes,
- const gtl::ArraySlice<DataType>& output_types) {
+NodeDef MakeMapNode(StringPiece name, StringPiece input_node_name,
+ StringPiece function_name,
+ gtl::ArraySlice<PartialTensorShape> output_shapes,
+ gtl::ArraySlice<DataType> output_types) {
return MakeMapNodeHelper(name, input_node_name, function_name, "MapDataset",
output_shapes, output_types);
}
-NodeDef MakeBatchNode(
- StringPiece name, StringPiece input_node_name,
- StringPiece input_batch_size_name,
- const gtl::ArraySlice<const gtl::ArraySlice<int>>& output_shapes,
- const gtl::ArraySlice<DataType>& output_types) {
- return NDef(name, "BatchDataset",
- {string(input_node_name), string(input_batch_size_name)},
- {{"output_types", output_types},
- {"output_shapes", MakeShapeListAttr(output_shapes)}});
+NodeDef MakeBatchNode(StringPiece name, StringPiece input_node_name,
+ StringPiece input_batch_size_name,
+ gtl::ArraySlice<PartialTensorShape> output_shapes,
+ gtl::ArraySlice<DataType> output_types) {
+ return NDef(
+ name, "BatchDataset",
+ {string(input_node_name), string(input_batch_size_name)},
+ {{"output_types", output_types}, {"output_shapes", output_shapes}});
}
-NodeDef MakeBatchV2Node(
- StringPiece name, StringPiece input_node_name,
- StringPiece input_batch_size_name, StringPiece input_drop_remainder_name,
- const gtl::ArraySlice<const gtl::ArraySlice<int>>& output_shapes,
- const gtl::ArraySlice<DataType>& output_types) {
- return NDef(name, "BatchDatasetV2",
- {string(input_node_name), string(input_batch_size_name),
- string(input_drop_remainder_name)},
- {{"output_types", output_types},
- {"output_shapes", MakeShapeListAttr(output_shapes)}});
+NodeDef MakeBatchV2Node(StringPiece name, StringPiece input_node_name,
+ StringPiece input_batch_size_name,
+ StringPiece input_drop_remainder_name,
+ gtl::ArraySlice<PartialTensorShape> output_shapes,
+ gtl::ArraySlice<DataType> output_types) {
+ return NDef(
+ name, "BatchDatasetV2",
+ {string(input_node_name), string(input_batch_size_name),
+ string(input_drop_remainder_name)},
+ {{"output_types", output_types}, {"output_shapes", output_shapes}});
}
-NodeDef MakeRangeNode(StringPiece name, const gtl::ArraySlice<string>& inputs) {
+NodeDef MakeRangeNode(StringPiece name, gtl::ArraySlice<string> inputs) {
return NDef(name, "RangeDataset", inputs,
- {{"output_shapes", MakeShapeListAttr({{}})},
+ {{"output_shapes", gtl::ArraySlice<TensorShape>({{}})},
{"output_types", gtl::ArraySlice<DataType>({DT_INT64})}});
}
@@ -184,7 +163,7 @@
item.graph = GDef(
{NDef("batch_size", "Const", {}, {{"value", 1}, {"dtype", DT_INT32}}),
NDef("input", "InputDataset", {},
- {{"output_shapes", MakeShapeListAttr({{}})}}),
+ {{"output_shapes", gtl::ArraySlice<TensorShape>({{}})}}),
MakeMapNode("map", "input", "XTimesTwo", {{}}, {DT_INT32}),
MakeBatchNode("batch", "map", "batch_size", {{-1}}, {DT_INT32})},
// FunctionLib
@@ -196,6 +175,37 @@
TF_ASSERT_OK(optimizer.Optimize(nullptr, item, &output));
}
+TEST(MapVectorizationTest, VectorizeWithFullyDefinedFunction) {
+ GrapplerItem item;
+ item.graph = GDef(
+ {NDef("start", "Const", {}, {{"value", 0}, {"dtype", DT_INT32}}),
+ NDef("stop", "Const", {}, {{"value", 10}, {"dtype", DT_INT32}}),
+ NDef("step", "Const", {}, {{"value", 1}, {"dtype", DT_INT32}}),
+ NDef("batch_size", "Const", {}, {{"value", 1}, {"dtype", DT_INT32}}),
+ MakeRangeNode("range", {"start", "stop", "step"}),
+ MakeMapNode("map", "range", "Func", {{}}, {DT_INT32}),
+ MakeBatchNode("batch", "map", "batch_size", {{-1}}, {DT_INT32})},
+ // FunctionLib
+ {FunctionDefHelper::Create(
+ "Func", {"x: int64", "y: int64"}, {"res: int64", "res2: int64"}, {},
+ {{{"o"}, "Mul", {"x", "x"}, {{"T", DT_INT64}}}},
+ {{"res", "o:z"}, {"res2", "o:z"}})});
+ MapVectorization optimizer;
+ GraphDef output;
+ TF_ASSERT_OK(optimizer.Optimize(nullptr, item, &output));
+
+ EXPECT_EQ(graph_utils::FindAllGraphNodesWithOp("MapDataset", output).size(),
+ 1);
+ EXPECT_EQ(graph_utils::FindAllGraphNodesWithOp("BatchDataset", output).size(),
+ 1);
+ const NodeDef& map_node =
+ output.node(graph_utils::FindGraphNodeWithOp("MapDataset", output));
+ const NodeDef& batch_node =
+ output.node(graph_utils::FindGraphNodeWithOp("BatchDataset", output));
+ EXPECT_EQ(map_node.input(0), batch_node.name());
+ EXPECT_EQ(batch_node.input(0), "range");
+}
+
} // namespace
} // namespace grappler
} // namespace tensorflow
diff --git a/tensorflow/core/grappler/optimizers/data/shuffle_and_repeat_fusion.cc b/tensorflow/core/grappler/optimizers/data/shuffle_and_repeat_fusion.cc
index cb0ff67..99c4afa 100644
--- a/tensorflow/core/grappler/optimizers/data/shuffle_and_repeat_fusion.cc
+++ b/tensorflow/core/grappler/optimizers/data/shuffle_and_repeat_fusion.cc
@@ -64,7 +64,7 @@
// Set `output_types` and `output_shapes` attributes.
for (auto key : {"output_shapes", "output_types"}) {
- (*new_node.mutable_attr())[key] = repeat_node.attr().at(key);
+ graph_utils::CopyAttribute(key, repeat_node, &new_node);
}
return new_node;
};
diff --git a/tensorflow/core/grappler/optimizers/data/vectorization/BUILD b/tensorflow/core/grappler/optimizers/data/vectorization/BUILD
index 1462cb2..37aa24b 100644
--- a/tensorflow/core/grappler/optimizers/data/vectorization/BUILD
+++ b/tensorflow/core/grappler/optimizers/data/vectorization/BUILD
@@ -9,13 +9,14 @@
VECTORIZER_DEPS = [
":vectorizer_registry",
- "//tensorflow/core/grappler/optimizers/data:function_utils",
+ "//tensorflow/core/grappler/optimizers/data:graph_utils",
] + tf_protos_all()
cc_library(
name = "vectorizer",
hdrs = ["vectorizer.h"],
deps = [
+ "//tensorflow/core:core_cpu",
"//tensorflow/core:lib",
] + tf_protos_all(),
)
diff --git a/tensorflow/core/grappler/optimizers/data/vectorization/cast_vectorizer.cc b/tensorflow/core/grappler/optimizers/data/vectorization/cast_vectorizer.cc
index c173973..3af6bab 100644
--- a/tensorflow/core/grappler/optimizers/data/vectorization/cast_vectorizer.cc
+++ b/tensorflow/core/grappler/optimizers/data/vectorization/cast_vectorizer.cc
@@ -14,7 +14,7 @@
==============================================================================*/
#include "tensorflow/core/framework/node_def.pb.h"
-#include "tensorflow/core/grappler/optimizers/data/function_utils.h"
+#include "tensorflow/core/grappler/optimizers/data/graph_utils.h"
#include "tensorflow/core/grappler/optimizers/data/vectorization/vectorizer_registry.h"
namespace tensorflow {
@@ -23,26 +23,21 @@
class CastVectorizer : public Vectorizer {
public:
- Status Vectorize(const NodeDef& node, gtl::ArraySlice<string> inputs,
- FunctionDef* outer_scope,
- std::map<string, string>* conversion_map) override {
- if (inputs.size() != 1) {
+ Status Vectorize(const Node& node, Graph* outer_scope,
+ std::vector<Port>* input_ports,
+ std::vector<Port>* output_ports) override {
+ Status s;
+ if (node.num_inputs() != 1) {
return errors::Internal("Cast op should only have one input.");
}
- // Add new Cast node
- NodeDef* new_cast_node = outer_scope->add_node_def();
- *new_cast_node = node;
- new_cast_node->clear_name();
- function_utils::SetUniqueFunctionNodeName(
- strings::StrCat("vectorized/", node.name()), outer_scope,
- new_cast_node);
- new_cast_node->set_input(0, inputs[0]);
+ // Add new Cast node with the same op and attrs as the original node
+ auto new_cast_node = outer_scope->AddNode(node.def(), &s);
+ TF_RETURN_IF_ERROR(s);
- // Add the output mapping to conversion map
- (*conversion_map)[strings::StrCat(node.name(), ":y:0")] =
- strings::StrCat(new_cast_node->name(), ":y:0");
-
+ // Add input and output mappings
+ input_ports->push_back({new_cast_node, 0});
+ output_ports->push_back({new_cast_node, 0});
return Status::OK();
}
};
diff --git a/tensorflow/core/grappler/optimizers/data/vectorization/unpack_vectorizer.cc b/tensorflow/core/grappler/optimizers/data/vectorization/unpack_vectorizer.cc
index 776d317..74ce520 100644
--- a/tensorflow/core/grappler/optimizers/data/vectorization/unpack_vectorizer.cc
+++ b/tensorflow/core/grappler/optimizers/data/vectorization/unpack_vectorizer.cc
@@ -14,7 +14,7 @@
==============================================================================*/
#include "tensorflow/core/framework/node_def.pb.h"
-#include "tensorflow/core/grappler/optimizers/data/function_utils.h"
+#include "tensorflow/core/grappler/optimizers/data/graph_utils.h"
#include "tensorflow/core/grappler/optimizers/data/vectorization/vectorizer_registry.h"
namespace tensorflow {
@@ -23,31 +23,29 @@
class UnpackVectorizer : public Vectorizer {
public:
- Status Vectorize(const NodeDef& node, gtl::ArraySlice<string> inputs,
- FunctionDef* outer_scope,
- std::map<string, string>* conversion_map) override {
- if (inputs.size() != 1) {
+ Status Vectorize(const Node& node, Graph* outer_scope,
+ std::vector<Port>* input_ports,
+ std::vector<Port>* output_ports) override {
+ Status s;
+ if (node.num_inputs() != 1) {
return errors::Internal("Unpack op should only have one input.");
}
- // Add new Unpack node
- NodeDef* new_unpack_node = outer_scope->add_node_def();
- *new_unpack_node = node;
- new_unpack_node->clear_name();
- function_utils::SetUniqueFunctionNodeName(
- strings::StrCat("vectorized/", node.name()), outer_scope,
- new_unpack_node);
+ // Add new Unpack node with the same op and attrs as the original node
+ auto new_unpack_node = outer_scope->AddNode(node.def(), &s);
+ TF_RETURN_IF_ERROR(s);
// Increment "axis" attr by 1:
- (*new_unpack_node->mutable_attr())["axis"].set_i(
- node.attr().at("axis").i() + 1);
- new_unpack_node->set_input(0, inputs[0]);
+ int new_axis = node.def().attr().at("axis").i() + 1;
+ new_unpack_node->AddAttr("axis", new_axis);
- // Add the output mappings to conversion map
- int num = new_unpack_node->attr().at("num").i();
+ // Add the input mappings
+ input_ports->push_back({new_unpack_node, 0});
+
+ // Add the output mappings
+ int num = node.def().attr().at("num").i();
for (int i = 0; i < num; ++i) {
- (*conversion_map)[strings::StrCat(node.name(), ":output:", i)] =
- strings::StrCat(new_unpack_node->name(), ":output:", i);
+ output_ports->push_back({new_unpack_node, i});
}
return Status::OK();
diff --git a/tensorflow/core/grappler/optimizers/data/vectorization/vectorizer.h b/tensorflow/core/grappler/optimizers/data/vectorization/vectorizer.h
index d341dbb..56eb88c 100644
--- a/tensorflow/core/grappler/optimizers/data/vectorization/vectorizer.h
+++ b/tensorflow/core/grappler/optimizers/data/vectorization/vectorizer.h
@@ -17,30 +17,33 @@
#define TENSORFLOW_CORE_GRAPPLER_OPTIMIZERS_DATA_VECTORIZATION_VECTORIZER_H_
#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/status.h"
-#include "tensorflow/core/lib/gtl/array_slice.h"
namespace tensorflow {
namespace grappler {
namespace vectorization_utils {
+// Describes a tensor with its operation Node and output position
+typedef std::pair<Node*, int> Port;
+
// Interface for vectorization of TensorFlow operations. See `CastVectorizer`
// for an example.
class Vectorizer {
public:
virtual ~Vectorizer() {}
- // Vectorizes an operation, `node`, by adding operation(s) to `outer_scope`
+ // Vectorizes an operation, `node`, by adding Node(s) to `outer_scope`
// that produce the same vector output(s) as executing `node`'s op
- // on elements of the vector inputs, and adding mappings to `conversion_map`
- // from old output tensor names to new (vectorized) output tensor names.
- // The new node(s) collectively have the same number of inputs and outputs as
- // the node being converted, and use the tensor names in `inputs` as their
- // inputs.
- virtual Status Vectorize(const NodeDef& node, gtl::ArraySlice<string> inputs,
- FunctionDef* outer_scope,
- std::map<string, string>* conversion_map) = 0;
+ // on elements of the vector inputs. The new Node(s) collectively have the
+ // same number of input and output ports as the node being converted.
+ // Adds mappings for the new nodes' input and output ports to `inputs` and
+ // `outputs` respectively, where the i'th Port in inputs/outputs
+ // corresponds to the i'th input/output port of the node to be converted.
+ virtual Status Vectorize(const Node& node, Graph* outer_scope,
+ std::vector<Port>* input_ports,
+ std::vector<Port>* output_ports) = 0;
};
} // namespace vectorization_utils
diff --git a/tensorflow/core/grappler/optimizers/data/vectorization/vectorizer_registry_test.cc b/tensorflow/core/grappler/optimizers/data/vectorization/vectorizer_registry_test.cc
index 86e3035..663ceba 100644
--- a/tensorflow/core/grappler/optimizers/data/vectorization/vectorizer_registry_test.cc
+++ b/tensorflow/core/grappler/optimizers/data/vectorization/vectorizer_registry_test.cc
@@ -24,9 +24,9 @@
class TestVectorizer : public Vectorizer {
public:
- Status Vectorize(const NodeDef& node, gtl::ArraySlice<string> inputs,
- FunctionDef* outer_scope,
- std::map<string, string>* conversion_map) override {
+ Status Vectorize(const Node& node, Graph* outer_scope,
+ std::vector<Port>* inputs,
+ std::vector<Port>* outputs) override {
return Status::OK();
}
};
@@ -39,10 +39,12 @@
auto vectorizer = VectorizerRegistry::Global()->Get("test_op");
EXPECT_NE(vectorizer, nullptr);
- FunctionDef function;
- NodeDef node;
- std::map<string, string> conversion_map;
- EXPECT_TRUE(vectorizer->Vectorize(node, {}, &function, &conversion_map).ok());
+ Graph g(OpRegistry::Global());
+ NodeDef node_def;
+ Status s;
+ Node* node = g.AddNode(node_def, &s);
+ std::vector<Port> inputs, outputs;
+ EXPECT_TRUE(vectorizer->Vectorize(*node, &g, &inputs, &outputs).ok());
}
} // namespace vectorization_utils
diff --git a/tensorflow/core/grappler/optimizers/data/vectorization_utils.cc b/tensorflow/core/grappler/optimizers/data/vectorization_utils.cc
index cb56b65..cea667f 100644
--- a/tensorflow/core/grappler/optimizers/data/vectorization_utils.cc
+++ b/tensorflow/core/grappler/optimizers/data/vectorization_utils.cc
@@ -14,13 +14,17 @@
==============================================================================*/
#include "tensorflow/core/grappler/optimizers/data/vectorization_utils.h"
+#include <memory>
#include "tensorflow/core/grappler/optimizers/data/vectorization/vectorizer_registry.h"
#include "absl/strings/str_join.h"
+#include "tensorflow/core/common_runtime/function.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/device_base.h"
#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/framework/graph_to_functiondef.h"
#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/core/framework/types.h"
@@ -36,255 +40,346 @@
namespace grappler {
namespace vectorization_utils {
-using function_utils::FunctionDefTensorDesc;
-
namespace {
-void AddMapDefunOutput(FunctionDef* map_defun_fn, NodeDef* map_defun_node,
- const string& output_retval, const DataType t) {
- // Set to unknown shape
- TensorShapeProto tensor_shape_proto;
- PartialTensorShape().AsProto(&tensor_shape_proto);
+// Describes a tensor with its operation Node and output position
+typedef std::pair<Node*, int> TensorDesc;
- function_utils::AddFunctionOutputWithUniqueName(
- "vectorized_out", output_retval, map_defun_fn, t);
+const char* const kRetValOp = "_Retval";
- *(*map_defun_node->mutable_attr())["output_shapes"]
- .mutable_list()
- ->add_shape() = tensor_shape_proto;
- (*map_defun_node->mutable_attr())["output_types"].mutable_list()->add_type(t);
+void ReplaceEdgeSources(const TensorDesc& old_src, const TensorDesc& new_src,
+ Graph* graph) {
+ // NOTE: We need two for loops here because we can't mutate the set of output
+ // edges as we iterate over them.
+ std::vector<const Edge*> edges_to_replace;
+ for (auto edge : old_src.first->out_edges()) {
+ if (edge->src_output() == old_src.second) {
+ edges_to_replace.push_back(edge);
+ }
+ }
+ for (auto edge : edges_to_replace) {
+ graph->AddEdge(new_src.first, new_src.second, edge->dst(),
+ edge->dst_input());
+ graph->RemoveEdge(edge);
+ }
}
-void RemoveMapDefunOutput(FunctionDef* outer_scope, FunctionDef* map_defun_fn,
- NodeDef* map_defun_node, int output_position) {
- DCHECK_LT(output_position, map_defun_fn->signature().output_arg_size())
+Status AddMapDefunOutput(FunctionBody* map_defun_fn, Node* map_defun_node,
+ const TensorDesc& output) {
+ // Note that we don't update MapDefun attrs as we go, only when we are done
+ DataType type = output.first->output_type(output.second);
+ int index = map_defun_fn->ret_nodes.size();
+
+ NodeDef ret_node_def;
+ ret_node_def.set_name("map_out");
+ ret_node_def.set_op(kRetValOp);
+ AddNodeAttr("T", type, &ret_node_def);
+ AddNodeAttr("index", index, &ret_node_def);
+
+ Status s;
+ Node* ret_node = map_defun_fn->graph->AddNode(ret_node_def, &s);
+ TF_RETURN_IF_ERROR(s);
+
+ map_defun_fn->graph->AddEdge(output.first, output.second, ret_node, 0);
+ map_defun_fn->ret_nodes.push_back(ret_node);
+ map_defun_fn->ret_types.push_back(type);
+
+ return s;
+}
+
+void RemoveMapDefunOutput(int output_position, Graph* outer_scope,
+ FunctionBody* map_defun_fn, Node* map_defun_node) {
+ // Note that we don't update MapDefun attrs as we go, only when we are done
+ DCHECK_LT(output_position, map_defun_fn->ret_nodes.size())
<< "Trying to remove output that doesn't exist. Output number: "
<< output_position;
- int num_later_outputs =
- map_defun_fn->signature().output_arg_size() - output_position - 1;
+ int num_later_outputs = map_defun_fn->ret_nodes.size() - output_position - 1;
- // Remove from map_defun_fn's ret dict and output args
- map_defun_fn->mutable_ret()->erase(
- map_defun_fn->signature().output_arg(output_position).name());
- map_defun_fn->mutable_signature()->mutable_output_arg()->DeleteSubrange(
- output_position, 1);
+ // Modify map_defun_fn's signature and remove the output node from its graph
+ map_defun_fn->graph->RemoveNode(map_defun_fn->ret_nodes[output_position]);
+ map_defun_fn->ret_nodes.erase(map_defun_fn->ret_nodes.begin() +
+ output_position);
+ map_defun_fn->ret_types.erase(map_defun_fn->ret_types.begin() +
+ output_position);
- // Renumber outputs that come after
+ // Renumber the nodes and edges that come after
for (int i = 0; i < num_later_outputs; ++i) {
- function_utils::ReplaceReferences(
- strings::StrCat(map_defun_node->name(),
- ":output:", output_position + i + 1),
- strings::StrCat(map_defun_node->name(),
- ":output:", output_position + i),
- outer_scope);
+ ReplaceEdgeSources({map_defun_node, output_position + i + 1},
+ {map_defun_node, output_position + i}, outer_scope);
+ // Each ret node has an "index" attr that has to be updated
+ map_defun_fn->ret_nodes[output_position + i]->AddAttr("index",
+ output_position + i);
}
- map_defun_node->mutable_attr()
- ->at("output_shapes")
- .mutable_list()
- ->mutable_shape()
- ->DeleteSubrange(output_position, 1);
- map_defun_node->mutable_attr()
- ->at("output_types")
- .mutable_list()
- ->mutable_type()
- ->ExtractSubrange(output_position, 1, nullptr);
-}
-
-int FindOutputToConvert(const FunctionDef& function,
- const std::set<string>& unconvertible,
- FunctionDefTensorDesc* f) {
- for (int i = function.signature().output_arg_size() - 1; i >= 0; --i) {
- const string& ret_key = function.signature().output_arg(i).name();
- *f = FunctionDefTensorDesc(function.ret().at(ret_key));
-
- if (unconvertible.find(f->node_name) == unconvertible.end()) {
- return i;
- }
- }
- return -1;
}
// Helper class that vectorizes the body of a MapDefun node, adding new
// operations to the graph that collectively compute the same value as what
// running the MapDefun function on slices of the input would produce.
-// Each instance of the class encapsulates all the data necessary to vectorize a
-// MapDefun op in place.
+// This class transforms the input FunctionDefs into their corresponding
+// Graph objects and works on the graphs directly, then converts them back
+// to FunctionDefs when GetResult is called.
class Vectorization {
public:
- Vectorization(FunctionDef* outer_scope, FunctionDef* map_defun_fn,
- NodeDef* map_defun_node)
- : outer_scope_(outer_scope),
- map_defun_fn_(map_defun_fn),
- map_defun_node_(map_defun_node) {}
+ explicit Vectorization(FunctionDefLibrary* lib)
+ : lib_(lib), lib_def_(OpRegistry::Global(), *lib) {}
- // Repeatedly tries to convert outputs of map_defun_fn_ into new nodes in
- // the outer_scope_, until there are no convertible outputs remaining.
- // This method is idempotent.
- void Vectorize();
+ // Adds the vectorized function and new map_defun_fn to lib, and points
+ // vectorized_function to the former. Returns an error status if
+ // the conversion between FunctionDef -> Graph -> FunctionDef failed anywhere
+ // along the way.
+ Status Vectorize(const FunctionDef& outer_scope,
+ const NodeDef& map_defun_node, FunctionDef** result);
private:
- // Vectorizes the map defun function's output at output_position
- Status ConvertOutput(int output_position, const FunctionDefTensorDesc& desc);
- // Given a descriptor of the original output tensor, gets a string
- // corresponding to the converted output tensor.
- Status ConvertOutputHelper(const FunctionDefTensorDesc& output_desc,
- string* converted);
- Status AddConversionMappingFromInput(
- const FunctionDefTensorDesc& output_desc);
+ // Converts FunctionDefs to Graphs.
+ Status Initialize(const FunctionDef& outer_scope,
+ const NodeDef& map_defun_node);
+
+ // Converts Graphs back to FunctionDefs and adds them to `lib_`.
+ Status GetResult(FunctionDef** vectorized_function);
+
+ // Repeatedly tries to convert outputs of `map_defun_fn_` into new nodes in
+ // `outer_scope_`, until there are no convertible outputs remaining.
+ void VectorizeHelper();
+
+ // Vectorizes map_defun_fn's output at output_position.
+ Status ConvertOutput(int output_position);
// Adds mappings from node's outputs tensors to converted output tensors,
// creating the necessary new node(s). Generally, the steps to convert an op
// are:
- // 1) Promote the inputs of the op inputs to outputs of the map_defun_fn_,
- // and modify map_defun_node_ attrs accordingly
- // 2) Create new node(s) in outer_scope_ that act on batched input tensors.
+ // 1) Create new node(s) in `outer_scope_` that act on batched input tensors.
// These operations collectively compute the same value as what running
// the original operation on slices of the input tensors would produce.
// For example, a Cast op in MapDefun translates to a Cast op in
- // outer_scope_, since the vectorized version of Cast is itself.
- // 3) Set inputs of new node(s) to the corresponding converted inputs (that
- // are now outputs of map_defun_node_)
- // 4) For each output of the old node, add the mapping of output strings to
- // the conversion map (eg "Cast:y:0" -> "Vectorize/Cast:y:0")
- Status AddConversionMappingFromOp(const NodeDef& node,
- const FunctionDefTensorDesc& output_desc);
+ // `outer_scope_`, since the vectorized version of Cast is itself.
+ // 2) Promote the inputs of the op inputs to outputs of the
+ // `map_defun_node_` and `map_defun_fn_`.
+ // 3) Add edges between the promoted inputs (that are now outputs of
+ // `map_defun_node`) and the inputs ports of the new node(s).
+ // 4) For each output of the old node, add the mapping of output tensors to
+ // the conversion map.
+ Status AddConversionMapping(Node* op_node);
- // Maps a tensor name to the name of the corresponding vectorized tensor. For
- // example, "Cast:y:0" -> "Vectorize/Cast:y:0"
- std::map<string, string> conversion_map_;
- // Unconvertible node names
- std::set<string> unconvertible_;
+ // Maps a tensor to the corresponding vectorized tensor. For example,
+ // {"Cast" Node*, 0} -> {"Vectorize/Cast" Node*, 0}
+ std::map<TensorDesc, TensorDesc> conversion_map_;
- FunctionDef* outer_scope_;
- FunctionDef* map_defun_fn_;
- NodeDef* map_defun_node_;
+ // Unconvertible ret nodes
+ std::set<Node*> unconvertible_;
+
+ FunctionDefLibrary* lib_; // Not owned
+ FunctionLibraryDefinition lib_def_;
+ // Note that FunctionBody has a pointer to a Graph object that corresponds
+ // to the function's subgraph, with additional kArgOp and kRetValOp nodes
+ // that denote that function arguments and return values. These nodes have the
+ // attrs "T" for the type, and "index" for the argument / retval index
+ // respectively. FunctionBody also keeps track of arg/ret_nodes and
+ // arg/ret_types, that should be ordered according to argument/output indices.
+ std::unique_ptr<Graph> outer_scope_;
+ std::unique_ptr<FunctionBody> map_defun_fn_;
+ Node* map_defun_node_ = nullptr; // Owned by `outer_scope`
+ Status status_;
};
-Status Vectorization::AddConversionMappingFromOp(
- const NodeDef& node, const FunctionDefTensorDesc& output_desc) {
- for (const string& input_name : node.input()) {
- if (IsControlInput(input_name)) {
+Status Vectorization::AddConversionMapping(Node* op_node) {
+ for (auto edge : op_node->in_edges()) {
+ if (edge->IsControlEdge()) {
return errors::InvalidArgument(
"Vectorizing outputs with control inputs is currently not "
"supported.");
}
}
- // TODO(rachelim): Have some mechanism for registering converters and some
- // uniform, simpler way to represent them.
-
- DataTypeVector types;
- const OpDef* op_def = nullptr;
- TF_RETURN_IF_ERROR(OpRegistry::Global()->LookUpOpDef(node.op(), &op_def));
- TF_RETURN_IF_ERROR(InputTypesForNode(node, *op_def, &types));
-
- std::vector<string> promoted_inputs;
- promoted_inputs.reserve(node.input_size());
- for (int i = 0; i < node.input_size(); ++i) {
- promoted_inputs.push_back(strings::StrCat(
- map_defun_node_->name(),
- ":output:", map_defun_fn_->signature().output_arg_size() + i));
- }
-
- auto vectorizer = VectorizerRegistry::Global()->Get(node.op());
+ auto vectorizer = VectorizerRegistry::Global()->Get(op_node->type_string());
if (vectorizer == nullptr) {
return errors::Unimplemented("No vectorizer registered for op: ",
- node.op());
+ op_node->type_string());
+ }
+ std::vector<Port> input_ports, output_ports;
+ input_ports.reserve(op_node->num_inputs());
+ output_ports.reserve(op_node->num_outputs());
+ TF_RETURN_IF_ERROR(vectorizer->Vectorize(*op_node, outer_scope_.get(),
+ &input_ports, &output_ports));
+
+ std::vector<const Edge*> input_edges;
+ TF_RETURN_IF_ERROR(op_node->input_edges(&input_edges));
+
+ if (op_node->num_outputs() != output_ports.size() ||
+ op_node->num_inputs() != input_ports.size() ||
+ input_edges.size() != input_ports.size()) {
+ return errors::Internal("Vectorizer inputs/outputs don't match.");
}
- TF_RETURN_IF_ERROR(vectorizer->Vectorize(node, promoted_inputs, outer_scope_,
- &conversion_map_));
+ // Promote the inputs of the op to MapDefun outputs and connect the edges
+ // accordingly.
+ for (size_t i = 0; i < op_node->num_inputs(); ++i) {
+ auto edge = input_edges[i];
+ TF_RETURN_IF_ERROR(AddMapDefunOutput(map_defun_fn_.get(), map_defun_node_,
+ {edge->src(), edge->src_output()}));
+ outer_scope_->AddEdge(map_defun_node_, map_defun_fn_->ret_nodes.size() - 1,
+ input_ports[i].first, input_ports[i].second);
+ }
- // If we get here, the conversion was successful, so we promote the inputs
- // of the ops to MapDefun outputs.
- for (int i = 0; i < types.size(); ++i) {
- AddMapDefunOutput(map_defun_fn_, map_defun_node_, node.input(i), types[i]);
+ // Add output mappings.
+ for (size_t i = 0; i < op_node->num_outputs(); ++i) {
+ conversion_map_.insert({{op_node, i}, std::move(output_ports[i])});
}
return Status::OK();
}
-Status Vectorization::AddConversionMappingFromInput(
- const FunctionDefTensorDesc& output_desc) {
- int input_index = function_utils::FindFunctionInputWithName(
- output_desc.node_name, *map_defun_fn_);
- if (input_index == -1) {
- return errors::Internal("Cannot convert non-existent input.");
- }
+Status Vectorization::ConvertOutput(int output_position) {
+ // ret_edge->src() is the actual op that generated the retval, and
+ // ret_edge->dst() is the retval node whose op is "_Retval"
+ const Edge* ret_edge;
+ TF_RETURN_IF_ERROR(
+ map_defun_fn_->ret_nodes[output_position]->input_edge(0, &ret_edge));
- conversion_map_[output_desc.full_str] = map_defun_node_->input(input_index);
- return Status::OK();
-}
-
-Status Vectorization::ConvertOutputHelper(
- const FunctionDefTensorDesc& output_desc, string* converted) {
- // It's possible the output already has a mapping, if it comes from a node
- // that has already been converted.
- if (auto found = gtl::FindOrNull(conversion_map_, output_desc.full_str)) {
- *converted = *found;
- return Status::OK();
- }
-
- int index = function_utils::FindFunctionNodeWithName(output_desc.node_name,
- *map_defun_fn_);
- if (index == -1) { // The output comes from an input
- TF_RETURN_IF_ERROR(AddConversionMappingFromInput(output_desc));
+ TensorDesc output({ret_edge->src(), ret_edge->src_output()});
+ TensorDesc converted_output;
+ if (auto found = gtl::FindOrNull(conversion_map_, output)) {
+ // It's possible the output already has a mapping, if it comes from a node
+ // that has already been converted.
+ converted_output = *found;
} else {
- TF_RETURN_IF_ERROR(AddConversionMappingFromOp(
- map_defun_fn_->node_def(index), output_desc));
+ TF_RETURN_IF_ERROR(AddConversionMapping(output.first));
+ converted_output = conversion_map_.at(output);
}
- *converted = conversion_map_.at(output_desc.full_str);
- return Status::OK();
-}
-Status Vectorization::ConvertOutput(int output_position,
- const FunctionDefTensorDesc& output_desc) {
- string converted_output_name;
- TF_RETURN_IF_ERROR(ConvertOutputHelper(output_desc, &converted_output_name));
-
- // Remove the old output and make everything that referenced it point
- // to the new string
- function_utils::ReplaceReferences(
- strings::StrCat(map_defun_node_->name(), ":output:", output_position),
- converted_output_name, outer_scope_);
- RemoveMapDefunOutput(outer_scope_, map_defun_fn_, map_defun_node_,
- output_position);
+ ReplaceEdgeSources({map_defun_node_, output_position}, converted_output,
+ outer_scope_.get());
+ RemoveMapDefunOutput(output_position, outer_scope_.get(), map_defun_fn_.get(),
+ map_defun_node_);
return Status::OK();
}
-void Vectorization::Vectorize() {
+Status Vectorization::Vectorize(const FunctionDef& outer_scope,
+ const NodeDef& map_defun_node,
+ FunctionDef** result) {
+ TF_RETURN_IF_ERROR(Initialize(outer_scope, map_defun_node));
+ VectorizeHelper();
+ return GetResult(result);
+}
+
+void Vectorization::VectorizeHelper() {
while (true) {
- FunctionDefTensorDesc desc;
- int output_position =
- FindOutputToConvert(*map_defun_fn_, unconvertible_, &desc);
+ int output_position = graph_utils::GetFirstElementIndexWithPredicate(
+ [this](Node* n) {
+ return this->unconvertible_.find(n) == this->unconvertible_.end();
+ },
+ map_defun_fn_->ret_nodes);
+
+ // No outputs left to convert
if (output_position == -1) break;
- if (!ConvertOutput(output_position, desc).ok()) {
- unconvertible_.insert(desc.node_name);
+ Status s = ConvertOutput(output_position);
+ if (!s.ok()) {
+ Node* output_node = map_defun_fn_->ret_nodes.at(output_position);
+ VLOG(2) << "Could not convert the output at node: "
+ << output_node->DebugString() << "\nError: " << s;
+ unconvertible_.insert(output_node);
}
}
// If we've converted all the outputs of the MapDefun function, we no longer
// need the MapDefun node and can delete it.
- if (map_defun_fn_->signature().output_arg_size() == 0) {
- outer_scope_->mutable_node_def()->DeleteSubrange(
- function_utils::FindFunctionNodeWithName(map_defun_node_->name(),
- *outer_scope_),
- 1);
- }
-
- if (!unconvertible_.empty()) {
- VLOG(2) << "The following nodes could not be converted: ["
- << absl::StrJoin(unconvertible_, ", ") << "].";
+ if (map_defun_fn_->ret_nodes.empty()) {
+ outer_scope_->RemoveNode(map_defun_node_);
+ } else {
+ // Update MapDefun node attrs accordingly
+ DCHECK_EQ(map_defun_fn_->ret_types.size(), map_defun_fn_->ret_nodes.size());
+ map_defun_node_->AddAttr(
+ "output_shapes",
+ std::vector<PartialTensorShape>(map_defun_fn_->ret_types.size()));
+ map_defun_node_->AddAttr("output_types", map_defun_fn_->ret_types);
}
}
+Status Vectorization::Initialize(const FunctionDef& outer_scope,
+ const NodeDef& map_defun_node) {
+ // Convert outer_scope and map_defun_fn to FunctionBodys so we can
+ // work on Graphs directly.
+ const FunctionDef* map_defun_fn =
+ lib_def_.Find(map_defun_node.attr().at("f").func().name());
+
+ if (map_defun_fn == nullptr) {
+ return errors::NotFound("Could not find function with name ",
+ map_defun_node.attr().at("f").func().name(),
+ " in function library.");
+ }
+
+ auto get_func_sig = [this](const string& op, const OpDef** sig) {
+ return this->lib_def_.LookUpOpDef(op, sig);
+ };
+
+ FunctionBody* outer_fn;
+ TF_RETURN_IF_ERROR(FunctionDefToBodyHelper(outer_scope, {}, &lib_def_,
+ get_func_sig, &outer_fn));
+ // We don't need outer_fn, just the graph
+ outer_scope_.reset(outer_fn->graph);
+ outer_fn->graph = nullptr;
+ delete outer_fn;
+
+ FunctionBody* tmp;
+ TF_RETURN_IF_ERROR(FunctionDefToBodyHelper(*map_defun_fn, {}, &lib_def_,
+ get_func_sig, &tmp));
+ map_defun_fn_.reset(tmp);
+
+ // Find the MapDefun node in outer_scope_
+ int node_id = graph_utils::GetFirstElementIndexWithPredicate(
+ [&map_defun_node](Node* n) { return n->name() == map_defun_node.name(); },
+ outer_scope_->nodes());
+ if (node_id == -1) {
+ return errors::NotFound("Could not find node with name ",
+ map_defun_node.name(), " in outer_scope.");
+ }
+ map_defun_node_ = outer_scope_->FindNodeId(node_id);
+
+ // Add mappings from map_defun_fn_ arg nodes to map_defun_node_ input nodes to
+ // the conversion map
+ for (auto arg_node : map_defun_fn_->arg_nodes) {
+ Node* input_node;
+ TF_RETURN_IF_ERROR(map_defun_node_->input_node(
+ arg_node->attrs().Find("index")->i(), &input_node));
+
+ conversion_map_.insert({{arg_node, 0}, {input_node, 0}});
+ }
+
+ return Status::OK();
+}
+
+Status Vectorization::GetResult(FunctionDef** vectorized_function) {
+ TF_RETURN_IF_ERROR(status_);
+
+ if (!map_defun_fn_->ret_nodes.empty()) {
+ FunctionDef* map_defun_fn = lib_->add_function();
+ graph_utils::SetUniqueGraphFunctionName("map_defun_fn", lib_, map_defun_fn);
+ TF_RETURN_IF_ERROR(GraphToFunctionDef(
+ *map_defun_fn_->graph, map_defun_fn->signature().name(), map_defun_fn));
+
+ AttrValue func_attr;
+ func_attr.mutable_func()->set_name(map_defun_fn->signature().name());
+ map_defun_node_->AddAttr("f", func_attr);
+ }
+
+ *vectorized_function = lib_->add_function();
+ graph_utils::SetUniqueGraphFunctionName("vectorized_fn", lib_,
+ *vectorized_function);
+ TF_RETURN_IF_ERROR(GraphToFunctionDef(
+ *outer_scope_, (*vectorized_function)->signature().name(),
+ *vectorized_function));
+ return Status::OK();
+}
+
} // namespace
-void VectorizeMapDefun(FunctionDef* outer_scope, FunctionDef* map_defun_fn,
- NodeDef* map_defun_node) {
- Vectorization(outer_scope, map_defun_fn, map_defun_node).Vectorize();
+Status VectorizeMapDefun(const FunctionDef& outer_scope,
+ const NodeDef& map_defun_node, FunctionDefLibrary* lib,
+ FunctionDef** result) {
+ *result = nullptr;
+ return Vectorization(lib).Vectorize(outer_scope, map_defun_node, result);
}
} // end namespace vectorization_utils
diff --git a/tensorflow/core/grappler/optimizers/data/vectorization_utils.h b/tensorflow/core/grappler/optimizers/data/vectorization_utils.h
index bb405fa..bd7d390 100644
--- a/tensorflow/core/grappler/optimizers/data/vectorization_utils.h
+++ b/tensorflow/core/grappler/optimizers/data/vectorization_utils.h
@@ -24,22 +24,28 @@
namespace grappler {
namespace vectorization_utils {
-// Given a function, `map_defun_fn`, that is mapped across some input vector
-// elements via a MapDefun operation, `VectorizeMapDefun` attempts to
-// vectorize the MapDefun by "lifting" operations from the `map_defun_fn` to the
-// `outer_scope`; that is, replacing `map_defun_fn` operations with new
-// `outer_scope` operations that produce the same vector output(s) as executing
-// the `map_defun_fn` operations on elements of vector input(s) would. If all
-// `map_defun_fn` operations are successfully lifted, `map_defun_node` is
-// eliminated from `outer_scope` altogether. However, if some operations cannot
-// be lifted, and this vectorization only succeeds partially, `map_defun_node`
-// remains to be used for operations that were not lifted.
+// Given a MapDefun node (`map_defun_node`) in a FunctionDef (`outer_scope`)
+// that maps a function in lib across some input vector elements,
+// `VectorizeMapDefun` attempts to create a vectorized version of `outer_scope`
+// by "lifting" operations from the MapDefun function to the new function
+// (`result`); that is, replacing operations in the MapDefun function with
+// operations that produce the same vector output(s) as executing the original
+// operations on elements of vector input(s) would. If all operations in the
+// MapDefun function are successfully lifted, `result` has no MapDefun node
+// altogether. However, if some operations cannot be lifted, and this
+// vectorization only succeeds partially, a MapDefun node remains in `result` to
+// be used for operations that were not lifted, and the modified MapDefun
+// function is added to `lib`. The newly vectorized function `result` is also
+// added to `lib`.
+//
+// Returns Status::OK() if the vectorization is completely or partially
+// successful. Otherwise, returns an error, and sets `result` to nullptr.
//
// Example:
// If the input to the `VectorizeMapDefun` function is a MapDefun
// whose `map_defun_fn` performs the Cast operation, the vectorization will
// eliminate the MapDefun. This is because the Cast operation supports
-// any tensor shape and can thus be lifted to the `outer_scope`.
+// any tensor shape and can thus be lifted to `result`.
//
// Before:
//
@@ -68,7 +74,7 @@
//
// After:
//
-// outer_scope +------+
+// result +------+
// +---------------+ Arg0 +---------+
// | +---+--+ |
// | | |
@@ -80,8 +86,9 @@
// +---------------+ Ret0 +---------+
// +------+
//
-void VectorizeMapDefun(FunctionDef* outer_scope, FunctionDef* map_defun_fn,
- NodeDef* map_defun_node);
+Status VectorizeMapDefun(const FunctionDef& outer_scope,
+ const NodeDef& map_defun_node, FunctionDefLibrary* lib,
+ FunctionDef** result);
} // end namespace vectorization_utils
} // end namespace grappler
diff --git a/tensorflow/core/grappler/optimizers/data/vectorization_utils_test.cc b/tensorflow/core/grappler/optimizers/data/vectorization_utils_test.cc
index e129fa9..1ff6221 100644
--- a/tensorflow/core/grappler/optimizers/data/vectorization_utils_test.cc
+++ b/tensorflow/core/grappler/optimizers/data/vectorization_utils_test.cc
@@ -17,6 +17,7 @@
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/grappler/optimizers/data/function_utils.h"
+#include "tensorflow/core/grappler/optimizers/data/graph_utils.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
@@ -60,6 +61,11 @@
return node;
}
+string GetRetval(const FunctionDef& function_def, int index) {
+ return function_def.ret().at(
+ function_def.signature().output_arg(index).name());
+}
+
// TODO(rachelim): Use FunctionDefHelper::Create instead
FunctionDef CreateFunction(
StringPiece name, const std::vector<std::pair<string, DataType>>& inputs,
@@ -85,7 +91,6 @@
return func;
}
-TEST(FunctionDefInputDescTest, ConstructedCorrectly) {}
// Before:
//
@@ -133,10 +138,15 @@
{{}, {}}, inner.signature().name(), &outer);
CHECK_NOTNULL(map_defun);
- VectorizeMapDefun(&outer, &inner, map_defun);
- EXPECT_TRUE(!function_utils::ContainsFunctionNodeWithOp("MapDefun", outer));
- EXPECT_EQ(outer.ret().at("mapdefun"), "ret0");
- EXPECT_EQ(outer.ret().at("mapdefun_0"), "ret1");
+ FunctionDefLibrary lib;
+ *lib.add_function() = outer;
+ *lib.add_function() = inner;
+ FunctionDef* vectorized;
+ EXPECT_TRUE(VectorizeMapDefun(outer, *map_defun, &lib, &vectorized).ok());
+ EXPECT_TRUE(
+ !function_utils::ContainsFunctionNodeWithOp("MapDefun", *vectorized));
+ EXPECT_EQ(GetRetval(*vectorized, 0), "ret0");
+ EXPECT_EQ(GetRetval(*vectorized, 1), "ret1");
}
// Before:
@@ -149,12 +159,12 @@
// | +-----------+ Arg0 +---+ Arg1 +----+ |
// | | +---+--+ +---+--+ | |
// | | | | | |
-// | | +------+ | +---v--+ | |
-// | | |Const | | | Op0 | | |
-// | | +---v--+ | +---+--+ | |
+// | | +------+ | | | |
+// | | |Const | | | | |
+// | | +---v--+ | | | |
// | | | | | | |
// | | | +---v--+ +---v--+ | |
-// | | +---| XOp1 | | XOp2 | | |
+// | | +---| XOp1 | | Cast | | |
// | | +---+--+ +---+--+ | |
// | | | | | |
// | | MapDefun +---v--+ +---v--+ | |
@@ -165,23 +175,50 @@
// +---------------+ Ret0 +---+ Ret1 +--------+
// +------+ +------+
//
-// where XOp1 and XOp2 are not convertible.
+// where XOp1 is not convertible.
//
// After:
//
-// No change because the ops are not convertible.
+//
+// +------+ +------+
+// +---------------+ Arg0 +---+ Arg1 +--------+
+// | +---+--+ +---+--+ |
+// | | | |
+// | +---v--+ | |
+// | +-----------+ Arg0 +-+ | |
+// | | +---+--+ | | |
+// | | | | | |
+// | | +------+ | | | |
+// | | |Const | | | | |
+// | | +---v--+ | | | |
+// | | | | | | |
+// | | | +---v--+ | +---v--+ |
+// | | +---| XOp1 | | | Cast | |
+// | | +---+--+ | +---+--+ |
+// | | | | | |
+// | | MapDefun +---v--+ | | |
+// | +-----------+ Ret0 +-+ | |
+// | +---+--+ | |
+// | | | |
+// | +---v--+ +---v--+ |
+// +---------------+ Ret0 +---+ Ret1 +--------+
+// +------+ +------+
//
TEST(VectorizeMapDefunTest, VectorizeDefunUnconvertible) {
FunctionDef inner =
CreateFunction("inner_function", {{"arg0", DT_INT32}, {"arg1", DT_INT32}},
{{"ret0", DT_INT32}, {"ret1", DT_INT32}},
- {{"ret0", "XOp1:output:0"}, {"ret1", "XOp2:output:0"}});
+ {{"ret0", "MatMul:product:0"}, {"ret1", "Cast:y:0"}});
+ // TODO(rachelim): If we ever write a converter for MatMul, we have to
+ // change this test.
NodeDef* x_op1 =
- function_utils::AddNode("XOp1", "XOp1", {"const", "arg0"}, {}, &inner);
+ function_utils::AddNode("MatMul", "MatMul", {"arg0", "arg0"}, {}, &inner);
CHECK_NOTNULL(x_op1);
+ graph_transforms::SetNodeAttr("T", DT_INT32, x_op1);
- NodeDef* x_op2 = function_utils::AddNode("XOp2", "XOp2", {"op1"}, {}, &inner);
- CHECK_NOTNULL(x_op2);
+ NodeDef* cast_node =
+ AddCastNode("Cast", {"arg1"}, DT_INT32, DT_INT32, false, &inner);
+ CHECK_NOTNULL(cast_node);
FunctionDef outer = CreateFunction(
"outer_function", {{"x", DT_INT32}, {"y", DT_INT32}},
@@ -193,12 +230,22 @@
{{}, {}}, inner.signature().name(), &outer);
CHECK_NOTNULL(map_defun);
- FunctionDef outer_copy(outer);
- FunctionDef inner_copy(inner);
- VectorizeMapDefun(&outer, &inner, map_defun);
- // They should be unchanged
- EXPECT_TRUE(FunctionDefsEqual(outer_copy, outer));
- EXPECT_TRUE(FunctionDefsEqual(inner_copy, inner));
+ FunctionDefLibrary lib;
+ *lib.add_function() = outer;
+ *lib.add_function() = inner;
+ FunctionDef* vectorized;
+ EXPECT_TRUE(VectorizeMapDefun(outer, *map_defun, &lib, &vectorized).ok());
+
+ auto map_defun_node = vectorized->node_def(
+ function_utils::FindFunctionNodeWithOp("MapDefun", *vectorized));
+ // The Cast node should be converted just fine.
+ EXPECT_EQ(GetRetval(*vectorized, 1), "Cast:y:0");
+
+ // The inner function should only have one retval.
+ FunctionLibraryDefinition lib_def(OpRegistry::Global(), lib);
+ const FunctionDef* map_defun_fn =
+ lib_def.Find(map_defun_node.attr().at("f").func().name());
+ EXPECT_EQ(map_defun_fn->signature().output_arg_size(), 1);
}
// Before:
@@ -257,14 +304,19 @@
inner.signature().name(), &outer);
CHECK_NOTNULL(map_defun);
- VectorizeMapDefun(&outer, &inner, map_defun);
- EXPECT_TRUE(!function_utils::ContainsFunctionNodeWithOp("MapDefun", outer));
- const NodeDef& cast_node =
- outer.node_def(function_utils::FindFunctionNodeWithOp("Cast", outer));
+ FunctionDefLibrary lib;
+ *lib.add_function() = outer;
+ *lib.add_function() = inner;
+ FunctionDef* vectorized;
+ EXPECT_TRUE(VectorizeMapDefun(outer, *map_defun, &lib, &vectorized).ok());
+ EXPECT_TRUE(
+ !function_utils::ContainsFunctionNodeWithOp("MapDefun", *vectorized));
+ const NodeDef& cast_node = vectorized->node_def(
+ function_utils::FindFunctionNodeWithOp("Cast", *vectorized));
EXPECT_EQ(cast_node.input(0), "x");
- EXPECT_EQ(outer.ret().at("mapdefun"),
+ EXPECT_EQ(GetRetval(*vectorized, 0),
strings::StrCat(cast_node.name(), ":y:0"));
- EXPECT_EQ(outer.node_def_size(), 1);
+ EXPECT_EQ(vectorized->node_def_size(), 1);
}
// Before:
@@ -330,16 +382,21 @@
{{}, {}}, inner.signature().name(), &outer);
CHECK_NOTNULL(map_defun);
- VectorizeMapDefun(&outer, &inner, map_defun);
- EXPECT_TRUE(!function_utils::ContainsFunctionNodeWithOp("MapDefun", outer));
- const NodeDef& cast_node =
- outer.node_def(function_utils::FindFunctionNodeWithOp("Cast", outer));
+ FunctionDefLibrary lib;
+ *lib.add_function() = outer;
+ *lib.add_function() = inner;
+ FunctionDef* vectorized;
+ EXPECT_TRUE(VectorizeMapDefun(outer, *map_defun, &lib, &vectorized).ok());
+ EXPECT_TRUE(
+ !function_utils::ContainsFunctionNodeWithOp("MapDefun", *vectorized));
+ const NodeDef& cast_node = vectorized->node_def(
+ function_utils::FindFunctionNodeWithOp("Cast", *vectorized));
EXPECT_EQ(cast_node.input(0), "x");
- EXPECT_EQ(outer.ret().at("mapdefun"),
+ EXPECT_EQ(GetRetval(*vectorized, 0),
strings::StrCat(cast_node.name(), ":y:0"));
- EXPECT_EQ(outer.ret().at("mapdefun_0"),
+ EXPECT_EQ(GetRetval(*vectorized, 1),
strings::StrCat(cast_node.name(), ":y:0"));
- EXPECT_EQ(outer.node_def_size(), 1);
+ EXPECT_EQ(vectorized->node_def_size(), 1);
}
// Before:
@@ -411,21 +468,26 @@
{{1}, {1}, {1}}, inner.signature().name(), &outer);
CHECK_NOTNULL(map_defun);
- VectorizeMapDefun(&outer, &inner, map_defun);
- EXPECT_TRUE(!function_utils::ContainsFunctionNodeWithOp("MapDefun", outer));
- const NodeDef& unpack_node =
- outer.node_def(function_utils::FindFunctionNodeWithOp("Unpack", outer));
+ FunctionDefLibrary lib;
+ *lib.add_function() = outer;
+ *lib.add_function() = inner;
+ FunctionDef* vectorized;
+ EXPECT_TRUE(VectorizeMapDefun(outer, *map_defun, &lib, &vectorized).ok());
+ EXPECT_TRUE(
+ !function_utils::ContainsFunctionNodeWithOp("MapDefun", *vectorized));
+ const NodeDef& unpack_node = vectorized->node_def(
+ function_utils::FindFunctionNodeWithOp("Unpack", *vectorized));
EXPECT_EQ(unpack_node.input(0), "x");
EXPECT_EQ(unpack_node.attr().at("axis").i(), 1);
EXPECT_EQ(unpack_node.attr().at("T").type(), DT_INT32);
EXPECT_EQ(unpack_node.attr().at("num").i(), 3);
- EXPECT_EQ(outer.ret().at("mapdefun"),
+ EXPECT_EQ(GetRetval(*vectorized, 0),
strings::StrCat(unpack_node.name(), ":output:0"));
- EXPECT_EQ(outer.ret().at("mapdefun_0"),
+ EXPECT_EQ(GetRetval(*vectorized, 1),
strings::StrCat(unpack_node.name(), ":output:1"));
- EXPECT_EQ(outer.ret().at("mapdefun_1"),
+ EXPECT_EQ(GetRetval(*vectorized, 2),
strings::StrCat(unpack_node.name(), ":output:2"));
- EXPECT_EQ(outer.node_def_size(), 1);
+ EXPECT_EQ(vectorized->node_def_size(), 1);
}
// Before:
@@ -486,7 +548,7 @@
{"ret1", "MyUnstack:output:1"},
{"ret2", "MyUnstack:output:2"}});
NodeDef* cast_op =
- AddCastNode("Cast", {"arg0"}, DT_INT32, DT_INT64, false, &inner);
+ AddCastNode("Cast", {"arg0"}, DT_INT32, DT_INT32, false, &inner);
CHECK_NOTNULL(cast_op);
NodeDef* unstack_op =
AddUnstackNode("MyUnstack", {"Cast:y:0"}, DT_INT32, 0, 3, &inner);
@@ -505,25 +567,30 @@
{{1}, {1}, {1}}, inner.signature().name(), &outer);
CHECK_NOTNULL(map_defun);
- VectorizeMapDefun(&outer, &inner, map_defun);
- EXPECT_TRUE(!function_utils::ContainsFunctionNodeWithOp("MapDefun", outer));
- const NodeDef& cast_node =
- outer.node_def(function_utils::FindFunctionNodeWithOp("Cast", outer));
+ FunctionDefLibrary lib;
+ *lib.add_function() = outer;
+ *lib.add_function() = inner;
+ FunctionDef* vectorized;
+ EXPECT_TRUE(VectorizeMapDefun(outer, *map_defun, &lib, &vectorized).ok());
+ EXPECT_TRUE(
+ !function_utils::ContainsFunctionNodeWithOp("MapDefun", *vectorized));
+ const NodeDef& cast_node = vectorized->node_def(
+ function_utils::FindFunctionNodeWithOp("Cast", *vectorized));
EXPECT_EQ(cast_node.input(0), "x");
- const NodeDef& unpack_node =
- outer.node_def(function_utils::FindFunctionNodeWithOp("Unpack", outer));
+ const NodeDef& unpack_node = vectorized->node_def(
+ function_utils::FindFunctionNodeWithOp("Unpack", *vectorized));
EXPECT_EQ(unpack_node.input(0), strings::StrCat(cast_node.name(), ":y:0"));
EXPECT_EQ(unpack_node.attr().at("axis").i(), 1);
EXPECT_EQ(unpack_node.attr().at("T").type(), DT_INT32);
EXPECT_EQ(unpack_node.attr().at("num").i(), 3);
- EXPECT_EQ(outer.ret().at("mapdefun"),
+ EXPECT_EQ(GetRetval(*vectorized, 0),
strings::StrCat(unpack_node.name(), ":output:0"));
- EXPECT_EQ(outer.ret().at("mapdefun_0"),
+ EXPECT_EQ(GetRetval(*vectorized, 1),
strings::StrCat(unpack_node.name(), ":output:1"));
- EXPECT_EQ(outer.ret().at("mapdefun_1"),
+ EXPECT_EQ(GetRetval(*vectorized, 2),
strings::StrCat(unpack_node.name(), ":output:2"));
- EXPECT_EQ(outer.node_def_size(), 2);
+ EXPECT_EQ(vectorized->node_def_size(), 2);
}
// Before:
@@ -561,9 +628,11 @@
FunctionDef inner =
CreateFunction("inner_function", {{"arg0", DT_INT32}},
{{"ret0", DT_INT64}}, {{"ret0", "Cast:y:0"}});
- // The attrs aren't relevant
- NodeDef* print_op =
- function_utils::AddNode("Print", "Print", {"arg0", "arg0"}, {}, &inner);
+ NodeDef* print_op = function_utils::AddNode(
+ "Print", "Print", {"arg0", "arg0"}, {/*attrs*/}, &inner);
+ graph_transforms::SetNodeAttr("T", DT_INT32, print_op);
+ graph_transforms::SetNodeAttr("U", gtl::ArraySlice<DataType>({DT_INT32}),
+ print_op);
CHECK_NOTNULL(print_op);
NodeDef* cast_op = AddCastNode("Cast", {"arg0", "^Print"}, DT_INT32, DT_INT64,
false, &inner);
@@ -578,11 +647,27 @@
inner.signature().name(), &outer);
CHECK_NOTNULL(map_defun);
- FunctionDef outer_copy(outer);
- FunctionDef inner_copy(inner);
- VectorizeMapDefun(&outer, &inner, map_defun);
+ FunctionDefLibrary lib;
+ *lib.add_function() = outer;
+ *lib.add_function() = inner;
+ FunctionDef* vectorized;
+ EXPECT_TRUE(VectorizeMapDefun(outer, *map_defun, &lib, &vectorized).ok());
// They should be unchanged
- EXPECT_TRUE(FunctionDefsEqual(outer_copy, outer));
+ // We check this somewhat manually as the names of nodes may have changed
+ EXPECT_EQ(vectorized->node_def_size(), 1);
+ const NodeDef& map_defun_node = vectorized->node_def(0);
+ EXPECT_EQ(map_defun_node.op(), "MapDefun");
+ FunctionLibraryDefinition lib_def(OpRegistry::Global(), lib);
+ const FunctionDef* map_defun_fn =
+ lib_def.Find(map_defun_node.attr().at("f").func().name());
+
+ const NodeDef& print_node = map_defun_fn->node_def(
+ function_utils::FindFunctionNodeWithOp("Print", *map_defun_fn));
+ const NodeDef& cast_node = map_defun_fn->node_def(
+ function_utils::FindFunctionNodeWithOp("Cast", *map_defun_fn));
+ string control_input = strings::StrCat("^", print_node.name());
+ EXPECT_TRUE(cast_node.input(0) == control_input ||
+ cast_node.input(1) == control_input);
}
// TODO(rachelim): More test cases when we get around to implementing them:
diff --git a/tensorflow/core/grappler/optimizers/meta_optimizer.cc b/tensorflow/core/grappler/optimizers/meta_optimizer.cc
index c59645e..c3d70a1 100644
--- a/tensorflow/core/grappler/optimizers/meta_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/meta_optimizer.cc
@@ -37,6 +37,7 @@
#include "tensorflow/core/grappler/utils/functions.h"
#include "tensorflow/core/grappler/utils/topological_sort.h"
#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/gtl/map_util.h"
#include "tensorflow/core/util/ptr_util.h"
namespace tensorflow {
@@ -115,6 +116,9 @@
Status MetaOptimizer::InitializeOptimizers(
std::vector<std::unique_ptr<GraphOptimizer>>* optimizers) const {
+ if (cfg_.disable_meta_optimizer()) {
+ return Status::OK();
+ }
if (!cfg_.disable_model_pruning()) {
optimizers->push_back(MakeUnique<ModelPruner>());
}
@@ -135,7 +139,7 @@
if (cfg_.remapping() != RewriterConfig::OFF) {
optimizers->push_back(MakeUnique<Remapper>(cfg_.remapping()));
}
- if (cfg_.pin_to_host_optimization() != RewriterConfig::OFF) {
+ if (cfg_.pin_to_host_optimization() == RewriterConfig::ON) {
optimizers->push_back(MakeUnique<PinToHostOptimizer>());
}
if (cfg_.arithmetic_optimization() != RewriterConfig::OFF) {
@@ -172,11 +176,12 @@
optimizers->push_back(MakeUnique<ScopedAllocatorOptimizer>(
cfg_.scoped_allocator_optimization(), cfg_.scoped_allocator_opts()));
}
- return InitializeCustomGraphOptimizers(optimizers);
+ return InitializeCustomGraphOptimizers(std::set<string>(), optimizers);
}
Status MetaOptimizer::InitializeOptimizersByName(
std::vector<std::unique_ptr<GraphOptimizer>>* optimizers) const {
+ std::set<string> initialized_custom_optimizers;
for (const string& optimizer_name : cfg_.optimizers()) {
auto optimizer = MakeNewOptimizer(optimizer_name);
if (optimizer) {
@@ -190,18 +195,26 @@
if (custom_optimizer) {
VLOG(2) << "Registered custom graph optimizer: " << optimizer_name;
- TF_RETURN_IF_ERROR(custom_optimizer->Init());
+ TF_RETURN_IF_ERROR(custom_optimizer->Init(
+ GetCustomGraphOptimizerConfig(optimizer_name)));
optimizers->push_back(std::move(custom_optimizer));
+ initialized_custom_optimizers.insert(optimizer_name);
} else {
VLOG(2) << "Can't register an optimizer by name: " << optimizer_name;
}
}
- return InitializeCustomGraphOptimizers(optimizers);
+ return InitializeCustomGraphOptimizers(initialized_custom_optimizers,
+ optimizers);
}
Status MetaOptimizer::InitializeCustomGraphOptimizers(
+ const std::set<string>& pre_initialized_optimizers,
std::vector<std::unique_ptr<GraphOptimizer>>* optimizers) const {
for (const auto& optimizer_config : cfg_.custom_optimizers()) {
+ if (pre_initialized_optimizers.find(optimizer_config.name()) !=
+ pre_initialized_optimizers.end()) {
+ continue;
+ }
// Initialize the ExperimentalImplementationSelector here instead of
// CustomizeOptimizer registry, due the static link issue in TensorRT for
// double registry.
@@ -237,6 +250,16 @@
return Status::OK();
}
+const RewriterConfig::CustomGraphOptimizer*
+MetaOptimizer::GetCustomGraphOptimizerConfig(const string& name) const {
+ for (const auto& config : cfg_.custom_optimizers()) {
+ if (config.name() == name) {
+ return &config;
+ }
+ }
+ return nullptr;
+}
+
Status MetaOptimizer::OptimizeGraph(Cluster* cluster, const GrapplerItem& item,
GraphDef* optimized_graph) {
int min_graph_nodes = cfg_.min_graph_nodes() == 0 ? kDefaultMinGraphNodes
@@ -391,6 +414,15 @@
FunctionLibraryDefinition flib(OpRegistry::Global(),
optimized_graph->library());
+ // Find functions for which we might need to compute a gradient at runtime.
+ gtl::FlatSet<string> differentiable_functions;
+ for (const NodeDef& node : optimized_graph->node()) {
+ if (IsSymbolicGradient(node)) {
+ const auto* f_attr = gtl::FindOrNull(node.attr(), "f");
+ if (f_attr) differentiable_functions.insert(f_attr->func().name());
+ }
+ }
+
// Optimize each function only once.
std::unordered_set<string> optimized_funcs;
bool optimize_function_library = true;
@@ -406,6 +438,8 @@
// Skip parametrized functions (function type or body is defined only at
// function call time by caller node attributes).
+ // They should be specialized to their instantiation type parameters by
+ // the function optimizer, before we can optimize function body.
if (IsParametrized(func)) continue;
VLOG(3) << "Optimize function: function=" << func_name;
@@ -420,6 +454,13 @@
TF_RETURN_IF_ERROR(MakeGrapplerFunctionItem(
func, flib, item.graph.versions().producer(), &func_item));
+ // If we need to compute the gradient of optimized function at runtime, we
+ // can't perform non-differentiable rewrites.
+ if (differentiable_functions.find(func_name) !=
+ differentiable_functions.end()) {
+ func_item.allowed_optimizations.non_differentiable_rewrites = false;
+ }
+
// Optimize function body graph.
GraphDef optimized_func_graph;
TF_RETURN_IF_ERROR(
@@ -470,6 +511,9 @@
}
bool MetaOptimizerEnabled(const RewriterConfig& cfg) {
+ if (cfg.disable_meta_optimizer()) {
+ return false;
+ }
return !cfg.disable_model_pruning() ||
cfg.layout_optimizer() != RewriterConfig::OFF ||
cfg.function_optimization() != RewriterConfig::OFF ||
@@ -483,7 +527,7 @@
cfg.memory_optimization() != RewriterConfig::NO_MEM_OPT ||
cfg.debug_stripper() == RewriterConfig::ON ||
cfg.scoped_allocator_optimization() == RewriterConfig::ON ||
- cfg.pin_to_host_optimization() != RewriterConfig::OFF ||
+ cfg.pin_to_host_optimization() == RewriterConfig::ON ||
!cfg.optimizers().empty() || !cfg.custom_optimizers().empty();
}
diff --git a/tensorflow/core/grappler/optimizers/meta_optimizer.h b/tensorflow/core/grappler/optimizers/meta_optimizer.h
index 831c5e3..99a0a33 100644
--- a/tensorflow/core/grappler/optimizers/meta_optimizer.h
+++ b/tensorflow/core/grappler/optimizers/meta_optimizer.h
@@ -54,7 +54,11 @@
std::vector<std::unique_ptr<GraphOptimizer>>* optimizers) const;
// Initialize active optimizers from RewriterConfig.custom_optimizers.
Status InitializeCustomGraphOptimizers(
+ const std::set<string>& pre_initialized_optimizers,
std::vector<std::unique_ptr<GraphOptimizer>>* optimizers) const;
+ // Returns the config for a custom graph optimizer. Null if none was found.
+ const RewriterConfig::CustomGraphOptimizer* GetCustomGraphOptimizerConfig(
+ const string& name) const;
// Run optimization pass over a single GrapplerItem. Meta optimizer might run
// multiple such passes: 1) for the main graph 2) for the function library
diff --git a/tensorflow/core/grappler/optimizers/meta_optimizer_test.cc b/tensorflow/core/grappler/optimizers/meta_optimizer_test.cc
index e74e0f7..3f3f433 100644
--- a/tensorflow/core/grappler/optimizers/meta_optimizer_test.cc
+++ b/tensorflow/core/grappler/optimizers/meta_optimizer_test.cc
@@ -25,6 +25,7 @@
#include "tensorflow/core/grappler/utils.h"
#include "tensorflow/core/grappler/utils/grappler_test.h"
#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/lib/gtl/map_util.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
@@ -71,6 +72,59 @@
REGISTER_GRAPH_OPTIMIZER(TestGraphOptimizer);
+class TestOptimizerWithParams : public TestOptimizer {
+ public:
+ Status Init(
+ const tensorflow::RewriterConfig_CustomGraphOptimizer* config) override {
+ CHECK(config != nullptr);
+ return Status::OK();
+ }
+};
+
+REGISTER_GRAPH_OPTIMIZER(TestOptimizerWithParams);
+
+// Record various properties of the GrapplerItems passed for optimization.
+class GrapplerItemPropertiesAccumulator : public CustomGraphOptimizer {
+ public:
+ static void SetAllowedOptimizations(
+ gtl::FlatMap<string, GrapplerItem::AllowedOptimizations>*
+ allowed_optimizations) {
+ allowed_optimizations_ = allowed_optimizations;
+ }
+ static void ResetAllowedOptimizations() { allowed_optimizations_ = nullptr; }
+
+ GrapplerItemPropertiesAccumulator() {}
+ string name() const override {
+ return "grappler_item_properties_accumulator";
+ }
+
+ Status Init(
+ const tensorflow::RewriterConfig_CustomGraphOptimizer* config) override {
+ return Status::OK();
+ }
+
+ Status Optimize(Cluster* cluster, const GrapplerItem& item,
+ GraphDef* optimized_graph) override {
+ *optimized_graph = item.graph;
+ if (allowed_optimizations_) {
+ allowed_optimizations_->insert({item.id, item.allowed_optimizations});
+ }
+ return Status::OK();
+ }
+
+ void Feedback(Cluster* cluster, const GrapplerItem& item,
+ const GraphDef& optimized_graph, double result) override {}
+
+ private:
+ static gtl::FlatMap<string, GrapplerItem::AllowedOptimizations>*
+ allowed_optimizations_;
+};
+
+gtl::FlatMap<string, GrapplerItem::AllowedOptimizations>*
+ GrapplerItemPropertiesAccumulator::allowed_optimizations_;
+
+REGISTER_GRAPH_OPTIMIZER(GrapplerItemPropertiesAccumulator);
+
class MetaOptimizerTest : public GrapplerTest {};
TEST_F(MetaOptimizerTest, RunsCustomOptimizer) {
@@ -90,6 +144,25 @@
EXPECT_TRUE(TestOptimizer::IsOptimized());
}
+TEST_F(MetaOptimizerTest, RunsCustomOptimizerWithParams) {
+ TrivialTestGraphInputYielder fake_input(4, 1, 10, false, {"CPU:0"});
+ GrapplerItem item;
+ CHECK(fake_input.NextItem(&item));
+
+ TestOptimizer::SetOptimized(false);
+ RewriterConfig rewriter_config;
+ rewriter_config.add_optimizers("TestOptimizerWithParams");
+ auto* custom_config = rewriter_config.add_custom_optimizers();
+ custom_config->set_name("TestOptimizerWithParams");
+ (*custom_config->mutable_parameter_map())["foo"] = AttrValue();
+
+ MetaOptimizer optimizer(nullptr, rewriter_config);
+ GraphDef output;
+ const Status status = optimizer.Optimize(nullptr, item, &output);
+ TF_EXPECT_OK(status);
+ EXPECT_TRUE(TestOptimizer::IsOptimized());
+}
+
TEST_F(MetaOptimizerTest, RunsCustomOptimizerAndCustomGraphOptimizer) {
TrivialTestGraphInputYielder fake_input(4, 1, 10, false, {"CPU:0"});
GrapplerItem item;
@@ -305,6 +378,89 @@
test::ExpectTensorEqual<int>(tensors_expected[1], tensors[1]);
}
+TEST_F(MetaOptimizerTest, OptimizeFunctionLibraryWithRestrictions) {
+ using test::function::NDef;
+ using FDH = FunctionDefHelper;
+
+ // We will record what type of optimizations meta optimizer allows for each
+ // GrapplerItem (main graph and graphs for each function).
+ gtl::FlatMap<string, GrapplerItem::AllowedOptimizations>
+ allowed_optimizations;
+ GrapplerItemPropertiesAccumulator::SetAllowedOptimizations(
+ &allowed_optimizations);
+
+ // Just record properties of optimized Grappler items.
+ RewriterConfig rewriter_config;
+ rewriter_config.set_meta_optimizer_iterations(RewriterConfig::TWO);
+ rewriter_config.add_optimizers("GrapplerItemPropertiesAccumulator");
+ rewriter_config.set_min_graph_nodes(-1);
+
+ MetaOptimizer optimizer(nullptr, rewriter_config);
+
+ // Define simple function library with two identical mul functions.
+ FunctionDef mul_func_1 = FunctionDefHelper::Create(
+ "MyMul1", {"x:float", "y:float"}, {"z:float"}, {},
+ {{{"mul"}, "Mul", {"x", "y"}, {}}},
+ /* Mapping between function returns and function node outputs. */
+ {{"z", "mul:z:0"}});
+
+ FunctionDef mul_func_2 = FunctionDefHelper::Create(
+ "MyMul2", {"x:float", "y:float"}, {"z:float"}, {},
+ {{{"mul"}, "Mul", {"x", "y"}, {}}},
+ /* Mapping between function returns and function node outputs. */
+ {{"z", "mul:z:0"}});
+
+ // Tensorflow graph:
+ //
+ // x0 = tf.Placeholder(tf.float);
+ // x1 = tf.Placeholder(tf.float);
+ // dy = tf.Placeholder(tf.float);
+ //
+ // mul_1 = MyMul1(x0, x1);
+ // mul_2 = MyMul2(x0, x1);
+ // dx = SymbolicGradient({x0, x1, dy}, f=MyMul2)
+ GrapplerItem item;
+ item.id = "main";
+ item.graph = test::function::GDef(
+ {NDef("x0", "Placeholder", {}, {{"dtype", DT_FLOAT}}, kDevice),
+ NDef("x1", "Placeholder", {}, {{"dtype", DT_FLOAT}}, kDevice),
+ NDef("dy", "Placeholder", {}, {{"dtype", DT_FLOAT}}, kDevice),
+ // Calls into function library
+ NDef("mul_1", "MyMul1", {"x0", "x1"}, {}, kDevice),
+ NDef("mul_2", "MyMul2", {"x0", "x1"}, {}, kDevice),
+ // Symbolic gradient of a MyMul2
+ NDef("dx", "SymbolicGradient", {"x0", "x1", "dy"},
+ {{"f", FDH::FunctionRef("MyMul2", {})},
+ {"Tin", DataTypeSlice{DT_FLOAT}},
+ {"Tout", DataTypeSlice{DT_FLOAT, DT_FLOAT}}},
+ kDevice)},
+ // FunctionLib
+ {mul_func_1, mul_func_2});
+ item.fetch = {"mul_1", "mul_2", "dx"};
+
+ GraphDef output;
+ TF_EXPECT_OK(optimizer.Optimize(nullptr, item, &output));
+
+ // Our custom optimizer must be called for the main graph and for the two
+ // functions.
+ ASSERT_EQ(allowed_optimizations.size(), 3);
+
+ auto allowed_optimizations_main =
+ gtl::FindOrNull(allowed_optimizations, "main");
+ ASSERT_NE(allowed_optimizations_main, nullptr);
+ EXPECT_TRUE(allowed_optimizations_main->non_differentiable_rewrites);
+
+ auto allowed_optimizations_my_mul_1 =
+ gtl::FindOrNull(allowed_optimizations, "MyMul1");
+ ASSERT_NE(allowed_optimizations_my_mul_1, nullptr);
+ EXPECT_TRUE(allowed_optimizations_my_mul_1->non_differentiable_rewrites);
+
+ auto allowed_optimizations_my_mul_2 =
+ gtl::FindOrNull(allowed_optimizations, "MyMul2");
+ ASSERT_NE(allowed_optimizations_my_mul_2, nullptr);
+ EXPECT_FALSE(allowed_optimizations_my_mul_2->non_differentiable_rewrites);
+}
+
} // namespace
} // namespace grappler
} // namespace tensorflow
diff --git a/tensorflow/core/grappler/optimizers/pin_to_host_optimizer.cc b/tensorflow/core/grappler/optimizers/pin_to_host_optimizer.cc
index 2190d38..89eb760 100644
--- a/tensorflow/core/grappler/optimizers/pin_to_host_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/pin_to_host_optimizer.cc
@@ -169,7 +169,13 @@
}
// All the nodes that should be blacklisted and not swapped.
-bool IsBlacklisted(const NodeDef& node) { return IsCollective(node); }
+bool IsBlacklisted(const NodeDef& node) {
+ return
+ // Collective ops should not be swapped.
+ IsCollective(node) ||
+ // NoOp breaks perf regression tests (probably due to group dependencies).
+ IsNoOp(node);
+}
} // end namespace internal
Status PinToHostOptimizer::Optimize(Cluster* cluster, const GrapplerItem& item,
diff --git a/tensorflow/core/grappler/optimizers/remapper.cc b/tensorflow/core/grappler/optimizers/remapper.cc
index 008a289..9ada8b7 100644
--- a/tensorflow/core/grappler/optimizers/remapper.cc
+++ b/tensorflow/core/grappler/optimizers/remapper.cc
@@ -168,11 +168,12 @@
Status Remapper::Optimize(Cluster* /*cluster*/, const GrapplerItem& item,
GraphDef* optimized_graph) {
GraphProperties properties(item);
- TF_RETURN_IF_ERROR(properties.InferStatically(false));
+ bool inferred_properties = false;
GraphView graph(const_cast<GraphDef*>(&item.graph));
// During inference, most of the inputs to FusedBatchNorm are constant, and we
// can therefore replace the op with a much cheaper set of primitives.
+ optimized_graph->mutable_node()->Reserve(item.graph.node_size());
for (const NodeDef& node : item.graph.node()) {
if (node.op() == "FusedBatchNorm" || node.op() == "FusedBatchNormV2") {
bool optimizable = (node.attr().count("T") == 0 ||
@@ -181,6 +182,11 @@
!node.attr().at("is_training").b());
if (optimizable) {
int const_inputs = 0;
+ if (!inferred_properties) {
+ // Infer properties lazily in case they are not needed.
+ TF_RETURN_IF_ERROR(properties.InferStatically(false));
+ inferred_properties = true;
+ }
const auto& props = properties.GetInputProperties(node.name());
for (const auto& prop : props) {
if (prop.has_value()) {
diff --git a/tensorflow/core/grappler/optimizers/shape_optimizer.cc b/tensorflow/core/grappler/optimizers/shape_optimizer.cc
index 4542d17..6ccb1cd 100644
--- a/tensorflow/core/grappler/optimizers/shape_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/shape_optimizer.cc
@@ -33,7 +33,7 @@
*optimized_graph = item.graph;
GraphProperties properties(item);
- TF_RETURN_IF_ERROR(properties.InferStatically(false));
+ bool inferred_properties = false;
GraphView graph(optimized_graph);
// The product of all the dimensions in a tensor shape can be expressed more
@@ -55,6 +55,11 @@
}
const GraphView::OutputPort reduce_indices =
graph.GetRegularFanin(GraphView::InputPort(fanout.node, 1));
+ if (!inferred_properties) {
+ // Infer properties lazily in case they are not needed.
+ TF_RETURN_IF_ERROR(properties.InferStatically(false));
+ inferred_properties = true;
+ }
const auto& prop =
properties.GetOutputProperties(reduce_indices.node->name());
if (prop.size() < reduce_indices.port_id) {
@@ -92,6 +97,11 @@
if (!IsSize(*input1.node) || !IsSize(*input2.node)) {
continue;
}
+ if (!inferred_properties) {
+ // Infer properties lazily in case they are not needed.
+ TF_RETURN_IF_ERROR(properties.InferStatically(false));
+ inferred_properties = true;
+ }
const auto& prop1 = properties.GetInputProperties(input1.node->name());
const auto& prop2 = properties.GetInputProperties(input2.node->name());
if (prop1.size() != 1 || prop2.size() != 1) {
diff --git a/tensorflow/core/grappler/utils.cc b/tensorflow/core/grappler/utils.cc
index db6e4e6..5867d01 100644
--- a/tensorflow/core/grappler/utils.cc
+++ b/tensorflow/core/grappler/utils.cc
@@ -156,45 +156,6 @@
return !name.empty() && name[0] == '^';
}
-string NodeName(const string& name) {
- int position;
- return ParseNodeName(name, &position);
-}
-
-int NodePosition(const string& name) {
- int position;
- ParseNodeNameAsStringPiece(name, &position);
- return position;
-}
-
-int NodePositionIfSameNode(const string& input_name, const string& node_name) {
- const bool is_ctrl = input_name[0] == '^';
- auto input_it = is_ctrl ? input_name.begin() + 1 : input_name.begin();
- auto node_it = node_name.begin();
- if (node_name.empty() ||
- std::distance(input_it, input_name.end()) < node_name.size()) {
- return -2;
- }
- while (node_it != node_name.end()) {
- if (*input_it++ != *node_it++) {
- return -2;
- }
- }
- if (input_it == input_name.end()) {
- return is_ctrl ? -1 : 0;
- } else if (*input_it++ == ':') {
- StringPiece remaining(&(*input_it),
- std::distance(input_it, input_name.end()));
- int position;
- if (!strings::safe_strto32(remaining, &position)) {
- return -2;
- }
- return is_ctrl ? -1 : position;
- } else {
- return -2;
- }
-}
-
string AddPrefixToNodeName(const string& name, const string& prefix,
const string& delimiter) {
if (!name.empty()) {
diff --git a/tensorflow/core/grappler/utils.h b/tensorflow/core/grappler/utils.h
index 296ee16..95126d4 100644
--- a/tensorflow/core/grappler/utils.h
+++ b/tensorflow/core/grappler/utils.h
@@ -29,7 +29,6 @@
#include "tensorflow/core/lib/core/stringpiece.h"
#include "tensorflow/core/lib/core/threadpool.h"
#include "tensorflow/core/lib/gtl/inlined_vector.h"
-#include "tensorflow/core/lib/strings/scanner.h"
namespace tensorflow {
namespace grappler {
@@ -102,40 +101,92 @@
// True iff 'name1' and 'name2' refer to the same input.
bool IsSameInput(const string& name1, const string& name2);
+// Returns the trailing position number (or zero if no number is present) if
+// NodeName(input_name) is equal to node_name. Returns -1 for control inputs.
+// Returns -2 if NodeName(input_name) is not equal to node_name.
+// Note: This function is used very heavily, and this hand-optimized
+// version is 3-4x faster than the version using Scanner, which it replaced.
+// This is worth the reduction in readability.
+inline int NodePositionIfSameNode(const string& input_name,
+ const string& node_name) {
+ if (input_name.empty()) return -2;
+ const bool is_ctrl = input_name[0] == '^';
+ auto input_it = is_ctrl ? input_name.begin() + 1 : input_name.begin();
+ auto node_it = node_name.begin();
+ if (node_name.empty() ||
+ std::distance(input_it, input_name.end()) < node_name.size()) {
+ return -2;
+ }
+ while (node_it != node_name.end()) {
+ if (*input_it++ != *node_it++) {
+ return -2;
+ }
+ }
+ if (input_it == input_name.end()) {
+ return is_ctrl ? -1 : 0;
+ } else if (*input_it++ == ':') {
+ StringPiece remaining(&(*input_it),
+ std::distance(input_it, input_name.end()));
+ int position;
+ if (!strings::safe_strto32(remaining, &position)) {
+ return -2;
+ }
+ return is_ctrl ? -1 : position;
+ } else {
+ return -2;
+ }
+}
+
// Return the node name corresponding to 'name' if name is valid, or the empty
// string otherwise.
-string NodeName(const string& name);
+inline StringPiece NodeNameAsStringPiece(const string& name) {
+ static const string empty;
+ if (name.empty()) return StringPiece(empty);
+ const auto begin_it = name[0] == '^' ? name.begin() + 1 : name.begin();
+ auto end_it = begin_it;
+ while (end_it != name.end() && *end_it != ':') {
+ ++end_it;
+ }
+ if (end_it != name.end() && *end_it != ':') {
+ return StringPiece(empty);
+ }
+ return StringPiece(&(*begin_it), std::distance(begin_it, end_it));
+}
-// Get the trailing position number ":{digits}" (if any) of a node name.
-// Returns -1 for control inputs.
-int NodePosition(const string& name);
+// Return the node name corresponding to 'name' if name is valid, or the empty
+// string otherwise.
+inline string NodeName(const string& name) {
+ return string(NodeNameAsStringPiece(name));
+}
+// Returns the node name and position in a single call.
inline StringPiece ParseNodeNameAsStringPiece(const string& name,
int* position) {
- // Strip the prefix '^' (if any), and strip the trailing ":{digits} (if any)
- // to get a node name.
- strings::Scanner scan(name);
- scan.ZeroOrOneLiteral("^")
- .RestartCapture()
- .One(strings::Scanner::LETTER_DIGIT_DOT_UNDERSCORE)
- .Any(strings::Scanner::LETTER_DIGIT_DASH_DOT_SLASH_UNDERSCORE);
- StringPiece capture;
- StringPiece remaining;
- if (scan.Peek(':') != ':' || !scan.GetResult(&remaining, &capture)) {
+ static const string empty;
+ if (name.empty()) {
*position = 0;
- static const string empty;
return StringPiece(empty);
- } else {
- if (name[0] == '^') {
- *position = -1;
- } else if (remaining.empty()) {
- *position = 0;
- } else {
- // Skip the first ':' character.
- CHECK(strings::safe_strto32(remaining.substr(1), position));
- }
- return capture;
}
+ const bool is_ctrl = name[0] == '^';
+ const auto begin_it = is_ctrl ? name.begin() + 1 : name.begin();
+ *position = is_ctrl ? -1 : 0;
+ auto end_it = begin_it;
+ while (end_it != name.end() && *end_it != ':') {
+ ++end_it;
+ }
+ const StringPiece node_name(&(*begin_it), std::distance(begin_it, end_it));
+ if (end_it != name.end()) {
+ if (*end_it != ':') {
+ return StringPiece(empty);
+ } else if (!is_ctrl) {
+ ++end_it;
+ StringPiece remaining(&(*end_it), std::distance(end_it, name.end()));
+ if (!strings::safe_strto32(remaining, position)) {
+ return StringPiece(empty);
+ }
+ }
+ }
+ return node_name;
}
// Returns the node name and position in a single call.
@@ -143,10 +194,11 @@
return string(ParseNodeNameAsStringPiece(name, position));
}
-// Returns NodePosition(input_name) if NodeName(input_name) == node_name.
-// Otherwise returns -2;
-// REQUIRES: inputs_name.size() > 0 && node_name.size() > 0.
-int NodePositionIfSameNode(const string& input_name, const string& node_name);
+inline int NodePosition(const string& name) {
+ int position;
+ ParseNodeNameAsStringPiece(name, &position);
+ return position;
+}
// Add a prefix to a node name with a custom delimiter.
string AddPrefixToNodeName(const string& name, const string& prefix,
diff --git a/tensorflow/core/grappler/utils/functions.cc b/tensorflow/core/grappler/utils/functions.cc
index a428aea..6861fb4 100644
--- a/tensorflow/core/grappler/utils/functions.cc
+++ b/tensorflow/core/grappler/utils/functions.cc
@@ -41,7 +41,8 @@
tensorflow::NameRangeMap outputs_range_map;
TF_RETURN_IF_ERROR(tensorflow::NameRangesForNode(
node, registration.op_def, nullptr, &outputs_range_map));
- connectivity->RegisterFunctionBodyOutputs(node.name(), outputs_range_map);
+ connectivity->RegisterFunctionBodyOutputs(node.name(),
+ std::move(outputs_range_map));
return Status::OK();
}
@@ -75,20 +76,22 @@
} // namespace
void GrapplerFunctionConnectivity::RegisterInputArgExpansion(
- const InputArgExpansion& input_arg_expansion) {
- const auto& input_name = input_arg_expansion.input_name;
+ InputArgExpansion input_arg_expansion) {
+ string input_name = input_arg_expansion.input_name;
const auto& placeholders = input_arg_expansion.placeholders;
- input_arg_expansions_.emplace(input_name, input_arg_expansion);
+
for (int i = 0; i < placeholders.size(); ++i) {
const string& placeholder = input_arg_expansion.placeholders[i];
- input_arg_placeholders_.emplace(
- placeholder, InputArgPlaceholder{input_name, /*position=*/i});
+ input_arg_placeholders_.insert(
+ {placeholder, InputArgPlaceholder{input_name, /*position=*/i}});
}
+ input_arg_expansions_.insert(
+ {std::move(input_name), std::move(input_arg_expansion)});
}
void GrapplerFunctionConnectivity::RegisterFunctionBodyOutputs(
- const string& node_name, const tensorflow::NameRangeMap& outputs) {
- function_body_outputs_[node_name] = outputs;
+ const string& node_name, tensorflow::NameRangeMap&& outputs) {
+ function_body_outputs_[node_name] = std::move(outputs);
}
Status GrapplerFunctionConnectivity::ExpandFunctionDefInput(
@@ -174,11 +177,12 @@
const auto& output_range = output->second;
if (position == -1) {
+ graph_def_inputs->reserve(graph_def_inputs->size() +
+ output_range.second - output_range.first);
// If position is not defined expand node output range
for (int i = output_range.first; i < output_range.second; ++i) {
- i == 0 ? graph_def_inputs->push_back(node_name)
- : graph_def_inputs->push_back(
- strings::StrCat(node_name, ":", i));
+ graph_def_inputs->push_back(
+ i == 0 ? node_name : strings::StrCat(node_name, ":", i));
}
} else {
if (position > (output_range.second - output_range.first)) {
@@ -187,9 +191,8 @@
" position: ", position, " (out of range)");
}
int pos = output_range.first + position;
- pos == 0 ? graph_def_inputs->push_back(node_name)
- : graph_def_inputs->push_back(
- strings::StrCat(node_name, ":", pos));
+ graph_def_inputs->push_back(
+ pos == 0 ? node_name : strings::StrCat(node_name, ":", pos));
}
return Status::OK();
@@ -211,8 +214,8 @@
}
function_body_node->clear_input();
- for (const string& expanded_input : expanded_inputs)
- function_body_node->add_input(expanded_input);
+ for (string& expanded_input : expanded_inputs)
+ function_body_node->add_input(std::move(expanded_input));
return Status::OK();
}
@@ -323,7 +326,7 @@
// Fill the feed nodes with input placeholders.
for (const InputArgExpansion& input_arg : input_arg_expansions_) {
for (const string& placeholder : input_arg.placeholders) {
- feed.emplace_back(placeholder, Tensor());
+ feed.push_back({placeholder, Tensor()});
input_arg_placeholders_.insert(placeholder);
}
}
@@ -460,7 +463,7 @@
auto it = func_instantiation_attr.find(placeholder);
if (it != func_instantiation_attr.end()) {
- body_parameters->emplace(placeholder, it->second);
+ body_parameters->insert({placeholder, it->second});
} else {
return errors::InvalidArgument("Can't resolve placeholder: ",
placeholder);
@@ -498,10 +501,6 @@
// GraphDef input format (name[:position])
GrapplerFunctionConnectivity connectivity;
- std::vector<InputArgExpansion> inputs;
- std::vector<OutputArgExpansion> outputs;
- std::vector<string> keep_nodes;
-
// Function body shares the library with the graph that instantiated it.
GraphDef function_body;
*function_body.mutable_library() = flib.ToProto();
@@ -518,6 +517,9 @@
}
}
+ std::vector<InputArgExpansion> inputs;
+ inputs.reserve(signature.input_arg_size());
+
// For each input argument create a placeholder in function body.
for (const OpDef::ArgDef& input : signature.input_arg()) {
if (!input.type_list_attr().empty() || !input.number_attr().empty()) {
@@ -542,9 +544,10 @@
/*is_ref*/ input.is_ref(),
/*placeholders=*/{input.name()}};
connectivity.RegisterInputArgExpansion(input_expansion);
- inputs.push_back(input_expansion);
+ 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();
@@ -572,6 +575,8 @@
TF_RETURN_IF_ERROR(connectivity.ExpandNodeInputs(&node));
}
+ std::vector<OutputArgExpansion> outputs;
+ outputs.reserve(signature.output_arg_size());
// Add function outputs
for (const OpDef::ArgDef& out : signature.output_arg()) {
std::vector<string> output_tensors;
@@ -589,8 +594,8 @@
OutputArgExpansion output{/*output_name=*/out.name(),
/*data_type=*/output_data_type,
/*is_ref=*/out.is_ref(),
- /*output_tensors=*/output_tensors};
- outputs.push_back(output);
+ /*output_tensors=*/std::move(output_tensors)};
+ outputs.push_back(std::move(output));
}
bool is_stateful = signature.is_stateful();
diff --git a/tensorflow/core/grappler/utils/functions.h b/tensorflow/core/grappler/utils/functions.h
index 733caf3..ef944ce 100644
--- a/tensorflow/core/grappler/utils/functions.h
+++ b/tensorflow/core/grappler/utils/functions.h
@@ -18,6 +18,7 @@
#include <memory>
#include <string>
+#include <unordered_map>
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/function.pb.h"
@@ -70,9 +71,9 @@
// and fold it back when doing backward conversion.
class GrapplerFunctionConnectivity {
public:
- void RegisterInputArgExpansion(const InputArgExpansion& input_arg_expansion);
+ void RegisterInputArgExpansion(InputArgExpansion input_arg_expansion);
void RegisterFunctionBodyOutputs(const string& node_name,
- const tensorflow::NameRangeMap& outputs);
+ tensorflow::NameRangeMap&& outputs);
// Expand input encoded in FunctionDef format (name[:output][:position]) into
// multiple inputs in GraphDef format (name[:position]).
diff --git a/tensorflow/core/grappler/utils_test.cc b/tensorflow/core/grappler/utils_test.cc
index 6b787a6..9b6c1f6 100644
--- a/tensorflow/core/grappler/utils_test.cc
+++ b/tensorflow/core/grappler/utils_test.cc
@@ -371,6 +371,25 @@
BM_NodePositionIfSameNode("blah", "foo/bar/baz", NoMatch_0);
BM_NodePositionIfSameNode("foo/bar/baz/gnu", "foo/bar/baz", NoMatch_end);
+#define BM_ParseNodeNameAsStringPiece(I, NAME) \
+ static void BM_ParseNodeNameAsStringPiece_##NAME(int iters) { \
+ string input = I; \
+ for (int i = 0; i < iters; ++i) { \
+ int position; \
+ const StringPiece name = ParseNodeNameAsStringPiece(input, &position); \
+ CHECK_GE(position, -1); \
+ CHECK(!name.empty()); \
+ } \
+ } \
+ BENCHMARK(BM_ParseNodeNameAsStringPiece_##NAME)
+
+BM_ParseNodeNameAsStringPiece("foo", foo);
+BM_ParseNodeNameAsStringPiece("foo/bar/baz", foo_bar_baz);
+BM_ParseNodeNameAsStringPiece("^foo/bar/baz", foo_bar_baz_ctrl);
+BM_ParseNodeNameAsStringPiece("foo:123", foo123);
+BM_ParseNodeNameAsStringPiece("foo/bar/baz:123", foo_bar_baz_123);
+BM_ParseNodeNameAsStringPiece("^foo/bar/baz:123", foo_bar_baz_123_ctrl);
+
} // namespace
} // namespace grappler
} // namespace tensorflow
diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index 1a3db2c..9439ab3 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -1197,8 +1197,10 @@
tf_cc_test(
name = "example_parsing_ops_test",
- size = "large",
+ size = "medium",
srcs = ["example_parsing_ops_test.cc"],
+ shard_count = 4,
+ tags = ["optonly"],
deps = [
":example_parsing_ops",
":ops_testutil",
@@ -2028,8 +2030,8 @@
":variable_ops",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
+ "//tensorflow/core:lib_internal",
"//tensorflow/core:resource_variable_ops_op_lib",
- "//third_party/eigen3",
],
)
@@ -4049,11 +4051,6 @@
)
SPARSE_DEPS = [
- ":bounds_check",
- ":cwise_op",
- ":fill_functor",
- ":scatter_functor",
- "//third_party/eigen3",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:sparse_ops_op_lib",
@@ -4086,7 +4083,9 @@
tf_kernel_library(
name = "sparse_cross_op",
prefix = "sparse_cross_op",
- deps = SPARSE_DEPS,
+ deps = SPARSE_DEPS + [
+ "//third_party/eigen3",
+ ],
)
tf_kernel_library(
@@ -4098,13 +4097,19 @@
tf_kernel_library(
name = "sparse_dense_binary_op_shared",
prefix = "sparse_dense_binary_op_shared",
- deps = SPARSE_DEPS,
+ deps = SPARSE_DEPS + [
+ ":cwise_op",
+ "//third_party/eigen3",
+ ],
)
tf_kernel_library(
name = "sparse_sparse_binary_op_shared",
prefix = "sparse_sparse_binary_op_shared",
- deps = SPARSE_DEPS,
+ deps = SPARSE_DEPS + [
+ ":cwise_op",
+ "//third_party/eigen3",
+ ],
)
tf_kernel_library(
@@ -4136,7 +4141,9 @@
tf_kernel_library(
name = "sparse_softmax",
prefix = "sparse_softmax",
- deps = SPARSE_DEPS,
+ deps = SPARSE_DEPS + [
+ "//third_party/eigen3",
+ ],
)
tf_kernel_library(
@@ -4148,25 +4155,37 @@
tf_kernel_library(
name = "sparse_tensor_dense_add_op",
prefix = "sparse_tensor_dense_add_op",
- deps = SPARSE_DEPS,
+ deps = SPARSE_DEPS + [
+ ":scatter_functor",
+ "//third_party/eigen3",
+ ],
)
tf_kernel_library(
name = "sparse_tensor_dense_matmul_op",
prefix = "sparse_tensor_dense_matmul_op",
- deps = SPARSE_DEPS,
+ deps = SPARSE_DEPS + [
+ ":bounds_check",
+ ":fill_functor",
+ "//third_party/eigen3",
+ ],
)
tf_kernel_library(
name = "sparse_to_dense_op",
prefix = "sparse_to_dense_op",
- deps = SPARSE_DEPS,
+ deps = SPARSE_DEPS + [
+ "//third_party/eigen3",
+ ],
)
tf_kernel_library(
name = "sparse_xent_op",
prefix = "sparse_xent_op",
- deps = SPARSE_DEPS,
+ deps = SPARSE_DEPS + [
+ ":bounds_check",
+ "//third_party/eigen3",
+ ],
)
tf_kernel_library(
@@ -4431,6 +4450,7 @@
":string_strip_op",
":string_to_hash_bucket_op",
":substr_op",
+ ":unicode_script_op",
],
)
@@ -5254,6 +5274,8 @@
"cwise_op_squared_difference.cc",
"cwise_op_sub.cc",
"cwise_op_tanh.cc",
+ "cwise_op_xlogy.cc",
+ "cwise_op_xdivy.cc",
"data_format_ops.cc",
"decode_wav_op.cc",
"deep_conv2d.cc",
@@ -5469,6 +5491,7 @@
"batch_kernels.*",
"regex_full_match_op.cc",
"regex_replace_op.cc",
+ "unicode_script_op.cc",
# Ops that are inherently incompatible with Android (e.g. tied to x86 platform).
"mkl_*",
"xsmm_*",
@@ -6414,6 +6437,12 @@
)
tf_mkl_kernel_library(
+ name = "mkl_slice_op",
+ prefix = "mkl_slice_op",
+ deps = ARRAY_DEPS + mkl_deps(),
+)
+
+tf_mkl_kernel_library(
name = "mkl_identity_op",
prefix = "mkl_identity_op",
deps = ARRAY_DEPS + mkl_deps(),
@@ -6557,6 +6586,16 @@
],
)
+tf_kernel_library(
+ name = "unicode_script_op",
+ srcs = ["unicode_script_op.cc"],
+ deps = [
+ "//tensorflow/core:framework",
+ "//tensorflow/core:string_ops_op_lib",
+ "@icu//:common",
+ ],
+)
+
# -----------------------------------------------------------------------------
# Google-internal targets. These must be at the end for syncrepo.
diff --git a/tensorflow/core/kernels/batching_util/BUILD b/tensorflow/core/kernels/batching_util/BUILD
index 792eb74..0d53240 100644
--- a/tensorflow/core/kernels/batching_util/BUILD
+++ b/tensorflow/core/kernels/batching_util/BUILD
@@ -1,7 +1,7 @@
# Description: Utilities.
package(
- default_visibility = ["//tensorflow:internal"],
+ default_visibility = ["//visibility:public"],
)
licenses(["notice"]) # Apache 2.0
@@ -12,7 +12,6 @@
name = "periodic_function_dynamic",
srcs = ["periodic_function.cc"],
hdrs = ["periodic_function.h"],
- visibility = ["//visibility:public"],
deps = [
"//tensorflow/core:framework_headers_lib",
"//tensorflow/core:protos_all_cc",
@@ -21,7 +20,6 @@
cc_library(
name = "periodic_function",
- visibility = ["//visibility:public"],
deps = [
":periodic_function_dynamic",
"//tensorflow/core:lib",
@@ -190,7 +188,6 @@
testonly = 1,
srcs = ["fake_clock_env.cc"],
hdrs = ["fake_clock_env.h"],
- visibility = ["//visibility:public"],
deps = [
"//tensorflow/core:lib",
"//tensorflow/core:tensorflow",
diff --git a/tensorflow/core/kernels/collective_ops.cc b/tensorflow/core/kernels/collective_ops.cc
index e0da911..82e2913 100644
--- a/tensorflow/core/kernels/collective_ops.cc
+++ b/tensorflow/core/kernels/collective_ops.cc
@@ -143,6 +143,7 @@
c->forward_input_or_allocate_output(
{0}, 0, c->input(0).shape(), &output),
done);
+ col_params_.instance.shape = c->input(0).shape();
}
if (!CanProceedWithCompute(c, col_exec, done)) return;
auto actual_done = [c, col_exec, done](const Status& s) {
@@ -171,7 +172,7 @@
OP_REQUIRES_OK(
c, c->GetAttr("instance_key", &col_params_.instance.instance_key));
OP_REQUIRES_OK(c, c->GetAttr("T", &col_params_.instance.data_type));
- OP_REQUIRES_OK(c, c->GetAttr("shape", &shape_));
+ OP_REQUIRES_OK(c, c->GetAttr("shape", &col_params_.instance.shape));
col_params_.is_source = true;
col_params_.instance.impl_details.subdiv_offsets = {0};
@@ -195,13 +196,14 @@
if (c->mutable_output(0) == nullptr) {
// Allocate the output tensor, trying to reuse the input.
Tensor* output = nullptr;
- OP_REQUIRES_OK_ASYNC(
- c, c->forward_input_or_allocate_output({0}, 0, shape_, &output),
- done);
+ OP_REQUIRES_OK_ASYNC(c,
+ c->forward_input_or_allocate_output(
+ {0}, 0, col_params_.instance.shape, &output),
+ done);
}
if (!CanProceedWithCompute(c, col_exec, done)) return;
OP_REQUIRES_ASYNC(
- c, shape_.IsSameSize(c->input(0).shape()),
+ c, col_params_.instance.shape.IsSameSize(c->input(0).shape()),
errors::Internal("Declared shape of op ", col_params_.name,
" does not match shape of input"),
done);
@@ -214,8 +216,6 @@
}
private:
- TensorShape shape_;
-
TF_DISALLOW_COPY_AND_ASSIGN(CollectiveBcastSendOpKernel);
};
@@ -234,7 +234,7 @@
OP_REQUIRES_OK(
c, c->GetAttr("instance_key", &col_params_.instance.instance_key));
OP_REQUIRES_OK(c, c->GetAttr("T", &col_params_.instance.data_type));
- OP_REQUIRES_OK(c, c->GetAttr("shape", &shape_));
+ OP_REQUIRES_OK(c, c->GetAttr("shape", &col_params_.instance.shape));
col_params_.is_source = false;
col_params_.instance.impl_details.subdiv_offsets = {0};
@@ -258,7 +258,8 @@
if (c->mutable_output(0) == nullptr) {
// No input, so must allocate output.
Tensor* output = nullptr;
- OP_REQUIRES_OK_ASYNC(c, c->allocate_output(0, shape_, &output), done);
+ OP_REQUIRES_OK_ASYNC(
+ c, c->allocate_output(0, col_params_.instance.shape, &output), done);
}
if (!CanProceedWithCompute(c, col_exec, done)) return;
@@ -270,8 +271,6 @@
}
private:
- TensorShape shape_;
-
TF_DISALLOW_COPY_AND_ASSIGN(CollectiveBcastRecvOpKernel);
};
diff --git a/tensorflow/core/kernels/conv_ops.cc b/tensorflow/core/kernels/conv_ops.cc
index 717a9f4..78856c4 100644
--- a/tensorflow/core/kernels/conv_ops.cc
+++ b/tensorflow/core/kernels/conv_ops.cc
@@ -264,150 +264,198 @@
};
#endif
+#define TF_REQUIRES(EXP, STATUS) \
+ do { \
+ if (!TF_PREDICT_TRUE(EXP)) return (STATUS); \
+ } while (false)
+
+Status InitConv2DParameters(const OpKernelConstruction* context,
+ Conv2DParameters* params) {
+ TF_RETURN_IF_ERROR(context->GetAttr("dilations", ¶ms->dilations));
+ TF_RETURN_IF_ERROR(context->GetAttr("strides", ¶ms->strides));
+ TF_RETURN_IF_ERROR(context->GetAttr("padding", ¶ms->padding));
+ string data_format_string;
+ TF_RETURN_IF_ERROR(context->GetAttr("data_format", &data_format_string));
+ TF_REQUIRES(FormatFromString(data_format_string, ¶ms->data_format),
+ errors::InvalidArgument("Invalid data format"));
+
+ const auto& strides = params->strides;
+ const auto& dilations = params->dilations;
+ const auto& data_format = params->data_format;
+
+ TF_REQUIRES(dilations.size() == 4,
+ errors::InvalidArgument("Sliding window dilations field must "
+ "specify 4 dimensions"));
+ TF_REQUIRES(strides.size() == 4,
+ errors::InvalidArgument("Sliding window strides field must "
+ "specify 4 dimensions"));
+ const int64 stride_n = GetTensorDim(strides, data_format, 'N');
+ const int64 stride_c = GetTensorDim(strides, data_format, 'C');
+ const int64 stride_h = GetTensorDim(strides, data_format, 'H');
+ const int64 stride_w = GetTensorDim(strides, data_format, 'W');
+ TF_REQUIRES(
+ stride_n == 1 && stride_c == 1,
+ errors::InvalidArgument("Current implementation does not yet support "
+ "strides in the batch and depth dimensions."));
+ TF_REQUIRES(stride_h > 0 && stride_w > 0,
+ errors::InvalidArgument(
+ "Row and column strides should be larger than 0."));
+
+ const int64 dilation_n = GetTensorDim(dilations, data_format, 'N');
+ const int64 dilation_c = GetTensorDim(dilations, data_format, 'C');
+ const int64 dilation_h = GetTensorDim(dilations, data_format, 'H');
+ const int64 dilation_w = GetTensorDim(dilations, data_format, 'W');
+ TF_REQUIRES(
+ dilation_n == 1 && dilation_c == 1,
+ errors::InvalidArgument("Current implementation does not yet support "
+ "dilations in the batch and depth dimensions."));
+ TF_REQUIRES(
+ dilation_h > 0 && dilation_w > 0,
+ errors::InvalidArgument("Dilated rates should be larger than 0."));
+
+ return Status::OK();
+}
+
+Status ComputeConv2DDimension(const Conv2DParameters& params,
+ const Tensor& input, const Tensor& filter,
+ Conv2DDimensions* dimensions) {
+ // Check that 2D convolution input and filter have exactly 4 dimensions.
+ TF_REQUIRES(input.dims() == 4,
+ errors::InvalidArgument("input must be 4-dimensional",
+ input.shape().DebugString()));
+ TF_REQUIRES(filter.dims() == 4,
+ errors::InvalidArgument("filter must be 4-dimensional: ",
+ filter.shape().DebugString()));
+ for (int i = 0; i < 3; i++) {
+ TF_REQUIRES(
+ FastBoundsCheck(filter.dim_size(i), std::numeric_limits<int>::max()),
+ errors::InvalidArgument("filter too large"));
+ }
+
+ // The last dimension for input is in_depth. Check that it is the same as the
+ // filter's in_depth or it is evenly divisible by filter's in_depth.
+ const int64 in_depth_raw = GetTensorDim(input, params.data_format, 'C');
+ const int64 patch_depth_raw = filter.dim_size(2);
+ TF_REQUIRES(FastBoundsCheck(in_depth_raw, std::numeric_limits<int>::max()),
+ errors::InvalidArgument("Input depth too large"));
+ TF_REQUIRES(FastBoundsCheck(patch_depth_raw, std::numeric_limits<int>::max()),
+ errors::InvalidArgument("Patch depth too large"));
+ const int in_depth = static_cast<int>(in_depth_raw);
+ const int patch_depth = static_cast<int>(patch_depth_raw);
+ TF_REQUIRES(in_depth % patch_depth == 0,
+ errors::InvalidArgument(
+ "input depth must be evenly divisible by filter depth: ",
+ in_depth, " vs ", patch_depth));
+
+ // The last dimension for filter is out_depth.
+ const int out_depth = static_cast<int>(filter.dim_size(3));
+
+ // The second dimension for input is rows/height.
+ // The first dimension for filter is rows/height.
+ const int64 input_rows_raw = GetTensorDim(input, params.data_format, 'H');
+ TF_REQUIRES(FastBoundsCheck(input_rows_raw, std::numeric_limits<int>::max()),
+ errors::InvalidArgument("Input rows too large"));
+ const int input_rows = static_cast<int>(input_rows_raw);
+ const int filter_rows = static_cast<int>(filter.dim_size(0));
+
+ // The third dimension for input is columns/width.
+ // The second dimension for filter is columns/width.
+ const int64 input_cols_raw = GetTensorDim(input, params.data_format, 'W');
+ TF_REQUIRES(FastBoundsCheck(input_cols_raw, std::numeric_limits<int>::max()),
+ errors::InvalidArgument("Input cols too large"));
+ const int input_cols = static_cast<int>(input_cols_raw);
+ const int filter_cols = static_cast<int>(filter.dim_size(1));
+
+ // The first dimension for input is batch.
+ const int64 batch_raw = GetTensorDim(input, params.data_format, 'N');
+ TF_REQUIRES(FastBoundsCheck(batch_raw, std::numeric_limits<int>::max()),
+ errors::InvalidArgument("batch is too large"));
+ const int batch = static_cast<int>(batch_raw);
+
+ // Take the stride and dilation from the second and third dimensions only (we
+ // do not support striding or dilation on the batch or depth dimension).
+ const int stride_rows = GetTensorDim(params.strides, params.data_format, 'H');
+ const int stride_cols = GetTensorDim(params.strides, params.data_format, 'W');
+ const int dilation_rows =
+ GetTensorDim(params.dilations, params.data_format, 'H');
+ const int dilation_cols =
+ GetTensorDim(params.dilations, params.data_format, 'W');
+
+ // Compute windowed output sizes for rows and columns.
+ int64 out_rows = 0, out_cols = 0, pad_rows = 0, pad_cols = 0;
+ TF_RETURN_IF_ERROR(GetWindowedOutputSizeV2(
+ input_rows, filter_rows, dilation_rows, stride_rows, params.padding,
+ &out_rows, &pad_rows));
+ TF_RETURN_IF_ERROR(GetWindowedOutputSizeV2(
+ input_cols, filter_cols, dilation_cols, stride_cols, params.padding,
+ &out_cols, &pad_cols));
+
+ dimensions->batch = batch;
+ dimensions->input_rows = input_rows;
+ dimensions->input_cols = input_cols;
+ dimensions->in_depth = in_depth;
+ dimensions->filter_rows = filter_rows;
+ dimensions->filter_cols = filter_cols;
+ dimensions->patch_depth = patch_depth;
+ dimensions->out_depth = out_depth;
+ dimensions->stride_rows = stride_rows;
+ dimensions->stride_cols = stride_cols;
+ dimensions->dilation_rows = dilation_rows;
+ dimensions->dilation_cols = dilation_cols;
+ dimensions->out_rows = out_rows;
+ dimensions->out_cols = out_cols;
+ dimensions->pad_rows = pad_rows;
+ dimensions->pad_cols = pad_cols;
+
+ return Status::OK();
+}
+
+#undef TF_REQUIRES
+
template <typename Device, typename T>
class Conv2DOp : public BinaryOp<T> {
public:
explicit Conv2DOp(OpKernelConstruction* context) : BinaryOp<T>(context) {
- OP_REQUIRES_OK(context, context->GetAttr("dilations", &dilations_));
- OP_REQUIRES_OK(context, context->GetAttr("strides", &strides_));
- string data_format;
- OP_REQUIRES_OK(context, context->GetAttr("data_format", &data_format));
- OP_REQUIRES(context, FormatFromString(data_format, &data_format_),
- errors::InvalidArgument("Invalid data format"));
+ OP_REQUIRES_OK(context, InitConv2DParameters(context, ¶ms_));
+
OP_REQUIRES_OK(context, context->GetAttr("use_cudnn_on_gpu", &use_cudnn_));
use_cudnn_ &= CanUseCudnn();
cudnn_use_autotune_ = CudnnUseAutotune();
- OP_REQUIRES(context, dilations_.size() == 4,
- errors::InvalidArgument("Sliding window dilations field must "
- "specify 4 dimensions"));
- OP_REQUIRES(context, strides_.size() == 4,
- errors::InvalidArgument("Sliding window strides field must "
- "specify 4 dimensions"));
- const int64 stride_n = GetTensorDim(strides_, data_format_, 'N');
- const int64 stride_c = GetTensorDim(strides_, data_format_, 'C');
- const int64 stride_h = GetTensorDim(strides_, data_format_, 'H');
- const int64 stride_w = GetTensorDim(strides_, data_format_, 'W');
- OP_REQUIRES(
- context, stride_n == 1 && stride_c == 1,
- errors::InvalidArgument("Current implementation does not yet support "
- "strides in the batch and depth dimensions."));
- OP_REQUIRES(context, stride_h > 0 && stride_w > 0,
- errors::InvalidArgument(
- "Row and column strides should be larger than 0."));
-
- const int64 dilation_n = GetTensorDim(dilations_, data_format_, 'N');
- const int64 dilation_c = GetTensorDim(dilations_, data_format_, 'C');
- const int64 dilation_h = GetTensorDim(dilations_, data_format_, 'H');
- const int64 dilation_w = GetTensorDim(dilations_, data_format_, 'W');
- OP_REQUIRES(context, dilation_n == 1 && dilation_c == 1,
- errors::InvalidArgument(
- "Current implementation does not yet support "
- "dilations in the batch and depth dimensions."));
- OP_REQUIRES(
- context, dilation_h > 0 && dilation_w > 0,
- errors::InvalidArgument("Dilated rates should be larger than 0."));
- OP_REQUIRES_OK(context, context->GetAttr("padding", &padding_));
}
void Compute(OpKernelContext* context) override {
// Input tensor is of the following dimensions:
// [ batch, in_rows, in_cols, in_depth ]
-
const Tensor& input = context->input(0);
// Input filter is of the following dimensions:
// [ filter_rows, filter_cols, in_depth, out_depth]
const Tensor& filter = context->input(1);
- // For 2D convolution, there should be 4 dimensions.
- OP_REQUIRES(context, input.dims() == 4,
- errors::InvalidArgument("input must be 4-dimensional",
- input.shape().DebugString()));
- OP_REQUIRES(context, filter.dims() == 4,
- errors::InvalidArgument("filter must be 4-dimensional: ",
- filter.shape().DebugString()));
+ Conv2DDimensions dimensions;
+ OP_REQUIRES_OK(context,
+ ComputeConv2DDimension(params_, input, filter, &dimensions));
- for (int i = 0; i < 3; i++) {
- OP_REQUIRES(
- context,
- FastBoundsCheck(filter.dim_size(i), std::numeric_limits<int>::max()),
- errors::InvalidArgument("filter too large"));
- }
-
- // The last dimension for input is in_depth. It must be the same as the
- // filter's in_depth or be evenly divisible by filter's in_depth.
- const int64 in_depth = GetTensorDim(input, data_format_, 'C');
- const int64 patch_depth = filter.dim_size(2);
- OP_REQUIRES(context, in_depth % patch_depth == 0,
- errors::InvalidArgument(
- "input depth must be evenly divisible by filter depth: ",
- in_depth, " vs ", patch_depth));
-
- // The last dimension for filter is out_depth.
- const int out_depth = static_cast<int>(filter.dim_size(3));
-
- // The second dimension for input is rows/height.
- // The first dimension for filter is rows/height.
- const int64 input_rows_raw = GetTensorDim(input, data_format_, 'H');
- OP_REQUIRES(
- context,
- FastBoundsCheck(input_rows_raw, std::numeric_limits<int>::max()),
- errors::InvalidArgument("Input rows too large"));
- const int input_rows = static_cast<int>(input_rows_raw);
- const int filter_rows = static_cast<int>(filter.dim_size(0));
-
- // The third dimension for input is columns/width.
- // The second dimension for filter is columns/width.
- const int64 input_cols_raw = GetTensorDim(input, data_format_, 'W');
- OP_REQUIRES(
- context,
- FastBoundsCheck(input_cols_raw, std::numeric_limits<int>::max()),
- errors::InvalidArgument("Input cols too large"));
- const int input_cols = static_cast<int>(input_cols_raw);
- const int filter_cols = static_cast<int>(filter.dim_size(1));
-
- // The first dimension for input is batch.
- const int64 batch_raw = GetTensorDim(input, data_format_, 'N');
- OP_REQUIRES(context,
- FastBoundsCheck(batch_raw, std::numeric_limits<int>::max()),
- errors::InvalidArgument("batch is too large"));
- const int batch = static_cast<int>(batch_raw);
-
- // For now we take the stride and dilation from the second and third
- // dimensions only (we do not support striding or dilation on the batch or
- // depth dimension).
- const int stride_rows = GetTensorDim(strides_, data_format_, 'H');
- const int stride_cols = GetTensorDim(strides_, data_format_, 'W');
-
- const int dilation_rows = GetTensorDim(dilations_, data_format_, 'H');
- const int dilation_cols = GetTensorDim(dilations_, data_format_, 'W');
-
- int64 out_rows = 0, out_cols = 0, pad_rows = 0, pad_cols = 0;
- OP_REQUIRES_OK(context, GetWindowedOutputSizeV2(
- input_rows, filter_rows, dilation_rows,
- stride_rows, padding_, &out_rows, &pad_rows));
- OP_REQUIRES_OK(context, GetWindowedOutputSizeV2(
- input_cols, filter_cols, dilation_cols,
- stride_cols, padding_, &out_cols, &pad_cols));
- TensorShape out_shape =
- ShapeFromFormat(data_format_, batch, out_rows, out_cols, out_depth);
+ TensorShape out_shape = ShapeFromFormat(
+ params_.data_format, dimensions.batch, dimensions.out_rows,
+ dimensions.out_cols, dimensions.out_depth);
// Output tensor is of the following dimensions:
// [ in_batch, out_rows, out_cols, out_depth ]
Tensor* output = nullptr;
OP_REQUIRES_OK(context, context->allocate_output(0, out_shape, &output));
- VLOG(2) << "Conv2D: in_depth = " << in_depth
- << ", patch_depth = " << patch_depth
- << ", input_cols = " << input_cols
- << ", filter_cols = " << filter_cols
- << ", input_rows = " << input_rows
- << ", filter_rows = " << filter_rows
- << ", stride_rows = " << stride_rows
- << ", stride_cols = " << stride_cols
- << ", dilation_rows = " << dilation_rows
- << ", dilation_cols = " << dilation_cols
- << ", out_depth = " << out_depth;
+ VLOG(2) << "Conv2D: in_depth = " << dimensions.in_depth
+ << ", patch_depth = " << dimensions.patch_depth
+ << ", input_cols = " << dimensions.input_cols
+ << ", filter_cols = " << dimensions.filter_cols
+ << ", input_rows = " << dimensions.input_rows
+ << ", filter_rows = " << dimensions.filter_rows
+ << ", stride_rows = " << dimensions.stride_rows
+ << ", stride_cols = " << dimensions.stride_cols
+ << ", dilation_rows = " << dimensions.dilation_rows
+ << ", dilation_cols = " << dimensions.dilation_cols
+ << ", out_depth = " << dimensions.out_depth;
// If there is nothing to compute, return.
if (out_shape.num_elements() == 0) {
@@ -416,36 +464,41 @@
#ifdef TENSORFLOW_USE_LIBXSMM_CONVOLUTIONS
if (LaunchXsmmConvOp<Device, T>::Run(
- context, input, filter, batch, input_rows, input_cols, in_depth,
- filter_rows, filter_cols, pad_rows, pad_cols, out_rows, out_cols,
- out_depth, dilation_rows, dilation_cols, stride_rows, stride_cols,
- output, data_format_)) {
+ context, input, filter, dimensions.batch, dimensions.input_rows,
+ dimensions.input_cols, dimensions.in_depth, dimensions.filter_rows,
+ dimensions.filter_cols, dimensions.pad_rows, dimensions.pad_cols,
+ dimensions.out_rows, dimensions.out_cols, dimensions.out_depth,
+ dimensions.dilation_rows, dimensions.dilation_cols,
+ dimensions.stride_rows, dimensions.stride_cols, output,
+ params_.data_format)) {
return;
}
#endif
if (LaunchDeepConvOp<Device, T>::Run(
- context, input, filter, batch, input_rows, input_cols, in_depth,
- filter_rows, filter_cols, pad_rows, pad_cols, out_rows, out_cols,
- out_depth, dilation_rows, dilation_cols, stride_rows, stride_cols,
- output, data_format_)) {
+ context, input, filter, dimensions.batch, dimensions.input_rows,
+ dimensions.input_cols, dimensions.in_depth, dimensions.filter_rows,
+ dimensions.filter_cols, dimensions.pad_rows, dimensions.pad_cols,
+ dimensions.out_rows, dimensions.out_cols, dimensions.out_depth,
+ dimensions.dilation_rows, dimensions.dilation_cols,
+ dimensions.stride_rows, dimensions.stride_cols, output,
+ params_.data_format)) {
return;
}
launcher_(context, use_cudnn_, cudnn_use_autotune_, input, filter,
- dilation_rows, dilation_cols, stride_rows, stride_cols, padding_,
- output, data_format_);
+ dimensions.dilation_rows, dimensions.dilation_cols,
+ dimensions.stride_rows, dimensions.stride_cols, params_.padding,
+ output, params_.data_format);
}
private:
- std::vector<int32> dilations_;
- std::vector<int32> strides_;
+ Conv2DParameters params_;
bool use_cudnn_;
- Padding padding_;
- TensorFormat data_format_;
- LaunchConv2DOp<Device, T> launcher_;
bool cudnn_use_autotune_;
+ LaunchConv2DOp<Device, T> launcher_;
+
TF_DISALLOW_COPY_AND_ASSIGN(Conv2DOp);
};
diff --git a/tensorflow/core/kernels/conv_ops.h b/tensorflow/core/kernels/conv_ops.h
index adf4601..7ec878e 100644
--- a/tensorflow/core/kernels/conv_ops.h
+++ b/tensorflow/core/kernels/conv_ops.h
@@ -66,6 +66,50 @@
string DebugString() { return "Im2ColBufferResource"; }
};
+// Convolution parameters specified by Op attributes.
+struct Conv2DParameters {
+ std::vector<int32> dilations;
+ std::vector<int32> strides;
+ Padding padding;
+ TensorFormat data_format;
+};
+
+// Convolution dimensions inferred from parameters, input and filter tensors.
+struct Conv2DDimensions {
+ int batch;
+ int input_rows;
+ int input_cols;
+ int in_depth;
+
+ int filter_rows;
+ int filter_cols;
+ int patch_depth;
+ int out_depth;
+
+ int stride_rows;
+ int stride_cols;
+
+ int dilation_rows;
+ int dilation_cols;
+
+ int64 out_rows;
+ int64 out_cols;
+ int64 pad_rows;
+ int64 pad_cols;
+};
+
+// Initializes and validates Conv2D parameters configured by OpKernel
+// attributes.
+Status InitConv2DParameters(const OpKernelConstruction* context,
+ Conv2DParameters* params);
+
+// Computes and validates convolutions dimensions from Conv2D parameters. If
+// parameters are valid, dimensions will be updated with derived convolution
+// dimensions, otherwise error will be returned.
+Status ComputeConv2DDimension(const Conv2DParameters& params,
+ const Tensor& input, const Tensor& filter,
+ Conv2DDimensions* dimensions);
+
} // namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_CONV_OPS_H_
diff --git a/tensorflow/core/kernels/cwise_op_gpu_xdivy.cu.cc b/tensorflow/core/kernels/cwise_op_gpu_xdivy.cu.cc
new file mode 100644
index 0000000..e4b21a6
--- /dev/null
+++ b/tensorflow/core/kernels/cwise_op_gpu_xdivy.cu.cc
@@ -0,0 +1,26 @@
+/* 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.
+==============================================================================*/
+
+#if GOOGLE_CUDA
+
+#include "tensorflow/core/kernels/cwise_ops_gpu_common.cu.h"
+
+namespace tensorflow {
+namespace functor {
+DEFINE_BINARY5(xdivy, Eigen::half, float, double, complex64, complex128);
+} // namespace functor
+} // namespace tensorflow
+
+#endif // GOOGLE_CUDA
diff --git a/tensorflow/core/kernels/cwise_op_gpu_xlogy.cu.cc b/tensorflow/core/kernels/cwise_op_gpu_xlogy.cu.cc
new file mode 100644
index 0000000..1e1b5a4
--- /dev/null
+++ b/tensorflow/core/kernels/cwise_op_gpu_xlogy.cu.cc
@@ -0,0 +1,26 @@
+/* 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
+
+#include "tensorflow/core/kernels/cwise_ops_gpu_common.cu.h"
+
+namespace tensorflow {
+namespace functor {
+DEFINE_BINARY5(xlogy, Eigen::half, float, double, complex64, complex128);
+} // namespace functor
+} // namespace tensorflow
+
+#endif // GOOGLE_CUDA
diff --git a/tensorflow/core/kernels/cwise_op_xdivy.cc b/tensorflow/core/kernels/cwise_op_xdivy.cc
new file mode 100644
index 0000000..6a6aec5
--- /dev/null
+++ b/tensorflow/core/kernels/cwise_op_xdivy.cc
@@ -0,0 +1,38 @@
+/* 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/cwise_ops_common.h"
+
+namespace tensorflow {
+REGISTER5(BinaryOp, CPU, "Xdivy", functor::xdivy, float, Eigen::half, double,
+ complex64, complex128);
+
+#if TENSORFLOW_USE_SYCL
+#define REGISTER_SYCL_KERNEL(TYPE) \
+ REGISTER_KERNEL_BUILDER( \
+ Name("Xdivy").Device(DEVICE_SYCL).TypeConstraint<TYPE>("T"), \
+ BinaryOp<SYCLDevice, functor::xdivy<TYPE>>);
+REGISTER_SYCL_KERNEL(float);
+REGISTER_SYCL_KERNEL(double);
+#undef REGISTER_SYCL_KERNEL
+
+#endif // TENSORFLOW_USE_SYCL
+
+#if GOOGLE_CUDA
+REGISTER5(BinaryOp, GPU, "Xdivy", functor::xdivy, float, Eigen::half, double,
+ complex64, complex128);
+#endif // GOOGLE_CUDA
+
+} // namespace tensorflow
diff --git a/tensorflow/core/kernels/cwise_op_xlogy.cc b/tensorflow/core/kernels/cwise_op_xlogy.cc
new file mode 100644
index 0000000..e71a910
--- /dev/null
+++ b/tensorflow/core/kernels/cwise_op_xlogy.cc
@@ -0,0 +1,41 @@
+/* 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/cwise_ops_common.h"
+
+namespace tensorflow {
+REGISTER5(BinaryOp, CPU, "Xlogy", functor::xlogy, float, Eigen::half, double,
+ complex64, complex128);
+
+#if TENSORFLOW_USE_SYCL
+#define REGISTER_SYCL_KERNEL(TYPE) \
+ REGISTER_KERNEL_BUILDER( \
+ Name("Xlogy").Device(DEVICE_SYCL).TypeConstraint<TYPE>("T"), \
+ BinaryOp<SYCLDevice, functor::xlogy<TYPE>>);
+REGISTER_SYCL_KERNEL(Eigen::half);
+REGISTER_SYCL_KERNEL(float);
+REGISTER_SYCL_KERNEL(double);
+REGISTER_SYCL_KERNEL(complex64);
+REGISTER_SYCL_KERNEL(complex128);
+#undef REGISTER_SYCL_KERNEL
+
+#endif // TENSORFLOW_USE_SYCL
+
+#if GOOGLE_CUDA
+REGISTER5(BinaryOp, GPU, "Xlogy", functor::xlogy, float, Eigen::half, double,
+ complex64, complex128);
+#endif // GOOGLE_CUDA
+
+} // namespace tensorflow
diff --git a/tensorflow/core/kernels/cwise_ops.h b/tensorflow/core/kernels/cwise_ops.h
index 22eb66e..66ba827 100644
--- a/tensorflow/core/kernels/cwise_ops.h
+++ b/tensorflow/core/kernels/cwise_ops.h
@@ -471,6 +471,45 @@
enum { Cost = Eigen::NumTraits<Scalar>::AddCost, PacketAccess = true };
};
+// TODO(srvasude): Add packet versions of this operation.
+template <typename Scalar>
+struct xlogy_op {
+ EIGEN_EMPTY_STRUCT_CTOR(xlogy_op)
+ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Scalar
+ operator()(const Scalar& x, const Scalar& y) const {
+ if (x == Scalar(0.)) {
+ return Scalar(0.);
+ }
+ return x * numext::log(y);
+ }
+};
+
+template <typename Scalar>
+struct functor_traits<xlogy_op<Scalar>> {
+ enum {
+ Cost = (sizeof(Scalar) == 4 ? 40 : 85) + Eigen::NumTraits<Scalar>::MulCost,
+ PacketAccess = false
+ };
+};
+
+template <typename Scalar>
+// TODO(srvasude): Add packet versions of this operation.
+struct xdivy_op {
+ EIGEN_EMPTY_STRUCT_CTOR(xdivy_op)
+ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Scalar
+ operator()(const Scalar& x, const Scalar& y) const {
+ if (x == Scalar(0.)) {
+ return Scalar(0.);
+ }
+ return x / y;
+ }
+};
+
+template <typename Scalar>
+struct functor_traits<xdivy_op<Scalar>> {
+ enum { Cost = Eigen::NumTraits<Scalar>::MulCost, PacketAccess = false };
+};
+
} // end namespace internal
} // end namespace Eigen
@@ -830,6 +869,12 @@
Eigen::internal::scalar_difference_op<T>>> {};
template <typename T>
+struct xdivy : base<T, Eigen::internal::xdivy_op<T>> {};
+
+template <typename T>
+struct xlogy : base<T, Eigen::internal::xlogy_op<T>> {};
+
+template <typename T>
struct less : base<T, Eigen::internal::less<T>, bool> {};
template <typename T>
diff --git a/tensorflow/core/kernels/data/BUILD b/tensorflow/core/kernels/data/BUILD
index 87efdff..6333853 100644
--- a/tensorflow/core/kernels/data/BUILD
+++ b/tensorflow/core/kernels/data/BUILD
@@ -765,6 +765,7 @@
":window_dataset_op",
":writer_ops",
":zip_dataset_op",
+ "//tensorflow/core/kernels/data/experimental:dataset_kernels",
],
)
diff --git a/tensorflow/core/kernels/data/experimental/BUILD b/tensorflow/core/kernels/data/experimental/BUILD
new file mode 100644
index 0000000..43406db
--- /dev/null
+++ b/tensorflow/core/kernels/data/experimental/BUILD
@@ -0,0 +1,139 @@
+# Description:
+# Contains experimental kernels for datasets and iterators.
+package(default_visibility = ["//tensorflow:internal"])
+
+licenses(["notice"]) # Apache 2.0
+
+exports_files(["LICENSE"])
+
+load(
+ "//tensorflow:tensorflow.bzl",
+ "tf_kernel_library",
+)
+
+cc_library(
+ name = "indexed_dataset_headers",
+ hdrs = ["indexed_dataset.h"],
+ deps = [
+ "//tensorflow/core:framework",
+ "//third_party/eigen3",
+ ],
+)
+
+tf_kernel_library(
+ name = "indexed_dataset",
+ srcs = [
+ "identity_indexed_dataset.cc",
+ "indexed_dataset.cc",
+ ],
+ deps = [
+ ":indexed_dataset_headers",
+ "//tensorflow/core:experimental_dataset_ops_op_lib",
+ "//tensorflow/core:framework",
+ "//tensorflow/core:lib",
+ "//third_party/eigen3",
+ ],
+)
+
+tf_kernel_library(
+ name = "prefetching_kernels",
+ srcs = ["prefetching_kernels.cc"],
+ deps = [
+ "//tensorflow/core:core_cpu_internal",
+ "//tensorflow/core:experimental_dataset_ops_op_lib",
+ "//tensorflow/core:framework",
+ "//tensorflow/core:lib",
+ "//tensorflow/core:lib_internal",
+ ],
+)
+
+tf_kernel_library(
+ name = "directed_interleave_dataset_op",
+ srcs = ["directed_interleave_dataset_op.cc"],
+ deps = [
+ "//tensorflow/core:experimental_dataset_ops_op_lib",
+ "//tensorflow/core:framework",
+ "//tensorflow/core:lib",
+ "//third_party/eigen3",
+ ],
+)
+
+tf_kernel_library(
+ name = "csv_dataset_op",
+ srcs = ["csv_dataset_op.cc"],
+ deps = [
+ "//tensorflow/core:experimental_dataset_ops_op_lib",
+ "//tensorflow/core:framework",
+ "//tensorflow/core:lib",
+ "//tensorflow/core:lib_internal",
+ ],
+)
+
+tf_kernel_library(
+ name = "ignore_errors_dataset_op",
+ srcs = ["ignore_errors_dataset_op.cc"],
+ deps = [
+ "//tensorflow/core:experimental_dataset_ops_op_lib",
+ "//tensorflow/core:framework",
+ "//third_party/eigen3",
+ ],
+)
+
+tf_kernel_library(
+ name = "lmdb_dataset_op",
+ srcs = ["lmdb_dataset_op.cc"],
+ deps = [
+ "//tensorflow/core:experimental_dataset_ops_op_lib",
+ "//tensorflow/core:framework",
+ "//tensorflow/core:lib",
+ "//third_party/eigen3",
+ "@lmdb",
+ ],
+)
+
+tf_kernel_library(
+ name = "threadpool_dataset_op",
+ srcs = ["threadpool_dataset_op.cc"],
+ deps = [
+ "//tensorflow/core:experimental_dataset_ops_op_lib",
+ "//tensorflow/core:framework",
+ "//tensorflow/core:lib",
+ "//third_party/eigen3",
+ ],
+)
+
+tf_kernel_library(
+ name = "unique_dataset_op",
+ srcs = ["unique_dataset_op.cc"],
+ deps = [
+ "//tensorflow/core:experimental_dataset_ops_op_lib",
+ "//tensorflow/core:framework",
+ "//tensorflow/core:lib",
+ "//third_party/eigen3",
+ ],
+)
+
+tf_kernel_library(
+ name = "assert_next_dataset_op",
+ srcs = ["assert_next_dataset_op.cc"],
+ deps = [
+ "//tensorflow/core:experimental_dataset_ops_op_lib",
+ "//tensorflow/core:framework",
+ "//third_party/eigen3",
+ ],
+)
+
+tf_kernel_library(
+ name = "dataset_kernels",
+ deps = [
+ ":assert_next_dataset_op",
+ ":csv_dataset_op",
+ ":directed_interleave_dataset_op",
+ ":ignore_errors_dataset_op",
+ ":indexed_dataset",
+ ":lmdb_dataset_op",
+ ":prefetching_kernels",
+ ":threadpool_dataset_op",
+ ":unique_dataset_op",
+ ],
+)
diff --git a/tensorflow/contrib/data/kernels/assert_next_dataset_op.cc b/tensorflow/core/kernels/data/experimental/assert_next_dataset_op.cc
similarity index 97%
rename from tensorflow/contrib/data/kernels/assert_next_dataset_op.cc
rename to tensorflow/core/kernels/data/experimental/assert_next_dataset_op.cc
index c19a609..3511cca 100644
--- a/tensorflow/contrib/data/kernels/assert_next_dataset_op.cc
+++ b/tensorflow/core/kernels/data/experimental/assert_next_dataset_op.cc
@@ -147,8 +147,9 @@
std::vector<PartialTensorShape> output_shapes_;
};
-REGISTER_KERNEL_BUILDER(Name("AssertNextDataset").Device(DEVICE_CPU),
- AssertNextDatasetOp);
+REGISTER_KERNEL_BUILDER(
+ Name("ExperimentalAssertNextDataset").Device(DEVICE_CPU),
+ AssertNextDatasetOp);
} // namespace
} // namespace data
diff --git a/tensorflow/contrib/data/kernels/csv_dataset_op.cc b/tensorflow/core/kernels/data/experimental/csv_dataset_op.cc
similarity index 99%
rename from tensorflow/contrib/data/kernels/csv_dataset_op.cc
rename to tensorflow/core/kernels/data/experimental/csv_dataset_op.cc
index 21ec50f..7451ca4 100644
--- a/tensorflow/contrib/data/kernels/csv_dataset_op.cc
+++ b/tensorflow/core/kernels/data/experimental/csv_dataset_op.cc
@@ -852,7 +852,8 @@
}; // class CSVDatasetOp
// Register the kernel implementation for CSVDataset.
-REGISTER_KERNEL_BUILDER(Name("CSVDataset").Device(DEVICE_CPU), CSVDatasetOp);
+REGISTER_KERNEL_BUILDER(Name("ExperimentalCSVDataset").Device(DEVICE_CPU),
+ CSVDatasetOp);
} // namespace
} // namespace data
diff --git a/tensorflow/contrib/data/kernels/directed_interleave_dataset_op.cc b/tensorflow/core/kernels/data/experimental/directed_interleave_dataset_op.cc
similarity index 98%
rename from tensorflow/contrib/data/kernels/directed_interleave_dataset_op.cc
rename to tensorflow/core/kernels/data/experimental/directed_interleave_dataset_op.cc
index a532162..c47a909 100644
--- a/tensorflow/contrib/data/kernels/directed_interleave_dataset_op.cc
+++ b/tensorflow/core/kernels/data/experimental/directed_interleave_dataset_op.cc
@@ -272,8 +272,9 @@
};
};
-REGISTER_KERNEL_BUILDER(Name("DirectedInterleaveDataset").Device(DEVICE_CPU),
- DirectedInterleaveDatasetOp);
+REGISTER_KERNEL_BUILDER(
+ Name("ExperimentalDirectedInterleaveDataset").Device(DEVICE_CPU),
+ DirectedInterleaveDatasetOp);
} // namespace
} // namespace data
diff --git a/tensorflow/contrib/data/kernels/identity_indexed_dataset.cc b/tensorflow/core/kernels/data/experimental/identity_indexed_dataset.cc
similarity index 95%
rename from tensorflow/contrib/data/kernels/identity_indexed_dataset.cc
rename to tensorflow/core/kernels/data/experimental/identity_indexed_dataset.cc
index c3cb45d..2141f11 100644
--- a/tensorflow/contrib/data/kernels/identity_indexed_dataset.cc
+++ b/tensorflow/core/kernels/data/experimental/identity_indexed_dataset.cc
@@ -13,7 +13,7 @@
limitations under the License.
==============================================================================*/
-#include "tensorflow/contrib/data/kernels/indexed_dataset.h"
+#include "tensorflow/core/kernels/data/experimental/indexed_dataset.h"
#include "tensorflow/core/lib/core/errors.h"
namespace tensorflow {
@@ -147,8 +147,9 @@
};
};
-REGISTER_KERNEL_BUILDER(Name("IdentityIndexedDataset").Device(DEVICE_CPU),
- IdentityIndexedDatasetOp);
+REGISTER_KERNEL_BUILDER(
+ Name("ExperimentalIdentityIndexedDataset").Device(DEVICE_CPU),
+ IdentityIndexedDatasetOp);
} // namespace
} // namespace data
diff --git a/tensorflow/contrib/data/kernels/ignore_errors_dataset_op.cc b/tensorflow/core/kernels/data/experimental/ignore_errors_dataset_op.cc
similarity index 96%
rename from tensorflow/contrib/data/kernels/ignore_errors_dataset_op.cc
rename to tensorflow/core/kernels/data/experimental/ignore_errors_dataset_op.cc
index beec344..b34377c 100644
--- a/tensorflow/contrib/data/kernels/ignore_errors_dataset_op.cc
+++ b/tensorflow/core/kernels/data/experimental/ignore_errors_dataset_op.cc
@@ -15,7 +15,6 @@
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/framework/partial_tensor_shape.h"
#include "tensorflow/core/framework/tensor.h"
-#include "tensorflow/core/lib/random/random.h"
namespace tensorflow {
namespace data {
@@ -133,8 +132,9 @@
};
};
-REGISTER_KERNEL_BUILDER(Name("IgnoreErrorsDataset").Device(DEVICE_CPU),
- IgnoreErrorsDatasetOp);
+REGISTER_KERNEL_BUILDER(
+ Name("ExperimentalIgnoreErrorsDataset").Device(DEVICE_CPU),
+ IgnoreErrorsDatasetOp);
} // namespace
} // namespace data
diff --git a/tensorflow/contrib/data/kernels/indexed_dataset.cc b/tensorflow/core/kernels/data/experimental/indexed_dataset.cc
similarity index 96%
rename from tensorflow/contrib/data/kernels/indexed_dataset.cc
rename to tensorflow/core/kernels/data/experimental/indexed_dataset.cc
index ced8ab0..75ea462 100644
--- a/tensorflow/contrib/data/kernels/indexed_dataset.cc
+++ b/tensorflow/core/kernels/data/experimental/indexed_dataset.cc
@@ -12,7 +12,7 @@
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
-#include "tensorflow/contrib/data/kernels/indexed_dataset.h"
+#include "tensorflow/core/kernels/data/experimental/indexed_dataset.h"
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/framework/tensor_shape.h"
@@ -361,12 +361,14 @@
};
REGISTER_KERNEL_BUILDER(
- Name("MaterializedIndexDatasetHandle").Device(DEVICE_CPU),
+ Name("ExperimentalMaterializedIndexDatasetHandle").Device(DEVICE_CPU),
MaterializedHandleOp);
-REGISTER_KERNEL_BUILDER(Name("IndexedDatasetMaterialize").Device(DEVICE_CPU),
- MaterializeDatasetOp);
-REGISTER_KERNEL_BUILDER(Name("IndexedDatasetGet").Device(DEVICE_CPU),
- IndexedDatasetGet);
+REGISTER_KERNEL_BUILDER(
+ Name("ExperimentalIndexedDatasetMaterialize").Device(DEVICE_CPU),
+ MaterializeDatasetOp);
+REGISTER_KERNEL_BUILDER(
+ Name("ExperimentalIndexedDatasetGet").Device(DEVICE_CPU),
+ IndexedDatasetGet);
} // namespace
} // namespace data
diff --git a/tensorflow/contrib/data/kernels/indexed_dataset.h b/tensorflow/core/kernels/data/experimental/indexed_dataset.h
similarity index 95%
rename from tensorflow/contrib/data/kernels/indexed_dataset.h
rename to tensorflow/core/kernels/data/experimental/indexed_dataset.h
index 7aa2d3f..27a8360 100644
--- a/tensorflow/contrib/data/kernels/indexed_dataset.h
+++ b/tensorflow/core/kernels/data/experimental/indexed_dataset.h
@@ -12,8 +12,8 @@
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
-#ifndef TENSORFLOW_CONTRIB_DATA_KERNELS_INDEXED_DATASET_H_
-#define TENSORFLOW_CONTRIB_DATA_KERNELS_INDEXED_DATASET_H_
+#ifndef TENSORFLOW_CORE_KERNELS_DATA_EXPERIMENTAL_INDEXED_DATASET_H_
+#define TENSORFLOW_CORE_KERNELS_DATA_EXPERIMENTAL_INDEXED_DATASET_H_
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/framework/op_kernel.h"
@@ -116,4 +116,4 @@
} // namespace data
} // namespace tensorflow
-#endif // TENSORFLOW_CONTRIB_DATA_KERNELS_INDEXED_DATASET_H_
+#endif // TENSORFLOW_CORE_KERNELS_DATA_EXPERIMENTAL_INDEXED_DATASET_H_
diff --git a/tensorflow/contrib/data/kernels/lmdb_dataset_op.cc b/tensorflow/core/kernels/data/experimental/lmdb_dataset_op.cc
similarity index 98%
rename from tensorflow/contrib/data/kernels/lmdb_dataset_op.cc
rename to tensorflow/core/kernels/data/experimental/lmdb_dataset_op.cc
index d233c1f..8a88d32 100644
--- a/tensorflow/contrib/data/kernels/lmdb_dataset_op.cc
+++ b/tensorflow/core/kernels/data/experimental/lmdb_dataset_op.cc
@@ -210,7 +210,8 @@
};
};
-REGISTER_KERNEL_BUILDER(Name("LMDBDataset").Device(DEVICE_CPU), LMDBDatasetOp);
+REGISTER_KERNEL_BUILDER(Name("ExperimentalLMDBDataset").Device(DEVICE_CPU),
+ LMDBDatasetOp);
} // namespace
} // namespace data
diff --git a/tensorflow/contrib/data/kernels/prefetching_kernels.cc b/tensorflow/core/kernels/data/experimental/prefetching_kernels.cc
similarity index 95%
rename from tensorflow/contrib/data/kernels/prefetching_kernels.cc
rename to tensorflow/core/kernels/data/experimental/prefetching_kernels.cc
index 96f1dd0..2c6179d 100644
--- a/tensorflow/contrib/data/kernels/prefetching_kernels.cc
+++ b/tensorflow/core/kernels/data/experimental/prefetching_kernels.cc
@@ -338,20 +338,20 @@
DataTypeVector output_types_;
};
-REGISTER_KERNEL_BUILDER(Name("FunctionBufferingResource")
+REGISTER_KERNEL_BUILDER(Name("ExperimentalFunctionBufferingResource")
.Device(DEVICE_CPU)
.HostMemory("resource")
.HostMemory("string_arg")
.HostMemory("target_device"),
FunctionBufferResourceHandleOp);
-REGISTER_KERNEL_BUILDER(Name("FunctionBufferingResource")
+REGISTER_KERNEL_BUILDER(Name("ExperimentalFunctionBufferingResource")
.Device(DEVICE_GPU)
.HostMemory("resource")
.HostMemory("string_arg")
.HostMemory("target_device"),
FunctionBufferResourceHandleOp);
#if TENSORFLOW_USE_SYCL
-REGISTER_KERNEL_BUILDER(Name("FunctionBufferingResource")
+REGISTER_KERNEL_BUILDER(Name("ExperimentalFunctionBufferingResource")
.Device(DEVICE_SYCL)
.HostMemory("resource")
.HostMemory("string_arg")
@@ -403,16 +403,16 @@
}
};
-REGISTER_KERNEL_BUILDER(Name("FunctionBufferingResourceGetNext")
+REGISTER_KERNEL_BUILDER(Name("ExperimentalFunctionBufferingResourceGetNext")
.Device(DEVICE_CPU)
.HostMemory("function_buffer_resource"),
FunctionBufferingResourceGetNextOp);
-REGISTER_KERNEL_BUILDER(Name("FunctionBufferingResourceGetNext")
+REGISTER_KERNEL_BUILDER(Name("ExperimentalFunctionBufferingResourceGetNext")
.Device(DEVICE_GPU)
.HostMemory("function_buffer_resource"),
FunctionBufferingResourceGetNextOp);
#if TENSORFLOW_USE_SYCL
-REGISTER_KERNEL_BUILDER(Name("FunctionBufferingResourceGetNext")
+REGISTER_KERNEL_BUILDER(Name("ExperimentalFunctionBufferingResourceGetNext")
.Device(DEVICE_SYCL)
.HostMemory("function_buffer_resource"),
FunctionBufferingResourceGetNextOp);
@@ -440,16 +440,16 @@
}
};
-REGISTER_KERNEL_BUILDER(Name("FunctionBufferingResourceReset")
+REGISTER_KERNEL_BUILDER(Name("ExperimentalFunctionBufferingResourceReset")
.Device(DEVICE_CPU)
.HostMemory("function_buffer_resource"),
FunctionBufferingResourceResetOp);
-REGISTER_KERNEL_BUILDER(Name("FunctionBufferingResourceReset")
+REGISTER_KERNEL_BUILDER(Name("ExperimentalFunctionBufferingResourceReset")
.Device(DEVICE_GPU)
.HostMemory("function_buffer_resource"),
FunctionBufferingResourceResetOp);
#if TENSORFLOW_USE_SYCL
-REGISTER_KERNEL_BUILDER(Name("FunctionBufferingResourceReset")
+REGISTER_KERNEL_BUILDER(Name("ExperimentalFunctionBufferingResourceReset")
.Device(DEVICE_SYCL)
.HostMemory("function_buffer_resource"),
FunctionBufferingResourceResetOp);
@@ -473,8 +473,9 @@
}
};
-REGISTER_KERNEL_BUILDER(Name("IteratorGetDevice").Device(DEVICE_CPU),
- IteratorGetDeviceOp);
+REGISTER_KERNEL_BUILDER(
+ Name("ExperimentalIteratorGetDevice").Device(DEVICE_CPU),
+ IteratorGetDeviceOp);
} // namespace
} // namespace data
diff --git a/tensorflow/contrib/data/kernels/threadpool_dataset_op.cc b/tensorflow/core/kernels/data/experimental/threadpool_dataset_op.cc
similarity index 97%
rename from tensorflow/contrib/data/kernels/threadpool_dataset_op.cc
rename to tensorflow/core/kernels/data/experimental/threadpool_dataset_op.cc
index 30fa97a..c80493d 100644
--- a/tensorflow/contrib/data/kernels/threadpool_dataset_op.cc
+++ b/tensorflow/core/kernels/data/experimental/threadpool_dataset_op.cc
@@ -209,10 +209,11 @@
};
};
-REGISTER_KERNEL_BUILDER(Name("ThreadPoolHandle").Device(DEVICE_CPU),
+REGISTER_KERNEL_BUILDER(Name("ExperimentalThreadPoolHandle").Device(DEVICE_CPU),
ThreadPoolHandleOp);
-REGISTER_KERNEL_BUILDER(Name("ThreadPoolDataset").Device(DEVICE_CPU),
- ThreadPoolDatasetOp);
+REGISTER_KERNEL_BUILDER(
+ Name("ExperimentalThreadPoolDataset").Device(DEVICE_CPU),
+ ThreadPoolDatasetOp);
} // namespace
} // namespace data
diff --git a/tensorflow/contrib/data/kernels/unique_dataset_op.cc b/tensorflow/core/kernels/data/experimental/unique_dataset_op.cc
similarity index 96%
rename from tensorflow/contrib/data/kernels/unique_dataset_op.cc
rename to tensorflow/core/kernels/data/experimental/unique_dataset_op.cc
index 57fc569..cd612e0 100644
--- a/tensorflow/contrib/data/kernels/unique_dataset_op.cc
+++ b/tensorflow/core/kernels/data/experimental/unique_dataset_op.cc
@@ -199,8 +199,9 @@
HANDLE_TYPE(DT_INT64);
HANDLE_TYPE(DT_STRING);
default:
- LOG(FATAL) << "UniqueDataset unhandled data type: "
- << DataTypeString(lhs.dtype());
+ DCHECK(false) << "UniqueDataset unhandled data type: "
+ << DataTypeString(lhs.dtype());
+ return false;
}
}
};
@@ -215,7 +216,7 @@
};
};
-REGISTER_KERNEL_BUILDER(Name("UniqueDataset").Device(DEVICE_CPU),
+REGISTER_KERNEL_BUILDER(Name("ExperimentalUniqueDataset").Device(DEVICE_CPU),
UniqueDatasetOp);
} // namespace
diff --git a/tensorflow/core/kernels/data/generator_dataset_op.cc b/tensorflow/core/kernels/data/generator_dataset_op.cc
index 71a3631..b4367d5 100644
--- a/tensorflow/core/kernels/data/generator_dataset_op.cc
+++ b/tensorflow/core/kernels/data/generator_dataset_op.cc
@@ -86,8 +86,6 @@
TF_RETURN_IF_ERROR(dataset()->init_func_->Instantiate(ctx));
TF_RETURN_IF_ERROR(dataset()->next_func_->Instantiate(ctx));
TF_RETURN_IF_ERROR(dataset()->finalize_func_->Instantiate(ctx));
- TF_RETURN_IF_ERROR(
- dataset()->init_func_->RunWithBorrowedArgs(ctx, {}, &state_));
return Status::OK();
}
@@ -96,6 +94,12 @@
bool* end_of_sequence) override {
mutex_lock l(mu_);
+ if (!initialized_) {
+ TF_RETURN_IF_ERROR(
+ dataset()->init_func_->RunWithBorrowedArgs(ctx, {}, &state_));
+ initialized_ = true;
+ }
+
if (finalized_) {
*end_of_sequence = true;
return Status::OK();
@@ -123,6 +127,7 @@
private:
mutex mu_;
+ bool initialized_ GUARDED_BY(mu_) = false;
bool finalized_ GUARDED_BY(mu_) = false;
std::vector<Tensor> state_ GUARDED_BY(mu_);
};
diff --git a/tensorflow/core/kernels/data/group_by_reducer_dataset_op.cc b/tensorflow/core/kernels/data/group_by_reducer_dataset_op.cc
index d6ee42a..e7244ee 100644
--- a/tensorflow/core/kernels/data/group_by_reducer_dataset_op.cc
+++ b/tensorflow/core/kernels/data/group_by_reducer_dataset_op.cc
@@ -30,8 +30,7 @@
class GroupByReducerDatasetOp : public UnaryDatasetOpKernel {
public:
explicit GroupByReducerDatasetOp(OpKernelConstruction* ctx)
- : UnaryDatasetOpKernel(ctx),
- graph_def_version_(ctx->graph_def_version()) {
+ : UnaryDatasetOpKernel(ctx) {
OP_REQUIRES_OK(ctx, ctx->GetAttr("key_func", &key_func_));
OP_REQUIRES_OK(ctx, ctx->GetAttr("init_func", &init_func_));
OP_REQUIRES_OK(ctx, ctx->GetAttr("reduce_func", &reduce_func_));
@@ -421,7 +420,6 @@
const std::vector<PartialTensorShape> output_shapes_;
};
- const int graph_def_version_;
DataTypeVector output_types_;
std::vector<PartialTensorShape> output_shapes_;
NameAttrList key_func_;
diff --git a/tensorflow/core/kernels/data/group_by_window_dataset_op.cc b/tensorflow/core/kernels/data/group_by_window_dataset_op.cc
index 8b417bb..14aefe5 100644
--- a/tensorflow/core/kernels/data/group_by_window_dataset_op.cc
+++ b/tensorflow/core/kernels/data/group_by_window_dataset_op.cc
@@ -31,8 +31,7 @@
class GroupByWindowDatasetOp : public UnaryDatasetOpKernel {
public:
explicit GroupByWindowDatasetOp(OpKernelConstruction* ctx)
- : UnaryDatasetOpKernel(ctx),
- graph_def_version_(ctx->graph_def_version()) {
+ : UnaryDatasetOpKernel(ctx) {
OP_REQUIRES_OK(ctx, ctx->GetAttr("key_func", &key_func_));
OP_REQUIRES_OK(ctx, ctx->GetAttr("reduce_func", &reduce_func_));
OP_REQUIRES_OK(ctx, ctx->GetAttr("window_size_func", &window_size_func_));
@@ -507,7 +506,6 @@
const std::vector<PartialTensorShape> output_shapes_;
};
- const int graph_def_version_;
DataTypeVector output_types_;
std::vector<PartialTensorShape> output_shapes_;
NameAttrList key_func_;
diff --git a/tensorflow/core/kernels/data/iterator_ops.cc b/tensorflow/core/kernels/data/iterator_ops.cc
index c0bc507..7a83366 100644
--- a/tensorflow/core/kernels/data/iterator_ops.cc
+++ b/tensorflow/core/kernels/data/iterator_ops.cc
@@ -659,6 +659,115 @@
BackgroundWorker background_worker_;
};
+class ReduceDatasetOp : public AsyncOpKernel {
+ public:
+ explicit ReduceDatasetOp(OpKernelConstruction* ctx)
+ : AsyncOpKernel(ctx),
+ background_worker_(
+ ctx->env(),
+ strings::StrCat("reduce_thread_", SanitizeThreadSuffix(name()))) {
+ OP_REQUIRES_OK(ctx, ctx->GetAttr("f", &reduce_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("use_inter_op_parallelism",
+ &use_inter_op_parallelism_));
+ }
+
+ void ComputeAsync(OpKernelContext* ctx, DoneCallback done) override {
+ // The call to `iterator->GetNext()` may block and depend on an
+ // inter-op thread pool thread, so we issue the call from the
+ // owned thread pool.
+ background_worker_.Schedule([this, ctx, done]() {
+ DatasetBase* dataset;
+ OP_REQUIRES_OK_ASYNC(
+ ctx, GetDatasetFromVariantTensor(ctx->input(0), &dataset), done);
+ OpInputList inputs;
+ OP_REQUIRES_OK_ASYNC(ctx, ctx->input_list("initial_state", &inputs),
+ done);
+ std::vector<Tensor> state(inputs.begin(), inputs.end());
+
+ std::unique_ptr<CapturedFunction> captured_func;
+ OP_REQUIRES_OK_ASYNC(
+ ctx,
+ CapturedFunction::Create(reduce_func_, ctx, "other_arguments",
+ use_inter_op_parallelism_, &captured_func),
+ done);
+
+ IteratorContext iter_ctx(ctx);
+ OP_REQUIRES_OK_ASYNC(ctx, captured_func->Instantiate(&iter_ctx), done);
+
+ std::unique_ptr<IteratorBase> iterator;
+ OP_REQUIRES_OK_ASYNC(
+ ctx, dataset->MakeIterator(&iter_ctx, "ReduceIterator", &iterator),
+ done);
+
+ // NOTE(jsimsa): We must destroy the iterator before calling `done()`, to
+ // avoid destruction races.
+ IteratorBase* raw_iterator = iterator.release();
+ auto cleanup = gtl::MakeCleanup([raw_iterator, done] {
+ delete raw_iterator;
+ done();
+ });
+
+ // Iterate through the input dataset.
+ Status status;
+ while (true) {
+ std::vector<Tensor> next_input_element;
+ bool end_of_input;
+ status = raw_iterator->GetNext(&iter_ctx, &next_input_element,
+ &end_of_input);
+ if (!status.ok() || end_of_input) {
+ break;
+ }
+
+ // Run the reduce function to update the current state.
+ std::vector<Tensor> args;
+ args.reserve(state.size() + next_input_element.size());
+ std::copy(state.begin(), state.end(), std::back_inserter(args));
+ std::copy(next_input_element.begin(), next_input_element.end(),
+ std::back_inserter(args));
+
+ std::vector<Tensor> reduce_func_output;
+ status =
+ captured_func->Run(&iter_ctx, std::move(args), &reduce_func_output);
+ if (!status.ok()) {
+ break;
+ }
+ std::swap(reduce_func_output, state);
+ }
+
+ if (!status.ok()) {
+ ctx->SetStatus(status);
+ return;
+ }
+ for (int i = 0; i < state.size(); ++i) {
+ OP_REQUIRES_ASYNC(
+ ctx, state[i].dtype() == output_types_[i],
+ errors::InvalidArgument(
+ "The result does not match the expected type for component ", i,
+ ". Expected: ", DataTypeString(output_types_[i]),
+ ". Actual: ", DataTypeString(state[i].dtype()), "."),
+ done);
+ OP_REQUIRES_ASYNC(
+ ctx, output_shapes_[i].IsCompatibleWith(state[i].shape()),
+ errors::InvalidArgument(
+ "The result does not match the expected shape for component ",
+ i, ". Expected: ", output_shapes_[i].DebugString(),
+ ". Actual: ", state[i].shape().DebugString(), "."),
+ done);
+ ctx->set_output(i, state[i]);
+ }
+ });
+ }
+
+ private:
+ NameAttrList reduce_func_;
+ DataTypeVector output_types_;
+ std::vector<PartialTensorShape> output_shapes_;
+ bool use_inter_op_parallelism_;
+ BackgroundWorker background_worker_;
+};
+
class OneShotIteratorOp : public AsyncOpKernel {
public:
explicit OneShotIteratorOp(OpKernelConstruction* ctx)
@@ -1146,6 +1255,8 @@
AnonymousIteratorHandleOp);
REGISTER_KERNEL_BUILDER(Name("DatasetToSingleElement").Device(DEVICE_CPU),
ToSingleElementOp);
+REGISTER_KERNEL_BUILDER(Name("ReduceDataset").Device(DEVICE_CPU),
+ ReduceDatasetOp);
REGISTER_KERNEL_BUILDER(Name("OneShotIterator").Device(DEVICE_CPU),
OneShotIteratorOp);
REGISTER_KERNEL_BUILDER(Name("IteratorGetNext").Device(DEVICE_CPU),
diff --git a/tensorflow/core/kernels/data/map_and_batch_dataset_op.cc b/tensorflow/core/kernels/data/map_and_batch_dataset_op.cc
index 2bbf4af..bf08970 100644
--- a/tensorflow/core/kernels/data/map_and_batch_dataset_op.cc
+++ b/tensorflow/core/kernels/data/map_and_batch_dataset_op.cc
@@ -37,6 +37,8 @@
// See documentation in ../ops/dataset_ops.cc for a high-level
// description of the following op.
+// TODO(b/116852688): Make coordination between the performance model and this
+// transformation more robust.
class MapAndBatchDatasetOp : public UnaryDatasetOpKernel {
public:
explicit MapAndBatchDatasetOp(OpKernelConstruction* ctx)
@@ -185,29 +187,31 @@
public:
explicit Iterator(const Params& params)
: DatasetIterator<Dataset>(params),
- num_parallel_calls_(params.dataset->num_parallel_calls_) {}
+ mu_(std::make_shared<mutex>()),
+ cond_var_(std::make_shared<condition_variable>()),
+ num_parallel_calls_(std::make_shared<model::SharedState>(
+ params.dataset->num_parallel_calls_, mu_, cond_var_)) {}
~Iterator() override {
- mutex_lock l(mu_);
+ mutex_lock l(*mu_);
// Cancel the runner thread.
cancelled_ = true;
- cond_var_.notify_all();
+ cond_var_->notify_all();
// Wait for all in-flight calls to complete.
while (num_calls_ > 0) {
- cond_var_.wait(l);
+ cond_var_->wait(l);
}
}
Status Initialize(IteratorContext* ctx) override {
- mutex_lock l(mu_);
+ mutex_lock l(*mu_);
AddConstantParameter(ctx, "batch_size", dataset()->batch_size_);
- if (num_parallel_calls_ == kAutoTune) {
- num_parallel_calls_ = 1;
- AddTunableParameter(ctx, "parallelism",
- &num_parallel_calls_ /* value */, 1 /* min */,
- port::NumSchedulableCPUs() /* max */, &cond_var_);
+ if (num_parallel_calls_->value == kAutoTune) {
+ num_parallel_calls_->value = 1;
+ AddTunableParameter(ctx, "parallelism", num_parallel_calls_, 1,
+ port::NumSchedulableCPUs());
} else {
- AddConstantParameter(ctx, "parallelism", num_parallel_calls_);
+ AddConstantParameter(ctx, "parallelism", num_parallel_calls_->value);
}
TF_RETURN_IF_ERROR(
dataset()->input_->MakeIterator(ctx, prefix(), &input_impl_));
@@ -219,27 +223,27 @@
bool* end_of_sequence) override {
std::shared_ptr<BatchResult> result;
{
- mutex_lock l(mu_);
+ mutex_lock l(*mu_);
EnsureRunnerThreadStarted(ctx);
while (batch_results_.empty() ||
batch_results_.front()->num_calls > 0) {
RecordStop(ctx);
- cond_var_.wait(l);
+ cond_var_->wait(l);
RecordStart(ctx);
}
std::swap(result, batch_results_.front());
batch_results_.pop_front();
- cond_var_.notify_all();
+ cond_var_->notify_all();
}
return ProcessResult(ctx, result, out_tensors, end_of_sequence);
}
protected:
Status SaveInternal(IteratorStateWriter* writer) override {
- mutex_lock l(mu_);
+ mutex_lock l(*mu_);
// Wait for all in-flight calls to complete.
while (num_calls_ > 0) {
- cond_var_.wait(l);
+ cond_var_->wait(l);
}
CHECK_EQ(num_calls_, 0);
TF_RETURN_IF_ERROR(SaveInput(writer, input_impl_));
@@ -255,7 +259,7 @@
Status RestoreInternal(IteratorContext* ctx,
IteratorStateReader* reader) override {
- mutex_lock l(mu_);
+ mutex_lock l(*mu_);
TF_RETURN_IF_ERROR(RestoreInput(ctx, reader, input_impl_));
TF_RETURN_IF_ERROR(
reader->ReadScalar(full_name("call_counter"), &call_counter_));
@@ -296,7 +300,7 @@
void Callback(const std::shared_ptr<IteratorContext>& ctx,
const std::shared_ptr<BatchResult>& result,
const std::shared_ptr<std::vector<Tensor>>& return_values,
- int64 offset, const Status& status) LOCKS_EXCLUDED(mu_) {
+ int64 offset, const Status& status) LOCKS_EXCLUDED(*mu_) {
result->UpdateStatus(status);
if (status.ok()) {
EnsureOutputAllocated(ctx, result, return_values);
@@ -332,16 +336,16 @@
}
void CallCompleted(const std::shared_ptr<BatchResult>& result)
- LOCKS_EXCLUDED(mu_) {
- mutex_lock l(mu_);
+ LOCKS_EXCLUDED(*mu_) {
+ mutex_lock l(*mu_);
num_calls_--;
result->num_calls--;
- cond_var_.notify_all();
+ cond_var_->notify_all();
}
void CallFunction(std::shared_ptr<IteratorContext> ctx,
const std::shared_ptr<BatchResult>& result,
- int64 offset) LOCKS_EXCLUDED(mu_) {
+ int64 offset) LOCKS_EXCLUDED(*mu_) {
// Get the next input element.
std::vector<Tensor> input_element;
bool end_of_input;
@@ -398,7 +402,7 @@
}
void EnsureRunnerThreadStarted(IteratorContext* ctx)
- EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+ EXCLUSIVE_LOCKS_REQUIRED(*mu_) {
if (!runner_thread_) {
std::shared_ptr<IteratorContext> ctx_copy(new IteratorContext(*ctx));
runner_thread_.reset(ctx->env()->StartThread(
@@ -474,14 +478,14 @@
}
void RunnerThread(const std::shared_ptr<IteratorContext>& ctx)
- LOCKS_EXCLUDED(mu_) {
+ LOCKS_EXCLUDED(*mu_) {
std::vector<std::pair<std::shared_ptr<BatchResult>, int64>> new_calls;
RecordStart(ctx.get());
auto stop_cleanup =
gtl::MakeCleanup([this, &ctx]() { RecordStop(ctx.get()); });
- new_calls.reserve(num_parallel_calls_);
- auto busy = [this]() EXCLUSIVE_LOCKS_REQUIRED(mu_) -> bool {
- int64 num_parallel_calls = num_parallel_calls_;
+ new_calls.reserve(num_parallel_calls_->value);
+ auto busy = [this]() EXCLUSIVE_LOCKS_REQUIRED(*mu_) -> bool {
+ int64 num_parallel_calls = num_parallel_calls_->value;
int64 max_batch_results =
(num_parallel_calls + dataset()->batch_size_ - 1) /
dataset()->batch_size_;
@@ -492,10 +496,10 @@
};
while (true) {
{
- mutex_lock l(mu_);
+ mutex_lock l(*mu_);
while (!cancelled_ && busy()) {
RecordStop(ctx.get());
- cond_var_.wait(l);
+ cond_var_->wait(l);
RecordStart(ctx.get());
}
@@ -522,7 +526,7 @@
}
Status ReadBatchResult(IteratorContext* ctx, IteratorStateReader* reader,
- size_t index) EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+ size_t index) EXCLUSIVE_LOCKS_REQUIRED(*mu_) {
batch_results_.emplace_back(new BatchResult(dataset()->batch_size_));
std::shared_ptr<BatchResult> result = batch_results_.back();
string prefix = strings::StrCat("batch_results_", index);
@@ -567,7 +571,7 @@
}
Status ReadStatus(IteratorStateReader* reader, const string& prefix,
- Status* status) EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+ Status* status) EXCLUSIVE_LOCKS_REQUIRED(*mu_) {
int64 code_int;
TF_RETURN_IF_ERROR(reader->ReadScalar(
full_name(strings::StrCat(prefix, "_code")), &code_int));
@@ -585,7 +589,7 @@
}
Status WriteBatchResult(IteratorStateWriter* writer, size_t index)
- EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+ EXCLUSIVE_LOCKS_REQUIRED(*mu_) {
std::shared_ptr<BatchResult> result = batch_results_[index];
string prefix = strings::StrCat("batch_results_", index);
mutex_lock l(result->mu);
@@ -626,7 +630,7 @@
}
Status WriteStatus(IteratorStateWriter* writer, const string& prefix,
- const Status& status) EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+ const Status& status) EXCLUSIVE_LOCKS_REQUIRED(*mu_) {
TF_RETURN_IF_ERROR(
writer->WriteScalar(full_name(strings::StrCat(prefix, "_code")),
static_cast<int64>(status.code())));
@@ -640,24 +644,24 @@
// Used for coordination between the main thread, the runner thread, and
// the callback threads.
- mutex mu_;
+ const std::shared_ptr<mutex> mu_;
// Used for coordination between the main thread, the runner thread, and
// the callback threads. In particular, the runner thread should only
- // schedule new calls when the number of in-flight calls is less than the
- // user specified level of parallelism and there are slots available in
- // the `batch_results_` buffer.
- condition_variable cond_var_;
+ // schedule new calls when the number of in-flight calls is less than
+ // `num_parallel_calls_->value` and there are slots available in the
+ // `batch_results_` buffer.
+ const std::shared_ptr<condition_variable> cond_var_;
// Identifies the maximum number of parallel calls.
- std::atomic<int64> num_parallel_calls_;
+ const std::shared_ptr<model::SharedState> num_parallel_calls_;
// Counts the number of outstanding calls for this batch.
- int64 num_calls_ GUARDED_BY(mu_) = 0;
+ int64 num_calls_ GUARDED_BY(*mu_) = 0;
// Counts the total number of calls.
- int64 call_counter_ GUARDED_BY(mu_) = 0;
+ int64 call_counter_ GUARDED_BY(*mu_) = 0;
std::unique_ptr<IteratorBase> input_impl_;
// Buffer for storing the (intermediate) batch results.
- std::deque<std::shared_ptr<BatchResult>> batch_results_ GUARDED_BY(mu_);
- std::unique_ptr<Thread> runner_thread_ GUARDED_BY(mu_);
- bool cancelled_ GUARDED_BY(mu_) = false;
+ std::deque<std::shared_ptr<BatchResult>> batch_results_ GUARDED_BY(*mu_);
+ std::unique_ptr<Thread> runner_thread_ GUARDED_BY(*mu_);
+ bool cancelled_ GUARDED_BY(*mu_) = false;
};
const DatasetBase* const input_;
diff --git a/tensorflow/core/kernels/data/multi_device_iterator_ops.cc b/tensorflow/core/kernels/data/multi_device_iterator_ops.cc
index 5f14396..d909b9e 100644
--- a/tensorflow/core/kernels/data/multi_device_iterator_ops.cc
+++ b/tensorflow/core/kernels/data/multi_device_iterator_ops.cc
@@ -134,19 +134,17 @@
void Reset() LOCKS_EXCLUDED(mu_) {
{
mutex_lock l(mu_);
- if (background_thread_finished_) {
- return;
- }
+ if (!background_thread_finished_) {
+ cancelled_ = true;
+ // Wake up the background thread.
+ for (int i = 0; i < size_; ++i) {
+ buffer_[i].cond_var.notify_all();
+ }
- cancelled_ = true;
- // Wake up the background thread.
- for (int i = 0; i < size_; ++i) {
- buffer_[i].cond_var.notify_all();
- }
-
- // Make sure background thread has finished first.
- while (!background_thread_finished_) {
- shutdown_cond_var_.wait(l);
+ // Make sure background thread has finished first.
+ while (!background_thread_finished_) {
+ shutdown_cond_var_.wait(l);
+ }
}
}
RunPendingCallbacks();
@@ -182,7 +180,7 @@
buffer_[shard_num].cond_var.notify_all();
}
} else {
- if (background_thread_finished_) {
+ if (end_of_iterator_) {
produced_output = true;
elem.end_of_sequence = true;
} else {
@@ -219,8 +217,12 @@
while (!buffer_[i].callbacks.empty()) {
if (buffer_[i].data.empty()) {
HostBufferElement elem;
- elem.status =
- errors::Cancelled("Cancelled and buffer not filled.");
+ if (end_of_iterator_) {
+ elem.end_of_sequence = true;
+ } else {
+ elem.status =
+ errors::Cancelled("Cancelled and buffer not filled.");
+ }
cancellation_elements.push_back(std::move(elem));
} else {
cancellation_elements.push_back(
@@ -293,6 +295,7 @@
{
mutex_lock l(mu_);
background_thread_finished_ = true;
+ end_of_iterator_ = true;
shutdown_cond_var_.notify_all();
}
RunPendingCallbacks();
@@ -312,6 +315,7 @@
std::unique_ptr<Thread> background_thread_ GUARDED_BY(mu_);
bool background_thread_finished_ GUARDED_BY(mu_) = false;
bool background_thread_started_ GUARDED_BY(mu_) = false;
+ bool end_of_iterator_ GUARDED_BY(mu_) = false;
bool cancelled_ GUARDED_BY(mu_) = false;
condition_variable shutdown_cond_var_ GUARDED_BY(mu_);
diff --git a/tensorflow/core/kernels/data/parallel_interleave_dataset_op.cc b/tensorflow/core/kernels/data/parallel_interleave_dataset_op.cc
index 2e6e046..6b6b3d6 100644
--- a/tensorflow/core/kernels/data/parallel_interleave_dataset_op.cc
+++ b/tensorflow/core/kernels/data/parallel_interleave_dataset_op.cc
@@ -1084,6 +1084,9 @@
// The above design choices were made with automated optimizations in mind,
// isolating the degree of parallelism as the single tunable knob of this
// implementation.
+//
+// TODO(b/116852688): Make coordination between the performance model and this
+// transformation more robust.
class ParallelInterleaveDatasetV2Op : public UnaryDatasetOpKernel {
public:
explicit ParallelInterleaveDatasetV2Op(OpKernelConstruction* ctx)
@@ -1214,7 +1217,10 @@
public:
explicit Iterator(const Params& params)
: DatasetIterator<Dataset>(params),
- num_parallel_calls_(params.dataset->num_parallel_calls_),
+ mu_(std::make_shared<mutex>()),
+ cond_var_(std::make_shared<condition_variable>()),
+ num_parallel_calls_(std::make_shared<model::SharedState>(
+ params.dataset->num_parallel_calls_, mu_, cond_var_)),
args_list_(params.dataset->cycle_length_),
current_elements_(params.dataset->cycle_length_),
element_in_use_(params.dataset->cycle_length_, false),
@@ -1224,25 +1230,24 @@
false /* low_latency_hint */)) {}
~Iterator() override {
- mutex_lock l(mu_);
+ mutex_lock l(*mu_);
// Cancel the runner thread.
cancelled_ = true;
- cond_var_.notify_all();
+ cond_var_->notify_all();
// Wait for all in-flight calls to complete.
while (num_calls_ > 0) {
- cond_var_.wait(l);
+ cond_var_->wait(l);
}
}
Status Initialize(IteratorContext* ctx) override {
- mutex_lock l(mu_);
- if (num_parallel_calls_ == kAutoTune) {
- num_parallel_calls_ = 1;
- AddTunableParameter(ctx, "parallelism",
- &num_parallel_calls_ /* value */, 1 /* min */,
- dataset()->cycle_length_ /* max */, &cond_var_);
+ mutex_lock l(*mu_);
+ if (num_parallel_calls_->value == kAutoTune) {
+ num_parallel_calls_->value = 1;
+ AddTunableParameter(ctx, "parallelism", num_parallel_calls_, 1,
+ dataset()->cycle_length_);
} else {
- AddConstantParameter(ctx, "parallelism", num_parallel_calls_);
+ AddConstantParameter(ctx, "parallelism", num_parallel_calls_->value);
}
AddConstantParameter(ctx, "cycle_length", dataset()->cycle_length_);
TF_RETURN_IF_ERROR(
@@ -1256,12 +1261,12 @@
std::shared_ptr<InvocationResult> result;
do {
{
- mutex_lock l(mu_);
+ mutex_lock l(*mu_);
EnsureRunnerThreadStarted(ctx);
while (invocation_results_.empty() &&
(!end_of_input_ || num_open_ > 0)) {
RecordStop(ctx);
- cond_var_.wait(l);
+ cond_var_->wait(l);
RecordStart(ctx);
}
if (!invocation_results_.empty()) {
@@ -1271,7 +1276,7 @@
*end_of_sequence = true;
return Status::OK();
}
- cond_var_.notify_all();
+ cond_var_->notify_all();
}
RecordStop(ctx);
result->notification.WaitForNotification();
@@ -1287,10 +1292,10 @@
protected:
Status SaveInternal(IteratorStateWriter* writer) override {
- mutex_lock l(mu_);
+ mutex_lock l(*mu_);
// Wait for all in-flight calls to complete.
while (num_calls_ > 0) {
- cond_var_.wait(l);
+ cond_var_->wait(l);
}
CHECK_EQ(num_calls_, 0);
TF_RETURN_IF_ERROR(SaveInput(writer, input_impl_));
@@ -1328,7 +1333,7 @@
Status RestoreInternal(IteratorContext* ctx,
IteratorStateReader* reader) override {
- mutex_lock l(mu_);
+ mutex_lock l(*mu_);
TF_RETURN_IF_ERROR(RestoreInput(ctx, reader, input_impl_));
int64 invocation_results_size;
TF_RETURN_IF_ERROR(reader->ReadScalar(
@@ -1381,7 +1386,7 @@
};
void EnsureRunnerThreadStarted(IteratorContext* ctx)
- EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+ EXCLUSIVE_LOCKS_REQUIRED(*mu_) {
if (!runner_thread_) {
std::shared_ptr<IteratorContext> new_ctx(new IteratorContext(*ctx));
runner_thread_.reset(ctx->env()->StartThread(
@@ -1398,7 +1403,7 @@
void FetchOutputs(
const std::shared_ptr<IteratorContext>& ctx, int64 cycle_index,
const std::vector<std::shared_ptr<InvocationResult>>& results)
- LOCKS_EXCLUDED(mu_) {
+ LOCKS_EXCLUDED(*mu_) {
RecordStart(ctx.get());
auto cleanup = gtl::MakeCleanup([this, ctx] { RecordStop(ctx.get()); });
bool end_of_input = false;
@@ -1421,14 +1426,14 @@
if (end_of_input) {
current_elements_[cycle_index].reset();
}
- mutex_lock l(mu_);
+ mutex_lock l(*mu_);
element_in_use_[cycle_index] = false;
num_calls_--;
if (end_of_input) {
args_list_[cycle_index].clear();
num_open_--;
}
- cond_var_.notify_all();
+ cond_var_->notify_all();
}
// Method responsible for 1) creating iterators out of input elements, 2)
@@ -1439,20 +1444,20 @@
void RunnerThread(const std::shared_ptr<IteratorContext>& ctx) {
RecordStart(ctx.get());
auto cleanup = gtl::MakeCleanup([this, ctx] { RecordStop(ctx.get()); });
- auto busy = [this]() EXCLUSIVE_LOCKS_REQUIRED(mu_) -> bool {
+ auto busy = [this]() EXCLUSIVE_LOCKS_REQUIRED(*mu_) -> bool {
return element_in_use_[cycle_index_] ||
- num_calls_ >= num_parallel_calls_ ||
+ num_calls_ >= num_parallel_calls_->value ||
invocation_results_.size() >=
dataset()->cycle_length_ * dataset()->block_length_;
};
while (true) {
- mutex_lock l(mu_);
+ mutex_lock l(*mu_);
// Wait until this thread is cancelled, the end of input has been
// reached, or the cycle element at the `cycle_index_` position is
// not in use and there is space in the `invocation_results_` queue.
while (!cancelled_ && (!end_of_input_ || num_open_ > 0) && busy()) {
RecordStop(ctx.get());
- cond_var_.wait(l);
+ cond_var_->wait(l);
RecordStart(ctx.get());
}
@@ -1506,13 +1511,13 @@
}
cycle_index_ = (cycle_index_ + 1) % dataset()->cycle_length_;
}
- cond_var_.notify_all();
+ cond_var_->notify_all();
}
}
Status WriteStatusLocked(IteratorStateWriter* writer, size_t index,
const Status& status)
- EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+ EXCLUSIVE_LOCKS_REQUIRED(*mu_) {
TF_RETURN_IF_ERROR(writer->WriteScalar(
CodeKey(index), static_cast<int64>(status.code())));
if (!status.ok()) {
@@ -1523,7 +1528,7 @@
}
Status ReadStatusLocked(IteratorStateReader* reader, size_t index,
- Status* status) EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+ Status* status) EXCLUSIVE_LOCKS_REQUIRED(*mu_) {
int64 code_int;
TF_RETURN_IF_ERROR(reader->ReadScalar(CodeKey(index), &code_int));
error::Code code = static_cast<error::Code>(code_int);
@@ -1550,7 +1555,7 @@
}
Status WriteCurrentElements(IteratorStateWriter* writer)
- EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+ EXCLUSIVE_LOCKS_REQUIRED(*mu_) {
for (int idx = 0; idx < current_elements_.size(); idx++) {
if (current_elements_[idx]) {
TF_RETURN_IF_ERROR(SaveInput(writer, current_elements_[idx]));
@@ -1569,7 +1574,7 @@
Status ReadCurrentElements(IteratorContext* ctx,
IteratorStateReader* reader)
- EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+ EXCLUSIVE_LOCKS_REQUIRED(*mu_) {
for (int idx = 0; idx < current_elements_.size(); idx++) {
if (reader->Contains(
full_name(strings::StrCat("args_size[", idx, "]")))) {
@@ -1597,7 +1602,7 @@
// Used for coordination between the main thread, the runner thread, and
// the worker threads.
- mutex mu_;
+ const std::shared_ptr<mutex> mu_;
// Used for coordination between the main thread, the runner thread, and
// the worker threads. In particular, the runner thread should only
@@ -1605,45 +1610,45 @@
// user specified level of parallelism, there are slots available in the
// `invocation_results_` buffer, the current cycle element is not in use,
// and there are elements left to be fetched.
- condition_variable cond_var_;
+ const std::shared_ptr<condition_variable> cond_var_;
// Identifies the maximum number of parallel calls.
- std::atomic<int64> num_parallel_calls_;
+ const std::shared_ptr<model::SharedState> num_parallel_calls_;
// Iterator for input elements.
- std::unique_ptr<IteratorBase> input_impl_ GUARDED_BY(mu_);
+ std::unique_ptr<IteratorBase> input_impl_ GUARDED_BY(*mu_);
// Identifies current cycle element.
int64 cycle_index_ = 0;
// Arguments for creating an iterator for cycle elements.
- std::vector<std::vector<Tensor>> args_list_ GUARDED_BY(mu_);
+ std::vector<std::vector<Tensor>> args_list_ GUARDED_BY(*mu_);
// Iterators for the current cycle elements. Concurrent access is
// protected by `element_in_use_`.
std::vector<std::unique_ptr<IteratorBase>> current_elements_;
// Identifies cycle elements that are in use by worker threads.
- std::vector<bool> element_in_use_ GUARDED_BY(mu_);
+ std::vector<bool> element_in_use_ GUARDED_BY(*mu_);
// Buffer for storing the invocation results.
std::deque<std::shared_ptr<InvocationResult>> invocation_results_
- GUARDED_BY(mu_);
+ GUARDED_BY(*mu_);
// Identifies whether end of input has been reached.
- bool end_of_input_ GUARDED_BY(mu_) = false;
+ bool end_of_input_ GUARDED_BY(*mu_) = false;
// Identifies the number of open iterators.
- int64 num_open_ GUARDED_BY(mu_) = 0;
+ int64 num_open_ GUARDED_BY(*mu_) = 0;
// Identifies the number of outstanding calls.
- int64 num_calls_ GUARDED_BY(mu_) = 0;
+ int64 num_calls_ GUARDED_BY(*mu_) = 0;
std::unique_ptr<thread::ThreadPool> thread_pool_;
- std::unique_ptr<Thread> runner_thread_ GUARDED_BY(mu_);
+ std::unique_ptr<Thread> runner_thread_ GUARDED_BY(*mu_);
// Identifies whether background activity should be cancelled.
- bool cancelled_ GUARDED_BY(mu_) = false;
+ bool cancelled_ GUARDED_BY(*mu_) = false;
};
const DatasetBase* const input_;
diff --git a/tensorflow/core/kernels/data/parallel_map_iterator.cc b/tensorflow/core/kernels/data/parallel_map_iterator.cc
index ee20249b..13bd4b6 100644
--- a/tensorflow/core/kernels/data/parallel_map_iterator.cc
+++ b/tensorflow/core/kernels/data/parallel_map_iterator.cc
@@ -27,6 +27,8 @@
namespace data {
namespace {
+// TODO(b/116852688): Make coordination between the performance model and this
+// transformation more robust.
class ParallelMapIterator : public DatasetBaseIterator {
public:
explicit ParallelMapIterator(
@@ -38,30 +40,32 @@
input_dataset_(input_dataset),
init_func_(std::move(init_func)),
map_func_(std::move(map_func)),
- num_parallel_calls_(num_parallel_calls) {}
+ 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_)) {}
~ParallelMapIterator() override {
- mutex_lock l(mu_);
+ mutex_lock l(*mu_);
// Cancel the runner thread.
cancelled_ = true;
- cond_var_.notify_all();
+ cond_var_->notify_all();
// Wait for all in-flight calls to complete.
while (num_calls_ > 0) {
- cond_var_.wait(l);
+ cond_var_->wait(l);
}
}
Status Initialize(IteratorContext* ctx) override {
- mutex_lock l(mu_);
- if (num_parallel_calls_ == kAutoTune) {
- num_parallel_calls_ = 1;
+ mutex_lock l(*mu_);
+ if (num_parallel_calls_->value == kAutoTune) {
+ num_parallel_calls_->value = 1;
// TODO(jsimsa): Surface the number of threads used by `ctx->runner()` and
// use it here for the maximum.
- AddTunableParameter(ctx, "parallelism", &num_parallel_calls_ /* value */,
- 1 /* min */, port::NumSchedulableCPUs() /* max */,
- &cond_var_);
+ AddTunableParameter(ctx, "parallelism", num_parallel_calls_, 1,
+ port::NumSchedulableCPUs());
} else {
- AddConstantParameter(ctx, "parallelism", num_parallel_calls_);
+ AddConstantParameter(ctx, "parallelism", num_parallel_calls_->value);
}
TF_RETURN_IF_ERROR(
input_dataset_->MakeIterator(ctx, prefix(), &input_impl_));
@@ -75,16 +79,16 @@
bool* end_of_sequence) override {
std::shared_ptr<InvocationResult> result;
{
- mutex_lock l(mu_);
+ mutex_lock l(*mu_);
EnsureRunnerThreadStarted(ctx);
while (invocation_results_.empty()) {
RecordStop(ctx);
- cond_var_.wait(l);
+ cond_var_->wait(l);
RecordStart(ctx);
}
std::swap(result, invocation_results_.front());
invocation_results_.pop_front();
- cond_var_.notify_all();
+ cond_var_->notify_all();
}
RecordStop(ctx);
result->notification.WaitForNotification();
@@ -94,28 +98,27 @@
protected:
Status SaveInternal(IteratorStateWriter* writer) override {
- mutex_lock l(mu_);
+ mutex_lock l(*mu_);
// Wait for all in-flight calls to complete.
while (num_calls_ > 0) {
- cond_var_.wait(l);
+ cond_var_->wait(l);
}
CHECK_EQ(num_calls_, 0);
TF_RETURN_IF_ERROR(SaveInput(writer, input_impl_));
TF_RETURN_IF_ERROR(writer->WriteScalar(full_name("invocation_results.size"),
invocation_results_.size()));
for (size_t i = 0; i < invocation_results_.size(); i++) {
- std::shared_ptr<InvocationResult> result = invocation_results_[i];
- TF_RETURN_IF_ERROR(WriteStatusLocked(writer, i, result->status));
+ const auto& result = *(invocation_results_[i]);
+ TF_RETURN_IF_ERROR(WriteStatusLocked(writer, i, result.status));
TF_RETURN_IF_ERROR(writer->WriteScalar(
full_name(strings::StrCat("invocation_results[", i, "].size")),
- result->return_values.size()));
- for (size_t j = 0; j < result->return_values.size(); j++) {
- TF_RETURN_IF_ERROR(
- writer->WriteTensor(full_name(strings::StrCat(
- "invocation_results[", i, "][", j, "]")),
- result->return_values[j]));
+ result.return_values.size()));
+ for (size_t j = 0; j < result.return_values.size(); j++) {
+ TF_RETURN_IF_ERROR(writer->WriteTensor(
+ full_name(strings::StrCat("invocation_results[", i, "][", j, "]")),
+ result.return_values[j]));
}
- if (result->end_of_input) {
+ if (result.end_of_input) {
TF_RETURN_IF_ERROR(writer->WriteScalar(
full_name(
strings::StrCat("invocation_results[", i, "].end_of_input")),
@@ -127,15 +130,15 @@
Status RestoreInternal(IteratorContext* ctx,
IteratorStateReader* reader) override {
- mutex_lock l(mu_);
+ mutex_lock l(*mu_);
TF_RETURN_IF_ERROR(RestoreInput(ctx, reader, input_impl_));
int64 invocation_results_size;
TF_RETURN_IF_ERROR(reader->ReadScalar(
full_name("invocation_results.size"), &invocation_results_size));
for (size_t i = 0; i < invocation_results_size; i++) {
- std::shared_ptr<InvocationResult> result(new InvocationResult());
- invocation_results_.push_back(result);
- TF_RETURN_IF_ERROR(ReadStatusLocked(reader, i, &result->status));
+ invocation_results_.push_back(std::make_shared<InvocationResult>());
+ auto& result = *invocation_results_.back();
+ TF_RETURN_IF_ERROR(ReadStatusLocked(reader, i, &result.status));
size_t num_return_values;
{
int64 size;
@@ -151,17 +154,16 @@
": ", size, " is not a valid value of type size_t."));
}
}
- result->return_values.reserve(num_return_values);
+ result.return_values.reserve(num_return_values);
for (size_t j = 0; j < num_return_values; j++) {
- result->return_values.emplace_back();
- TF_RETURN_IF_ERROR(
- reader->ReadTensor(full_name(strings::StrCat(
- "invocation_results[", i, "][", j, "]")),
- &result->return_values.back()));
+ result.return_values.emplace_back();
+ TF_RETURN_IF_ERROR(reader->ReadTensor(
+ full_name(strings::StrCat("invocation_results[", i, "][", j, "]")),
+ &result.return_values.back()));
}
- result->end_of_input = reader->Contains(full_name(
+ result.end_of_input = reader->Contains(full_name(
strings::StrCat("invocation_results[", i, "].end_of_input")));
- result->notification.Notify();
+ result.notification.Notify();
}
return Status::OK();
}
@@ -175,7 +177,7 @@
};
void EnsureRunnerThreadStarted(IteratorContext* ctx)
- EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+ EXCLUSIVE_LOCKS_REQUIRED(*mu_) {
if (!runner_thread_) {
std::shared_ptr<IteratorContext> ctx_copy(new IteratorContext(*ctx));
runner_thread_.reset(ctx->env()->StartThread(
@@ -185,18 +187,18 @@
}
void CallCompleted(const std::shared_ptr<InvocationResult>& result)
- LOCKS_EXCLUDED(mu_) {
+ LOCKS_EXCLUDED(*mu_) {
{
- mutex_lock l(mu_);
+ mutex_lock l(*mu_);
num_calls_--;
- cond_var_.notify_all();
+ cond_var_->notify_all();
}
result->notification.Notify();
}
void CallFunction(const std::shared_ptr<IteratorContext>& ctx,
const std::shared_ptr<InvocationResult>& result)
- LOCKS_EXCLUDED(mu_) {
+ LOCKS_EXCLUDED(*mu_) {
// Get the next input element.
std::vector<Tensor> input_element;
result->status =
@@ -239,29 +241,29 @@
RecordStart(ctx.get());
auto cleanup = gtl::MakeCleanup([this, ctx] { RecordStop(ctx.get()); });
std::vector<std::shared_ptr<InvocationResult>> new_calls;
- new_calls.reserve(num_parallel_calls_);
- auto busy = [this]() EXCLUSIVE_LOCKS_REQUIRED(mu_) -> bool {
- int64 num_parallel_calls = num_parallel_calls_;
+ new_calls.reserve(num_parallel_calls_->value);
+ auto busy = [this]() EXCLUSIVE_LOCKS_REQUIRED(*mu_) -> bool {
+ int64 num_parallel_calls = num_parallel_calls_->value;
return num_calls_ >= num_parallel_calls ||
invocation_results_.size() >= num_parallel_calls;
};
while (true) {
{
- mutex_lock l(mu_);
+ mutex_lock l(*mu_);
while (!cancelled_ && busy()) {
RecordStop(ctx.get());
- cond_var_.wait(l);
+ cond_var_->wait(l);
RecordStart(ctx.get());
}
if (cancelled_) {
return;
}
while (!busy()) {
- invocation_results_.emplace_back(new InvocationResult());
+ invocation_results_.push_back(std::make_shared<InvocationResult>());
new_calls.push_back(invocation_results_.back());
num_calls_++;
}
- cond_var_.notify_all();
+ cond_var_->notify_all();
}
for (const auto& call : new_calls) {
CallFunction(ctx, call);
@@ -271,7 +273,8 @@
}
Status WriteStatusLocked(IteratorStateWriter* writer, size_t index,
- const Status& status) EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+ const Status& status)
+ EXCLUSIVE_LOCKS_REQUIRED(*mu_) {
TF_RETURN_IF_ERROR(
writer->WriteScalar(CodeKey(index), static_cast<int64>(status.code())));
if (!status.ok()) {
@@ -282,7 +285,7 @@
}
Status ReadStatusLocked(IteratorStateReader* reader, size_t index,
- Status* status) EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+ Status* status) EXCLUSIVE_LOCKS_REQUIRED(*mu_) {
int64 code_int;
TF_RETURN_IF_ERROR(reader->ReadScalar(CodeKey(index), &code_int));
error::Code code = static_cast<error::Code>(code_int);
@@ -312,23 +315,23 @@
const std::function<Status(IteratorContext*)> init_func_;
const ParallelMapIteratorFunction map_func_;
// Used for coordination between the main thread and the runner thread.
- mutex mu_;
+ const std::shared_ptr<mutex> mu_;
// Used for coordination between the main thread and the runner thread. In
// particular, the runner thread should only schedule new calls when the
// number of in-flight calls is less than the user specified level of
// parallelism and there are slots available in the `invocation_results_`
// buffer.
- condition_variable cond_var_;
+ const std::shared_ptr<condition_variable> cond_var_;
// Identifies the maximum number of parallel calls.
- std::atomic<int64> num_parallel_calls_;
+ const std::shared_ptr<model::SharedState> num_parallel_calls_;
// Counts the number of outstanding calls.
- int64 num_calls_ GUARDED_BY(mu_) = 0;
+ int64 num_calls_ GUARDED_BY(*mu_) = 0;
std::unique_ptr<IteratorBase> input_impl_;
// Buffer for storing the invocation results.
std::deque<std::shared_ptr<InvocationResult>> invocation_results_
- GUARDED_BY(mu_);
- std::unique_ptr<Thread> runner_thread_ GUARDED_BY(mu_);
- bool cancelled_ GUARDED_BY(mu_) = false;
+ GUARDED_BY(*mu_);
+ std::unique_ptr<Thread> runner_thread_ GUARDED_BY(*mu_);
+ bool cancelled_ GUARDED_BY(*mu_) = false;
};
} // namespace
diff --git a/tensorflow/core/kernels/data/scan_dataset_op.cc b/tensorflow/core/kernels/data/scan_dataset_op.cc
index dbe31f3..2a911aa 100644
--- a/tensorflow/core/kernels/data/scan_dataset_op.cc
+++ b/tensorflow/core/kernels/data/scan_dataset_op.cc
@@ -32,8 +32,7 @@
class ScanDatasetOp : public UnaryDatasetOpKernel {
public:
explicit ScanDatasetOp(OpKernelConstruction* ctx)
- : UnaryDatasetOpKernel(ctx),
- graph_def_version_(ctx->graph_def_version()) {
+ : UnaryDatasetOpKernel(ctx) {
OP_REQUIRES_OK(ctx, ctx->GetAttr("f", &func_));
OP_REQUIRES_OK(ctx, ctx->GetAttr("Tstate", &state_types_));
OP_REQUIRES_OK(ctx, ctx->GetAttr("output_types", &output_types_));
@@ -258,7 +257,6 @@
const std::vector<PartialTensorShape> output_shapes_;
};
- const int graph_def_version_;
DataTypeVector state_types_;
DataTypeVector output_types_;
std::vector<PartialTensorShape> output_shapes_;
diff --git a/tensorflow/core/kernels/data/stats_aggregator_dataset_op.cc b/tensorflow/core/kernels/data/stats_aggregator_dataset_op.cc
index f5314f7..c8abfb9 100644
--- a/tensorflow/core/kernels/data/stats_aggregator_dataset_op.cc
+++ b/tensorflow/core/kernels/data/stats_aggregator_dataset_op.cc
@@ -34,16 +34,18 @@
&stats_aggregator_resource));
core::ScopedUnref unref_stats_aggregator(stats_aggregator_resource);
- *output = new Dataset(ctx, input, stats_aggregator_resource);
+ *output = new Dataset(ctx, input, ctx->input(1), stats_aggregator_resource);
}
private:
class Dataset : public DatasetBase {
public:
explicit Dataset(OpKernelContext* ctx, const DatasetBase* input,
+ const Tensor& resource_handle,
StatsAggregatorResource* stats_aggregator_resource)
: DatasetBase(DatasetContext(ctx)),
input_(input),
+ resource_handle_(resource_handle),
stats_aggregator_resource_(stats_aggregator_resource) {
input_->Ref();
stats_aggregator_resource_->Ref();
@@ -75,8 +77,13 @@
Status AsGraphDefInternal(SerializationContext* ctx,
DatasetGraphDefBuilder* b,
Node** output) const override {
- return errors::Unimplemented("%s does not support serialization",
- DebugString());
+ Node* input_graph_node = nullptr;
+ TF_RETURN_IF_ERROR(b->AddInputDataset(ctx, input_, &input_graph_node));
+ Node* resource_handle_node = nullptr;
+ TF_RETURN_IF_ERROR(b->AddTensor(resource_handle_, &resource_handle_node));
+ TF_RETURN_IF_ERROR(b->AddDataset(
+ this, {input_graph_node, resource_handle_node}, output));
+ return Status::OK();
}
private:
@@ -111,16 +118,14 @@
protected:
Status SaveInternal(IteratorStateWriter* writer) override {
- mutex_lock l(mu_);
- TF_RETURN_IF_ERROR(SaveInput(writer, input_impl_));
- return Status::OK();
+ return errors::Unimplemented(dataset()->DebugString(),
+ " does not support checkpointing");
}
Status RestoreInternal(IteratorContext* ctx,
IteratorStateReader* reader) override {
- mutex_lock l(mu_);
- TF_RETURN_IF_ERROR(RestoreInput(ctx, reader, input_impl_));
- return Status::OK();
+ return errors::Unimplemented(dataset()->DebugString(),
+ " does not support checkpointing");
}
private:
@@ -129,6 +134,7 @@
};
const DatasetBase* const input_;
+ const Tensor resource_handle_;
StatsAggregatorResource* stats_aggregator_resource_;
};
};
diff --git a/tensorflow/core/kernels/fuzzing/parse_tensor_op_fuzz.cc b/tensorflow/core/kernels/fuzzing/parse_tensor_op_fuzz.cc
index c90ad2c..ada1235 100644
--- a/tensorflow/core/kernels/fuzzing/parse_tensor_op_fuzz.cc
+++ b/tensorflow/core/kernels/fuzzing/parse_tensor_op_fuzz.cc
@@ -31,9 +31,37 @@
}
void FuzzImpl(const uint8_t* data, size_t size) final {
+ // We need to be sure that we don't request too many elements (i.e., we
+ // don't make ASAN OOM). In theory, a tensor shape can have arbitrary large
+ // number of elements, up to the limit of the memory available to the OS.
+ // However, due to the tracing done in ASAN, after 2^32 bytes of requested
+ // memory we would get a crash in the fuzzer (see b/34190148). Hence, let's
+ // try parsing the proto here, check that the size (if valid) is below a
+ // maximum threshold (using 2^20 for convenience), and then run the
+ // remainder of the fuzzer testing. Of course, this duplicates some work
+ // but it's better than repeating the investigation whenever Autofuzz
+ // detects another similar OOM.
+ string as_string = string(reinterpret_cast<const char*>(data), size);
+ TensorProto proto;
+ if (!ParseProtoUnlimited(&proto, as_string)) {
+ LOG(WARNING) << "Unable to parse proto of tensor\n";
+ return;
+ }
+ if (!TensorShape::IsValid(proto.tensor_shape())) {
+ LOG(WARNING) << "Invalid tensor shape\n";
+ return;
+ }
+ TensorShape shape(proto.tensor_shape());
+ const int64 num_elements = shape.num_elements();
+ const int64 max_num_elements = 1 << 20;
+ if (num_elements > max_num_elements) {
+ LOG(WARNING) << "Requiring a tensor with too many elements\n";
+ return;
+ }
+
+ // Now we can do the actual fuzz implementation
Tensor input_tensor(tensorflow::DT_STRING, TensorShape({}));
- input_tensor.scalar<string>()() =
- string(reinterpret_cast<const char*>(data), size);
+ input_tensor.scalar<string>()() = as_string;
// TODO(b/32704451): Don't just ignore the ::tensorflow::Status object!
RunOneInput(input_tensor).IgnoreError();
}
diff --git a/tensorflow/core/kernels/mkl_conv_grad_filter_ops.cc b/tensorflow/core/kernels/mkl_conv_grad_filter_ops.cc
index 52157ed..f406ad2 100644
--- a/tensorflow/core/kernels/mkl_conv_grad_filter_ops.cc
+++ b/tensorflow/core/kernels/mkl_conv_grad_filter_ops.cc
@@ -853,7 +853,7 @@
// MKL DNN allocates large buffers when a conv gradient filter primtive is
// created. So we don't cache conv backward primitives when the env
- // variable TF_MKL_OPTIMIZE_PRIMITVE_MEMUSE is set to true.
+ // variable TF_MKL_OPTIMIZE_PRIMITIVE_MEMUSE is set to true.
bool do_not_cache = MklPrimitiveFactory<T>::IsPrimitiveMemOptEnabled();
conv_bwd_filter = MklConvBwdFilterPrimitiveFactory<T>::Get(
convBwdFilterDims, do_not_cache);
diff --git a/tensorflow/core/kernels/mkl_conv_grad_input_ops.cc b/tensorflow/core/kernels/mkl_conv_grad_input_ops.cc
index c38c9cc..a501ce2 100644
--- a/tensorflow/core/kernels/mkl_conv_grad_input_ops.cc
+++ b/tensorflow/core/kernels/mkl_conv_grad_input_ops.cc
@@ -713,7 +713,7 @@
TFPaddingToMklDnnPadding(this->padding_));
// We don't cache those primitves if the env variable
- // TF_MKL_OPTIMIZE_PRIMITVE_MEMUSE is true and if primitve descriptor
+ // TF_MKL_OPTIMIZE_PRIMITIVE_MEMUSE is true and if primitve descriptor
// includes potentialy large buffers. MKL DNN allocates buffers
// in the following cases
// 1. Legacy CPU without AVX512/AVX2, or
diff --git a/tensorflow/core/kernels/mkl_conv_ops.cc b/tensorflow/core/kernels/mkl_conv_ops.cc
index 184e0cb..b332eda 100644
--- a/tensorflow/core/kernels/mkl_conv_ops.cc
+++ b/tensorflow/core/kernels/mkl_conv_ops.cc
@@ -901,7 +901,7 @@
// In some cases, primitve descriptor includes potentialy large buffers,
// we don't cache those primitves if the env variable
- // TF_MKL_OPTIMIZE_PRIMITVE_MEMUSE is true. MKL DNN allocates buffers
+ // TF_MKL_OPTIMIZE_PRIMITIVE_MEMUSE is true. MKL DNN allocates buffers
// in the following cases
// 1. Legacy CPU without AVX512/AVX2, or
// 2. 1x1 convolution with stride != 1
diff --git a/tensorflow/core/kernels/mkl_slice_op.cc b/tensorflow/core/kernels/mkl_slice_op.cc
new file mode 100644
index 0000000..d63e14a
--- /dev/null
+++ b/tensorflow/core/kernels/mkl_slice_op.cc
@@ -0,0 +1,358 @@
+/* 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.
+==============================================================================*/
+
+// See docs in ../ops/array_ops.cc.
+
+#ifdef INTEL_MKL
+#ifndef INTEL_MKL_ML_ONLY
+
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/register_types.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/kernels/ops_util.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/gtl/array_slice.h"
+#include "tensorflow/core/platform/prefetch.h"
+#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
+
+#include "mkldnn.hpp"
+#include "tensorflow/core/util/mkl_util.h"
+
+using mkldnn::stream;
+using mkldnn::view;
+
+namespace tensorflow {
+
+namespace {
+
+gtl::InlinedVector<int64, 4> IntTensorToInt64Vec(const Tensor& tensor) {
+ gtl::InlinedVector<int64, 4> out;
+ if (tensor.dtype() == DT_INT32) {
+ for (int64 i = 0; i < tensor.NumElements(); ++i) {
+ out.push_back(tensor.flat<int32>()(i));
+ }
+ } else if (tensor.dtype() == DT_INT64) {
+ for (int64 i = 0; i < tensor.NumElements(); ++i) {
+ out.push_back(tensor.flat<int64>()(i));
+ }
+ } else {
+ // tensor must be either int32 or int64
+ DCHECK(false);
+ }
+ return out;
+}
+
+} // namespace
+
+typedef Eigen::ThreadPoolDevice CPUDevice;
+
+// A version of SharedValidation (slice_op.h) written for input that is in
+// either Mkl layout or Tensorflow layout.
+// A shared code to validate input shapes and check for identity, which is not dependent on the type of T.
+// We do this to reduce code size by not duplicating all this for all T (float, double, int32, etc.)
+static void ValidateMklInputs(OpKernelContext* context, bool* is_identity,
+ gtl::InlinedVector<int64, 4>* begin,
+ gtl::InlinedVector<int64, 4>* size) {
+ const int kInputTensorIndex = 0;
+ const int kInputBeginIndex = 1;
+ const int kInputSizeIndex = 2;
+ const Tensor& input = MklGetInput(context, kInputTensorIndex);
+ const Tensor& begin_tensor = MklGetInput(context, kInputBeginIndex);
+ const Tensor& size_tensor = MklGetInput(context, kInputSizeIndex);
+
+ MklDnnShape input_mkl_shape, begin_mkl_shape, size_mkl_shape;
+ GetMklShape(context, kInputTensorIndex, &input_mkl_shape);
+ GetMklShape(context, kInputBeginIndex, &begin_mkl_shape);
+ GetMklShape(context, kInputSizeIndex, &size_mkl_shape);
+
+ // Begin and size tensors cannot be in MklDnn layout.
+ DCHECK_EQ(begin_mkl_shape.IsMklTensor(), false);
+ DCHECK_EQ(size_mkl_shape.IsMklTensor(), false);
+
+ TensorShape input_tf_shape = input_mkl_shape.IsMklTensor()
+ ? input_mkl_shape.GetTfShape()
+ : input.shape();
+ const int input_dims = input_tf_shape.dims();
+
+ OP_REQUIRES(
+ context, context->op_kernel().IsLegacyVector(begin_tensor.shape()) &&
+ context->op_kernel().IsLegacyVector(size_tensor.shape()) &&
+ begin_tensor.NumElements() == input_dims &&
+ size_tensor.NumElements() == input_dims,
+ errors::InvalidArgument(
+ "Expected begin and size arguments to be 1-D tensors of size ",
+ input_dims, ", but got shapes ", begin_tensor.shape().DebugString(),
+ " and ", size_tensor.shape().DebugString(), " instead."));
+
+ *begin = IntTensorToInt64Vec(begin_tensor);
+ *size = IntTensorToInt64Vec(size_tensor);
+ for (int i = 0; i < input_dims; ++i) {
+ if ((*size)[i] == -1) {
+ // A size[i] of -1 means "all elements from begin[i] to dim_size(i)".
+ (*size)[i] = input_tf_shape.dim_size(i) - (*begin)[i];
+ }
+ }
+
+ *is_identity = true;
+ for (int i = 0; i < input_dims; ++i) {
+ int64 b = (*begin)[i];
+ int64 s = (*size)[i];
+ if (input_tf_shape.dim_size(i) == 0) {
+ OP_REQUIRES(
+ context, b == 0 && s == 0,
+ errors::InvalidArgument("Expected begin[", i, "] == 0 (got ", b,
+ ") and size[", i, "] == 0 ", "(got ", s,
+ ") when ", "input.dim_size(", i, ") == 0"));
+ } else {
+ OP_REQUIRES(context, 0 <= b && b <= input_tf_shape.dim_size(i),
+ errors::InvalidArgument("Expected begin[", i, "] in [0, ",
+ input_tf_shape.dim_size(i),
+ "], but got ", b));
+ OP_REQUIRES(context, 0 <= s && b + s <= input_tf_shape.dim_size(i),
+ errors::InvalidArgument("Expected size[", i, "] in [0, ",
+ input_tf_shape.dim_size(i) - b,
+ "], but ", "got ", s));
+ }
+ const bool take_all = (b == 0) && (s == input_tf_shape.dim_size(i));
+ (*is_identity) &= take_all;
+ }
+}
+
+// A version of SharedSliceCommonCases function written for input tensor
+// that may be in MklDnn layout or in Tensorflow layout.
+template <typename T>
+static void CheckCommonCasesForMklInputs(OpKernelContext* context,
+ gtl::InlinedVector<int64, 4>* begin,
+ gtl::InlinedVector<int64, 4>* size,
+ bool* done) {
+ bool is_identity = true;
+ *done = false;
+
+ ValidateMklInputs(context, &is_identity, begin, size);
+ if (!context->status().ok()) return;
+
+ const Tensor& input = MklGetInput(context, 0);
+ MklDnnShape input_mkl_shape;
+ GetMklShape(context, 0, &input_mkl_shape);
+
+ if (is_identity) {
+ VLOG(1) << "Slice identity";
+ context->set_output(0, input);
+ // Mkl metadata tensor in this case can just be forwarded from input to
+ // output.
+ AllocateOutputSetMklShape(context, 0, input_mkl_shape);
+ *done = true;
+ }
+}
+
+// MKL-DNN implementation of Slice
+template <typename Device, typename T>
+class MklDnnSliceOp : public OpKernel {
+ public:
+ explicit MklDnnSliceOp(OpKernelConstruction* context) : OpKernel(context) {}
+
+ ~MklDnnSliceOp() {}
+
+ void Compute(OpKernelContext* context) override {
+ gtl::InlinedVector<int64, 4> begin;
+ gtl::InlinedVector<int64, 4> size;
+ bool done = false;
+
+ CheckCommonCasesForMklInputs<T>(context, &begin, &size, &done);
+ if (!context->status().ok() || done == true) return;
+
+ // Though MKL-DNN supports more than 8 dimension and
+ // less than 12 dimension tensor.
+ // But we are mimicking functionality of Eigen Slice op for CPU.
+ if (begin.size() >= 8) {
+ OP_REQUIRES(
+ context, false,
+ errors::Unimplemented("MklDnnSliceOp : Unhandled input dimensions"));
+ }
+
+ ComputeMklDnnSlice(context, begin, size);
+ }
+
+ private:
+ // Slice op implemented using MKL-DNN APIs.
+ void ComputeMklDnnSlice(OpKernelContext* context,
+ const gtl::InlinedVector<int64, 4>& begin,
+ const gtl::InlinedVector<int64, 4>& size) {
+ try {
+ // MKL-DNN API usage below is guided by description at:
+ // https://github.com/01org/mkl-dnn/issues/69
+ //
+ // Relevant part of the description is copied below:
+ //
+ // Let's say you want to copy a part of memory into another buffer (and
+ // probably change the format). Then your steps are:
+ //
+ // 1. create memory primitive descriptor in_mem_pd and memory primitive
+ // in_mem_p for the entire source data.
+ // 2. create view primitive descriptor in_submem_pd based on in_mem_pd,
+ // initial offsets, and sub-sizes
+ // 3. create memory primitive descriptor out_mem_pd and memory primitive
+ // out_mem_p for the output (the logical sizes should match sub-sizes
+ // used in step 2, but the format might be arbitrary)
+ // 4. create reorder primitive descriptor reorder_pd based on in_submem_pd
+ // and out_mem_pd
+ // 5. create reorder primitive itself based on reorder_pd, in_mem_p, and
+ // out_mem_p.
+ //
+ // Please notice that there is no view primitive. There is only view
+ // primitive descriptor. And the reorder uses source memory as input but
+ // traverses it according to a view in_submem_pd.
+
+ auto cpu_engine = engine(engine::cpu, 0);
+ MklDnnData<T> src(&cpu_engine);
+ MklDnnData<T> output(&cpu_engine);
+
+ // Populate offsets and sizes in memory::dims format based on vector.
+ memory::dims begin_dims = {};
+ begin_dims.resize(begin.size());
+ for (size_t i = 0; i < begin.size(); ++i) begin_dims[i] = begin[i];
+ memory::dims size_dims = {};
+ bool empty = false;
+ size_dims.resize(size.size());
+ for (size_t i = 0; i < size.size(); ++i) {
+ size_dims[i] = size[i];
+ if (size_dims[i] == 0) empty = true;
+ }
+
+ Tensor* output_tensor = nullptr;
+ MklDnnShape output_mkl_shape;
+
+ // If no dimension is selected in slice, the result should be empty.
+ // Just return an empty output tensor, and a dummy Mkl-shape tensor.
+ if (empty) { // for empty dims
+ auto shape_to = MklDnnDimsToTFShape(size_dims);
+ AllocateOutputSetMklShape(context, 0, &output_tensor, shape_to,
+ output_mkl_shape);
+ return;
+ }
+
+ // Step 1 (as per above description) - Create memory for user data.
+ // We use blocked format here to describe input tensor.
+ const Tensor& input_tensor = MklGetInput(context, 0);
+ MklDnnShape input_mkl_shape;
+ GetMklShape(context, 0, &input_mkl_shape);
+
+ if (input_mkl_shape.IsMklTensor()) {
+ auto input_mkl_format = input_mkl_shape.GetTfDataFormat();
+ auto input_tf_format = MklDnnDataFormatToTFDataFormat(input_mkl_format);
+ begin_dims = MklDnnDimsInNCHW(begin_dims, input_tf_format);
+ size_dims = MklDnnDimsInNCHW(size_dims, input_tf_format);
+ auto input_md = input_mkl_shape.GetMklLayout();
+ src.SetUsrMem(input_md, &input_tensor);
+ } else {
+ // Initialize input dimensions and strides to be used when input is not
+ // in MklDnn layout.
+ memory::dims input_dims, input_strides;
+ input_dims = TFShapeToMklDnnDims(input_tensor.shape());
+ input_strides = CalculateTFStrides(input_dims);
+ // Create input memory descriptor.
+ auto input_md =
+ MklDnnData<T>::CreateBlockedMemDesc(input_dims, input_strides);
+ src.SetUsrMem(input_md, &input_tensor);
+ }
+
+ // Step 2 - create view primitive descriptor
+ auto view_pd =
+ view::primitive_desc(src.GetUsrMemPrimDesc(), size_dims, begin_dims)
+ .dst_primitive_desc();
+ auto output_strides = CalculateTFStrides(size_dims);
+ auto output_md =
+ MklDnnData<T>::CreateBlockedMemDesc(size_dims, output_strides);
+ auto output_pd = memory::primitive_desc(output_md, cpu_engine);
+
+ // Step 3 - Create memory for output. If input is in MklDnn layout, then
+ // output is also in MklDnn layout. Otherwise, output is in Tensorflow
+ // layout.
+ AllocateOutputTensor(context, input_mkl_shape, &output_pd, size_dims,
+ &output_tensor, &output_mkl_shape);
+ DCHECK(output_tensor);
+ DCHECK_EQ(input_mkl_shape.IsMklTensor(), output_mkl_shape.IsMklTensor());
+ output.SetUsrMem(output_md, output_tensor);
+
+ std::vector<primitive> net;
+ // Step 4 - create reorder primitive desc between view_pd and output_pd.
+ auto reorder_pd =
+ reorder::primitive_desc(view_pd, output.GetUsrMemPrimDesc());
+ // Step 5 - create reorder primitive itself.
+ net.push_back(reorder(reorder_pd, *src.GetUsrMem(), *output.GetUsrMem()));
+ // Execute the reorder primitive.
+ stream(stream::kind::eager).submit(net).wait();
+ } 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));
+ }
+ }
+
+ private:
+ void AllocateOutputTensor(OpKernelContext* context,
+ const MklDnnShape& input_mkl_shape,
+ memory::primitive_desc* output_pd,
+ const memory::dims& output_dims,
+ Tensor** output_tensor,
+ MklDnnShape* output_mkl_shape) {
+ DCHECK(output_tensor);
+ DCHECK(output_mkl_shape);
+
+ TensorShape output_tf_shape;
+
+ if (input_mkl_shape.IsMklTensor()) {
+ // Since input tensor is in Mkl layout, output tensor will be in Mkl
+ // layout.
+
+ // Allocate shape of Mkl tensor.
+ output_mkl_shape->SetMklTensor(true);
+ output_mkl_shape->SetMklLayout(output_pd);
+ output_mkl_shape->SetElemType(MklDnnType<T>());
+ output_mkl_shape->SetTfLayout(input_mkl_shape.GetDimension(), output_dims,
+ input_mkl_shape.GetTfDataFormat());
+
+ output_tf_shape.AddDim((output_pd->get_size() / sizeof(T)) + 1);
+ } else {
+ // If input is not in Mkl layout, then output won't be in Mkl layout.
+ output_mkl_shape->SetMklTensor(false);
+ output_tf_shape = MklDnnDimsToTFShape(output_dims);
+ }
+
+ AllocateOutputSetMklShape(context, 0, output_tensor, output_tf_shape,
+ *output_mkl_shape);
+ }
+};
+
+// MKL-DNN Slice registration
+#define REGISTER_MKL_SLICE(type) \
+ REGISTER_KERNEL_BUILDER(Name("_MklSlice") \
+ .Device(DEVICE_CPU) \
+ .TypeConstraint<type>("T") \
+ .HostMemory("begin") \
+ .HostMemory("size") \
+ .Label(mkl_op_registry::kMklOpLabel), \
+ MklDnnSliceOp<CPUDevice, type>);
+
+TF_CALL_float(REGISTER_MKL_SLICE);
+#undef REGISTER_MKL_SLICE
+
+} // namespace tensorflow
+
+#endif // INTEL_MKL_DNN
+#endif // INTEL_MKL
diff --git a/tensorflow/core/kernels/partitioned_function_ops.cc b/tensorflow/core/kernels/partitioned_function_ops.cc
index fc1c900..fdb4c84 100644
--- a/tensorflow/core/kernels/partitioned_function_ops.cc
+++ b/tensorflow/core/kernels/partitioned_function_ops.cc
@@ -97,7 +97,13 @@
OP_REQUIRES_ASYNC(ctx, fbody != nullptr,
errors::Internal("Could not find handle ", handle),
done);
+ // We need to pass global op_registry as default_registry when creating
+ // graph. So that graph optimization passes can lookup all possible ops
+ // by name.
auto graph = tensorflow::MakeUnique<Graph>(fbody->graph->flib_def());
+ FunctionLibraryDefinition global_flib(OpRegistry::Global(), {});
+ TF_CHECK_OK(
+ graph.get()->AddFunctionLibrary(global_flib.ToProto()));
CopyGraph(*fbody->graph, graph.get());
OP_REQUIRES_OK_ASYNC(ctx, PinResourceArgs(graph.get(), args), done);
@@ -250,9 +256,11 @@
VLOG(3) << "Partitioned function '" << func_.name() << "', yielding "
<< partitions.size() << " shards.";
- const FunctionLibraryDefinition* flib_def = &graph->flib_def();
for (const auto& partition : partitions) {
- std::unique_ptr<Graph> subgraph(new Graph(flib_def));
+ std::unique_ptr<Graph> subgraph(new Graph(graph->flib_def()));
+ FunctionLibraryDefinition global_flib(OpRegistry::Global(), {});
+ TF_CHECK_OK(
+ subgraph.get()->AddFunctionLibrary(global_flib.ToProto()));
GraphConstructorOptions opts;
opts.allow_internal_ops = true;
opts.expect_device_spec = true;
diff --git a/tensorflow/core/kernels/resource_variable_ops.cc b/tensorflow/core/kernels/resource_variable_ops.cc
index 26705a8..23d7698 100644
--- a/tensorflow/core/kernels/resource_variable_ops.cc
+++ b/tensorflow/core/kernels/resource_variable_ops.cc
@@ -51,7 +51,9 @@
#define EIGEN_USE_GPU
#endif
-#include "tensorflow/core/kernels/resource_variable_ops.h"
+#include <memory>
+#include <vector>
+
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/framework/resource_mgr.h"
@@ -60,10 +62,12 @@
#include "tensorflow/core/kernels/bounds_check.h"
#include "tensorflow/core/kernels/dense_update_functor.h"
#include "tensorflow/core/kernels/gather_functor.h"
+#include "tensorflow/core/kernels/resource_variable_ops.h"
#include "tensorflow/core/kernels/scatter_functor.h"
#include "tensorflow/core/kernels/training_op_helpers.h"
#include "tensorflow/core/kernels/variable_ops.h"
#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/refcount.h"
#include "tensorflow/core/platform/mem.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/types.h"
@@ -72,6 +76,8 @@
namespace tensorflow {
REGISTER_RESOURCE_HANDLE_KERNEL(Var);
+REGISTER_KERNEL_BUILDER(Name("_VarHandlesOp").Device(DEVICE_CPU),
+ ResourceHandlesOp<Var>);
ReadVariableOp::ReadVariableOp(OpKernelConstruction* c) : OpKernel(c) {
OP_REQUIRES_OK(c, c->GetAttr("dtype", &dtype_));
@@ -101,13 +107,58 @@
ctx->set_output(0, t);
}
+ReadVariablesOp::ReadVariablesOp(OpKernelConstruction* c) : OpKernel(c) {
+ int n;
+ OP_REQUIRES_OK(c, c->GetAttr("N", &n));
+ OP_REQUIRES_OK(c, c->GetAttr("dtypes", &dtypes_));
+ OP_REQUIRES(c, n == dtypes_.size(),
+ errors::InvalidArgument(
+ "Mismatched number of arguments to ReadVariablesOp (", n,
+ " vs. ", dtypes_.size(), ")"));
+}
+
+void ReadVariablesOp::Compute(OpKernelContext* ctx) {
+ std::vector<std::unique_ptr<Var, core::RefCountDeleter>> variables(
+ dtypes_.size());
+ std::vector<const ResourceHandle*> handles(dtypes_.size());
+ for (size_t i = 0; i < dtypes_.size(); ++i) {
+ handles[i] = &HandleFromInput(ctx, i);
+ }
+ const auto status = LookupResources(ctx, handles, &variables);
+ OP_REQUIRES(ctx, status.ok(),
+ errors::FailedPrecondition(
+ "Error while reading resource variable. This could mean that "
+ "the variable was uninitialized. ",
+ status.ToString()));
+
+ for (size_t i = 0; i < dtypes_.size(); ++i) {
+ // We're acquiring a reference to the underlying buffer while
+ // 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(),
+ 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);
+ }
+}
+
REGISTER_KERNEL_BUILDER(Name("ReadVariableOp").Device(DEVICE_CPU),
ReadVariableOp);
+REGISTER_KERNEL_BUILDER(Name("_ReadVariablesOp").Device(DEVICE_CPU),
+ ReadVariablesOp);
#if GOOGLE_CUDA
REGISTER_KERNEL_BUILDER(
Name("ReadVariableOp").Device(DEVICE_GPU).HostMemory("resource"),
ReadVariableOp);
+REGISTER_KERNEL_BUILDER(
+ Name("_ReadVariablesOp").Device(DEVICE_GPU).HostMemory("resources"),
+ ReadVariablesOp);
#define REGISTER_GPU_KERNELS(type) \
namespace functor { \
@@ -122,11 +173,20 @@
.HostMemory("resource") \
.TypeConstraint<type>("dtype"), \
ResourceHandleOp<Var>)
-
TF_CALL_GPU_ALL_TYPES(REGISTER_GPU_KERNELS);
TF_CALL_int64(REGISTER_GPU_KERNELS);
TF_CALL_variant(REGISTER_GPU_KERNELS);
#undef REGISTER_GPU_KERNELS
+
+REGISTER_KERNEL_BUILDER(Name("_VarHandlesOp")
+ .Device(DEVICE_GPU)
+ .HostMemory("resources")
+ .TypeConstraint("dtypes",
+ {DT_INT64, DT_COMPLEX64,
+ DT_COMPLEX128, DT_HALF, DT_FLOAT,
+ DT_DOUBLE, DT_BOOL, DT_VARIANT}),
+ ResourceHandlesOp<Var>);
+
#endif // GOOGLE_CUDA
template <typename T>
diff --git a/tensorflow/core/kernels/resource_variable_ops.h b/tensorflow/core/kernels/resource_variable_ops.h
index 9b60106..cffb732 100644
--- a/tensorflow/core/kernels/resource_variable_ops.h
+++ b/tensorflow/core/kernels/resource_variable_ops.h
@@ -28,6 +28,16 @@
DataType dtype_;
};
+class ReadVariablesOp : public OpKernel {
+ public:
+ explicit ReadVariablesOp(OpKernelConstruction* c);
+ void Compute(OpKernelContext* ctx) override;
+ bool IsExpensive() override { return false; }
+
+ private:
+ DataTypeVector dtypes_;
+};
+
class DestroyResourceOp : public OpKernel {
public:
explicit DestroyResourceOp(OpKernelConstruction* ctx);
diff --git a/tensorflow/core/kernels/slice_op.cc b/tensorflow/core/kernels/slice_op.cc
index 97f77e4..a006c69 100644
--- a/tensorflow/core/kernels/slice_op.cc
+++ b/tensorflow/core/kernels/slice_op.cc
@@ -228,191 +228,6 @@
}
};
-#ifdef INTEL_MKL
-template <typename Device, typename T>
-class MklSliceOp : public OpKernel {
- public:
- explicit MklSliceOp(OpKernelConstruction* context) : OpKernel(context) {}
-
- void Compute(OpKernelContext* context) override {
- TensorShape output_shape;
- gtl::InlinedVector<int64, 4> begin;
- gtl::InlinedVector<int64, 4> size;
- Tensor* result = nullptr;
- bool done = false;
- SharedSliceCommonCases<T>(context, &output_shape, &begin, &size, &result,
- &done);
- if (!context->status().ok() || done == true) return;
-
- const Tensor& input = context->input(0);
- const int input_dims = input.dims();
-
- if (output_shape.num_elements() > 0) {
- if (std::is_same<Device, CPUDevice>::value && input_dims == 2 &&
- DataTypeCanUseMemcpy(DataTypeToEnum<T>::v())) {
- auto input = context->input(0).tensor<T, 2>();
- auto output = result->tensor<T, 2>();
- // TODO(agarwal): Consider multi-threading this loop for cases where
- // size[0] is very large.
- for (int i = 0; i < size[0]; ++i) {
- const int64 row = begin[0] + i;
- if (i + 1 < size[0]) {
- port::prefetch<port::PREFETCH_HINT_T0>(&output(i + 1, 0));
- port::prefetch<port::PREFETCH_HINT_T0>(&input(row + 1, begin[1]));
- }
- memcpy(&output(i, 0), &input(row, begin[1]), size[1] * sizeof(T));
- }
- return;
- }
-#define HANDLE_DIM(NDIM) \
- if (input_dims == NDIM) { \
- HandleCase<NDIM>(context, begin, size, result); \
- return; \
- }
-
- HANDLE_DIM(1);
- HANDLE_DIM(2);
- HANDLE_DIM(3);
- HANDLE_DIM(4);
- HANDLE_DIM(5);
- HANDLE_DIM(6);
- HANDLE_DIM(7);
-
-#undef HANDLE_DIM
-
- OP_REQUIRES(
- context, false,
- errors::Unimplemented("SliceOp : Unhandled input dimensions"));
- }
- }
-
- private:
- // Helper function for DoesSliceShapeDifferInOnly1D. Checks if the following
- // criteria matches for slice_dim: if indices for slice are 0 in all dims
- // except slice_dim and if sizes of all the dimensions of the slice are same
- // as the sizes of all the dimensions of the input except slice_dim, then
- // returns True. Otherwise, returns False.
- bool DoesSliceShapeDifferInOnly1DHelper(const TensorShape& input_shape,
- const gtl::ArraySlice<int64>& begin,
- const gtl::ArraySlice<int64>& size,
- int slice_dim) {
- for (int dim = 0; dim < 4; dim++) {
- if (dim != slice_dim &&
- (begin[dim] != 0 || size[dim] != input_shape.dim_size(dim))) {
- return false;
- }
- }
- return true;
- }
-
- // Is 'input' tensor being sliced over a single dimension out of 4?
- //
- // This check is applicable in the context of Slice of a 4-D tensor in
- // NHWC or NCHW format over channel dimension.
- //
- // If indices for slice are 0 in all dims except one dimension and if sizes of
- // all dimensions of slice are same as sizes of all dimensions of inputs
- // except that dimension, then we are slicing over a single dimension.
- //
- // Returns True if Slicing over a single dimension, and sets slice_dim
- // to the number of the dimension that satisfies criteria.
- bool DoesSliceShapeDifferInOnly1D(const TensorShape& input_shape,
- const gtl::ArraySlice<int64>& begin,
- const gtl::ArraySlice<int64>& size,
- int* slice_dim) {
- for (int dim = 0; dim < 4; dim++) {
- if (DoesSliceShapeDifferInOnly1DHelper(input_shape, begin, size, dim)) {
- *slice_dim = dim;
- return true;
- }
- }
- return false;
- }
-
- template <int NDIM>
- void HandleCase(OpKernelContext* context, const gtl::ArraySlice<int64>& begin,
- const gtl::ArraySlice<int64>& size, Tensor* result) {
- int slice_dim = -1;
- TensorShape in_shape = context->input(0).shape();
- // Special case for handling 4-D tensor slice when shape of the slice
- // differs from the input tensor in only 1 out of 4 dimensions.
- // This case arises in the context of Slice of 4-D tensor in NHWC or NCHW
- // format over channel dimension.
- if (NDIM == 4 &&
- DoesSliceShapeDifferInOnly1D(in_shape, begin, size, &slice_dim)) {
- size_t in_strides[4] = {
- (size_t)in_shape.dim_size(1) * in_shape.dim_size(2) *
- in_shape.dim_size(3),
- (size_t)in_shape.dim_size(2) * in_shape.dim_size(3),
- (size_t)in_shape.dim_size(3), (size_t)1};
-
- size_t out_strides[4] = {(size_t)size[1] * size[2] * size[3],
- (size_t)size[2] * size[3], (size_t)size[3],
- (size_t)1};
-
- T* in_buf = const_cast<T*>(
- const_cast<const T*>(context->input(0).flat<T>().data()));
- T* op_buf = result->flat<T>().data();
-
- if (slice_dim == 1) {
- /* data format = NCHW */
-
-#pragma omp parallel for
- for (ssize_t d0 = begin[0]; d0 < begin[0] + size[0]; d0++) {
- T* ip = in_buf + (d0 * in_strides[0]);
- T* op = op_buf + ((d0 - begin[0]) * out_strides[0]);
-#pragma omp parallel for
- for (ssize_t d1 = begin[1]; d1 < begin[1] + size[1]; d1++) {
- T* ip1 = ip + (d1 * in_strides[1]);
- T* op1 = op + ((d1 - begin[1]) * out_strides[1]);
- // For NCHW, H and W will be contiguous. So we can copy
- // both with one memcpy.
- memcpy(static_cast<void*>(op1), static_cast<void*>(ip1),
- sizeof(T) * in_strides[1]);
- }
- }
- return;
- } else if (slice_dim == 3) {
- /* data_format = NHWC */
-
-#pragma omp parallel for
- for (ssize_t d0 = begin[0]; d0 < begin[0] + size[0]; d0++) {
- T* ip = in_buf + (d0 * in_strides[0]);
- T* op = op_buf + ((d0 - begin[0]) * out_strides[0]);
-#pragma omp parallel for
- for (ssize_t d1 = begin[1]; d1 < begin[1] + size[1]; d1++) {
- T* ip1 = ip + (d1 * in_strides[1]);
- T* op1 = op + ((d1 - begin[1]) * out_strides[1]);
-#pragma omp parallel for
- for (ssize_t d2 = begin[2]; d2 < begin[2] + size[2]; d2++) {
- T* ip2 = ip1 + (d2 * in_strides[2]);
- T* ip3 = ip2 + begin[3];
- T* op2 = op1 + ((d2 - begin[2]) * out_strides[2]);
- T* op3 = op2;
- memcpy(static_cast<void*>(op3), static_cast<void*>(ip3),
- sizeof(T) * size[3]);
- }
- }
- }
- return;
- }
- // slice_dim is not 1 or 3, then we fallback to Eigen implementation.
- }
-
- Eigen::DSizes<Eigen::DenseIndex, NDIM> indices;
- Eigen::DSizes<Eigen::DenseIndex, NDIM> sizes;
- for (int i = 0; i < NDIM; ++i) {
- indices[i] = begin[i];
- sizes[i] = size[i];
- }
-
- functor::Slice<Device, T, NDIM>()(
- context->eigen_device<Device>(), result->tensor<T, NDIM>(),
- context->input(0).tensor<T, NDIM>(), indices, sizes);
- }
-};
-#endif // INTEL_MKL
-
// Forward declarations of the functor specializations for declared in the
// sharded source files.
namespace functor {
@@ -440,15 +255,6 @@
#undef DECLARE_CPU_SPEC
} // namespace functor
-#if defined(INTEL_MKL) && defined(ENABLE_MKL)
-#define REGISTER_SLICE(type) \
- REGISTER_KERNEL_BUILDER(Name("Slice") \
- .Device(DEVICE_CPU) \
- .TypeConstraint<type>("T") \
- .HostMemory("begin") \
- .HostMemory("size"), \
- MklSliceOp<CPUDevice, type>)
-#else
#define REGISTER_SLICE(type) \
REGISTER_KERNEL_BUILDER(Name("Slice") \
.Device(DEVICE_CPU) \
@@ -456,7 +262,6 @@
.HostMemory("begin") \
.HostMemory("size"), \
SliceOp<CPUDevice, type>)
-#endif // INTEL_MKL && ENABLE_MKL
TF_CALL_POD_STRING_TYPES(REGISTER_SLICE);
TF_CALL_QUANTIZED_TYPES(REGISTER_SLICE);
diff --git a/tensorflow/core/kernels/strided_slice_op.cc b/tensorflow/core/kernels/strided_slice_op.cc
index f0575de..3e8a4c5 100644
--- a/tensorflow/core/kernels/strided_slice_op.cc
+++ b/tensorflow/core/kernels/strided_slice_op.cc
@@ -149,7 +149,7 @@
// NDIM and T
if (is_simple_slice && std::is_same<Device, CPUDevice>::value &&
input_dims == 2 && processing_shape.dims() == 2 &&
- final_shape.dims() == 2) {
+ final_shape.dims() == 2 && new_axis_mask == 0) {
MemCpyFunctor<T> functor;
if (functor.Copy(input, begin, end, result)) {
return;
diff --git a/tensorflow/core/kernels/training_op_helpers.cc b/tensorflow/core/kernels/training_op_helpers.cc
index 83b83fc..4262a54 100644
--- a/tensorflow/core/kernels/training_op_helpers.cc
+++ b/tensorflow/core/kernels/training_op_helpers.cc
@@ -15,14 +15,16 @@
#include "tensorflow/core/kernels/training_op_helpers.h"
+#include "tensorflow/core/util/ptr_util.h"
+
namespace tensorflow {
-mutex* GetTrainingVariableMutex(OpKernelContext* ctx, int input) {
+mutex* GetTrainingVariableMutex(OpKernelContext* ctx, int input,
+ Var** maybe_resource) {
+ *maybe_resource = nullptr;
if (ctx->input_dtype(input) == DT_RESOURCE) {
- Var* var;
- if (LookupResource(ctx, HandleFromInput(ctx, input), &var).ok()) {
- core::ScopedUnref scoped_unref(var);
- return var->mu();
+ if (LookupResource(ctx, HandleFromInput(ctx, input), maybe_resource).ok()) {
+ return (*maybe_resource)->mu();
} else {
ctx->CtxFailureWithWarning(
errors::Internal("Invalid variable reference."));
@@ -33,12 +35,13 @@
}
// MaybeLockVariableInputMutexesInOrder is a helper function to acquire mutexes
-// in address order to mitigate deadlock. Returns a vector of 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.
-std::vector<mutex_lock> MaybeLockVariableInputMutexesInOrder(
+// 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) {
@@ -47,14 +50,16 @@
break;
}
}
- std::vector<mutex_lock> locks;
if (!do_lock && !any_resource) {
- return locks;
+ return VariableInputLockHolder({}, {});
}
+ std::vector<Var*> vars;
std::vector<mutex*> mutexes;
std::vector<int> acquire_order;
for (auto input : input_ids) {
- mutex* mutex = GetTrainingVariableMutex(ctx, input);
+ 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());
@@ -64,13 +69,19 @@
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) {
- mutex* mu = GetTrainingVariableMutex(ctx, input);
+ Var* var;
+ mutex* mu = GetTrainingVariableMutex(ctx, input, &var);
+ core::ScopedUnref scoped_unref(var);
if (mu != nullptr) {
- locks.emplace_back(*mu);
+ locks->emplace_back(*mu);
}
}
- return locks;
+ return VariableInputLockHolder(std::move(vars), std::move(locks));
}
void MaybeForwardRefInputToRefOutput(OpKernelContext* ctx, int input,
diff --git a/tensorflow/core/kernels/training_op_helpers.h b/tensorflow/core/kernels/training_op_helpers.h
index 071cb37..9f173a8 100644
--- a/tensorflow/core/kernels/training_op_helpers.h
+++ b/tensorflow/core/kernels/training_op_helpers.h
@@ -23,9 +23,42 @@
namespace tensorflow {
-mutex* GetTrainingVariableMutex(OpKernelContext* ctx, int input);
+// 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);
-std::vector<mutex_lock> MaybeLockVariableInputMutexesInOrder(
+// 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(VariableInputLockHolder&& other)
+ : vars_(std::move(other.vars_)), locks_(std::move(other.locks_)) {}
+
+ ~VariableInputLockHolder() {
+ // Release the locks before unreffing the Vars, because each lock
+ // is potentially borrowed from a Var in vars_.
+ locks_.reset();
+ for (Var* var : vars_) {
+ var->Unref();
+ }
+ }
+
+ private:
+ std::vector<Var*> vars_;
+ // 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_;
+};
+
+VariableInputLockHolder MaybeLockVariableInputMutexesInOrder(
OpKernelContext* ctx, bool do_lock, const std::vector<int>& input_ids);
void MaybeForwardRefInputToRefOutput(OpKernelContext* ctx, int input,
diff --git a/tensorflow/core/kernels/training_ops.cc b/tensorflow/core/kernels/training_ops.cc
index 9a07ded..acf162d 100644
--- a/tensorflow/core/kernels/training_ops.cc
+++ b/tensorflow/core/kernels/training_ops.cc
@@ -561,7 +561,9 @@
}
void Compute(OpKernelContext* ctx) override {
- mutex* mu = GetTrainingVariableMutex(ctx, 0);
+ Var* resource;
+ mutex* mu = GetTrainingVariableMutex(ctx, 0, &resource);
+ core::ScopedUnref scoped_unref(resource);
if (use_exclusive_lock_ && mu != nullptr) {
mutex_lock l1(*mu);
// Don't try to acquire a lock on the second ref as they share the same
@@ -710,7 +712,9 @@
}
void Compute(OpKernelContext* ctx) override {
- mutex* mu = GetTrainingVariableMutex(ctx, 0);
+ Var* var;
+ mutex* mu = GetTrainingVariableMutex(ctx, 0, &var);
+ core::ScopedUnref scoped_unref(var);
// mu_accum is actually the same mutex as mu_var since currently we use a
// global mutex.
//
diff --git a/tensorflow/core/kernels/unicode_script_op.cc b/tensorflow/core/kernels/unicode_script_op.cc
new file mode 100644
index 0000000..085e397
--- /dev/null
+++ b/tensorflow/core/kernels/unicode_script_op.cc
@@ -0,0 +1,53 @@
+/* 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 "unicode/errorcode.h" // TF:icu
+#include "unicode/uscript.h" // TF:icu
+#include "tensorflow/core/framework/op_kernel.h"
+
+namespace tensorflow {
+
+class UnicodeScriptOp : public OpKernel {
+ public:
+ explicit UnicodeScriptOp(OpKernelConstruction* context) : OpKernel(context) {}
+
+ void Compute(OpKernelContext* context) override {
+ const Tensor* input_tensor;
+ OP_REQUIRES_OK(context, context->input("input", &input_tensor));
+ const auto& input_flat = input_tensor->flat<int32>();
+
+ Tensor* output_tensor = nullptr;
+ OP_REQUIRES_OK(context,
+ context->allocate_output("output", input_tensor->shape(),
+ &output_tensor));
+ auto output_flat = output_tensor->flat<int32>();
+
+ icu::ErrorCode status;
+ for (int i = 0; i < input_flat.size(); i++) {
+ UScriptCode script_code = uscript_getScript(input_flat(i), status);
+ if (status.isSuccess()) {
+ output_flat(i) = script_code;
+ } else {
+ output_flat(i) = -1;
+ status.reset();
+ }
+ }
+ }
+};
+
+REGISTER_KERNEL_BUILDER(Name("UnicodeScript").Device(DEVICE_CPU),
+ UnicodeScriptOp);
+
+} // namespace tensorflow
diff --git a/tensorflow/core/ops/array_ops.cc b/tensorflow/core/ops/array_ops.cc
index 442686c..c9f80df 100644
--- a/tensorflow/core/ops/array_ops.cc
+++ b/tensorflow/core/ops/array_ops.cc
@@ -1531,37 +1531,6 @@
.Attr("out_type: {int32, int64} = DT_INT32")
.SetShapeFn(shape_inference::ScalarShape);
-namespace {
-
-// This SliceHelper processes the output shape of the `slice`
-// when the tensor of `sizes` is available.
-template <typename T>
-Status SliceHelper(InferenceContext* c, ShapeHandle begin_value,
- const Tensor* sizes_value,
- std::vector<DimensionHandle>* dims) {
- auto sizes_vec = sizes_value->vec<T>();
- for (int i = 0; i < sizes_value->NumElements(); ++i) {
- DimensionHandle dim = c->Dim(c->input(0), i);
- if (sizes_vec(i) != -1) {
- auto dim_val = c->Value(dim);
- if (sizes_vec(i) < 0) {
- return errors::InvalidArgument(
- "Out of bounds slicing on dimension ", i, " of length ", dim_val,
- ": sizes vector cannot be < -1, but was ", sizes_vec(i));
- }
-
- dims->emplace_back(c->MakeDim(sizes_vec(i)));
- } else {
- DimensionHandle result;
- TF_RETURN_IF_ERROR(c->Subtract(dim, c->Dim(begin_value, i), &result));
- dims->emplace_back(result);
- }
- }
-
- return Status::OK();
-}
-} // namespace
-
// --------------------------------------------------------------------------
REGISTER_OP("Slice")
.Input("input: T")
@@ -1570,83 +1539,22 @@
.Output("output: T")
.Attr("T: type")
.Attr("Index: {int32,int64}")
- .SetShapeFn([](InferenceContext* c) {
- ShapeHandle input = c->input(0);
- ShapeHandle begin_shape;
- TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 1, &begin_shape));
- ShapeHandle sizes_shape;
- TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 1, &sizes_shape));
+ .SetShapeFn(shape_inference::SliceShape);
- // Merge to check compatibility of begin and sizes tensors.
- TF_RETURN_IF_ERROR(c->Merge(begin_shape, sizes_shape, &begin_shape));
-
- DimensionHandle ndims = c->Dim(begin_shape, 0);
- if (c->ValueKnown(ndims)) {
- TF_RETURN_IF_ERROR(c->WithRank(input, c->Value(ndims), &input));
- }
-
- // NOTE(mrry): Use MakeShapeFromShapeTensor to handle partially-known
- // values, even though the `begin` value does not represent a shape.
- ShapeHandle begin_value;
- TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensor(1, &begin_value));
-
- // We check the tensor value here and will only use
- // `MakeShapeFromShapeTensor` when `sizes_value` is null.
- // The reason is that `sizes`might contain -1, which can't
- // be represented (-1 in the ShapeHandle would mean "unknown".
- const Tensor* sizes_value = c->input_tensor(2);
-
- if (sizes_value != nullptr) {
- TF_RETURN_IF_ERROR(
- c->WithRank(begin_value, sizes_value->NumElements(), &begin_value));
- std::vector<DimensionHandle> dims;
- // If the begin and sizes tensors are available, then
- // we can be precise about the shape of the output.
- if (sizes_value->dtype() == DT_INT64) {
- TF_RETURN_IF_ERROR(
- SliceHelper<int64>(c, begin_value, sizes_value, &dims));
- } else {
- TF_RETURN_IF_ERROR(
- SliceHelper<int32>(c, begin_value, sizes_value, &dims));
- }
-
- c->set_output(0, c->MakeShape(dims));
- return Status::OK();
- } else {
- // In case `sizes` is not available (`sizes_value` is null),
- // we could try to use `MakeShapeFromShapeTensor` here.
- // If sizes contain -1, we will simply consider it as `Unknown`.
- // This is less than ideal but still an improvement of shape inference.
- // The following is an example that returns [None, 1, None] with this
- // code path:
- // z = tf.zeros((1, 2, 3))
- // m = tf.slice(z, [0, 0, 0], [tf.constant(1) + 0, 1, -1])
- // m.get_shape().as_list()
- ShapeHandle sizes_value;
- TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensor(2, &sizes_value));
- if (c->RankKnown(sizes_value)) {
- TF_RETURN_IF_ERROR(
- c->WithRank(begin_value, c->Rank(sizes_value), &begin_value));
- std::vector<DimensionHandle> dims;
- dims.reserve(c->Rank(sizes_value));
- for (int i = 0; i < c->Rank(sizes_value); ++i) {
- dims.emplace_back(c->Dim(sizes_value, i));
- }
- c->set_output(0, c->MakeShape(dims));
- return Status::OK();
- }
-
- // We might know the rank of the input.
- if (c->RankKnown(input)) {
- c->set_output(0, c->UnknownShapeOfRank(c->Rank(input)));
- return Status::OK();
- } else {
- return shape_inference::UnknownShape(c);
- }
- }
-
- return Status::OK();
- });
+#ifdef INTEL_MKL
+REGISTER_OP("_MklSlice")
+ .Input("input: T")
+ .Input("begin: Index")
+ .Input("size: Index")
+ .Input("mkl_input: uint8")
+ .Input("mkl_begin: uint8")
+ .Input("mkl_size: uint8")
+ .Output("output: T")
+ .Output("mkl_output: uint8")
+ .Attr("T: type")
+ .Attr("Index: {int32,int64}")
+ .SetShapeFn(shape_inference::SliceShape);
+#endif
REGISTER_OP("StridedSlice")
.Input("input: T")
diff --git a/tensorflow/core/ops/compat/ops_history.v1.pbtxt b/tensorflow/core/ops/compat/ops_history.v1.pbtxt
index 86d4c6b..43c14d8 100644
--- a/tensorflow/core/ops/compat/ops_history.v1.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history.v1.pbtxt
@@ -21532,6 +21532,421 @@
}
}
op {
+ name: "ExperimentalAssertNextDataset"
+ input_arg {
+ name: "input_dataset"
+ type: DT_VARIANT
+ }
+ input_arg {
+ name: "transformations"
+ type: DT_STRING
+ }
+ output_arg {
+ name: "handle"
+ type: DT_VARIANT
+ }
+ attr {
+ name: "output_types"
+ type: "list(type)"
+ has_minimum: true
+ minimum: 1
+ }
+ attr {
+ name: "output_shapes"
+ type: "list(shape)"
+ has_minimum: true
+ minimum: 1
+ }
+}
+op {
+ name: "ExperimentalCSVDataset"
+ input_arg {
+ name: "filenames"
+ type: DT_STRING
+ }
+ input_arg {
+ name: "compression_type"
+ type: DT_STRING
+ }
+ input_arg {
+ name: "buffer_size"
+ type: DT_INT64
+ }
+ input_arg {
+ name: "header"
+ type: DT_BOOL
+ }
+ input_arg {
+ name: "field_delim"
+ type: DT_STRING
+ }
+ input_arg {
+ name: "use_quote_delim"
+ type: DT_BOOL
+ }
+ input_arg {
+ name: "na_value"
+ type: DT_STRING
+ }
+ input_arg {
+ name: "select_cols"
+ type: DT_INT64
+ }
+ input_arg {
+ name: "record_defaults"
+ type_list_attr: "output_types"
+ }
+ output_arg {
+ name: "handle"
+ type: DT_VARIANT
+ }
+ attr {
+ name: "output_types"
+ type: "list(type)"
+ has_minimum: true
+ minimum: 1
+ allowed_values {
+ list {
+ type: DT_FLOAT
+ type: DT_DOUBLE
+ type: DT_INT32
+ type: DT_INT64
+ type: DT_STRING
+ }
+ }
+ }
+ attr {
+ name: "output_shapes"
+ type: "list(shape)"
+ has_minimum: true
+ minimum: 1
+ }
+ is_stateful: true
+}
+op {
+ name: "ExperimentalDirectedInterleaveDataset"
+ input_arg {
+ name: "selector_input_dataset"
+ type: DT_VARIANT
+ }
+ input_arg {
+ name: "data_input_datasets"
+ type: DT_VARIANT
+ number_attr: "N"
+ }
+ output_arg {
+ name: "handle"
+ type: DT_VARIANT
+ }
+ 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: "N"
+ type: "int"
+ has_minimum: true
+ minimum: 1
+ }
+}
+op {
+ name: "ExperimentalFunctionBufferingResource"
+ input_arg {
+ name: "string_arg"
+ type: DT_STRING
+ }
+ input_arg {
+ name: "target_device"
+ type: DT_STRING
+ }
+ output_arg {
+ name: "resource"
+ type: DT_RESOURCE
+ }
+ attr {
+ name: "shared_name"
+ type: "string"
+ }
+ attr {
+ name: "container"
+ type: "string"
+ }
+ attr {
+ name: "f"
+ type: "func"
+ }
+ attr {
+ name: "buffer_size"
+ type: "int"
+ }
+ attr {
+ name: "output_types"
+ type: "list(type)"
+ }
+ is_stateful: true
+}
+op {
+ name: "ExperimentalFunctionBufferingResourceGetNext"
+ input_arg {
+ name: "function_buffer_resource"
+ type: DT_RESOURCE
+ }
+ output_arg {
+ name: "output"
+ type_list_attr: "output_types"
+ }
+ attr {
+ name: "output_types"
+ type: "list(type)"
+ has_minimum: true
+ minimum: 1
+ }
+ is_stateful: true
+}
+op {
+ name: "ExperimentalFunctionBufferingResourceReset"
+ input_arg {
+ name: "function_buffer_resource"
+ type: DT_RESOURCE
+ }
+ is_stateful: true
+}
+op {
+ name: "ExperimentalIdentityIndexedDataset"
+ input_arg {
+ name: "size"
+ type: DT_UINT64
+ }
+ output_arg {
+ name: "handle"
+ type: DT_VARIANT
+ }
+ is_stateful: true
+}
+op {
+ name: "ExperimentalIgnoreErrorsDataset"
+ input_arg {
+ name: "input_dataset"
+ type: DT_VARIANT
+ }
+ output_arg {
+ name: "handle"
+ type: DT_VARIANT
+ }
+ attr {
+ name: "output_types"
+ type: "list(type)"
+ has_minimum: true
+ minimum: 1
+ }
+ attr {
+ name: "output_shapes"
+ type: "list(shape)"
+ has_minimum: true
+ minimum: 1
+ }
+}
+op {
+ name: "ExperimentalIndexedDatasetGet"
+ input_arg {
+ name: "materialized"
+ type: DT_RESOURCE
+ }
+ input_arg {
+ name: "index"
+ type: DT_UINT64
+ }
+ output_arg {
+ name: "components"
+ type_list_attr: "output_types"
+ }
+ attr {
+ name: "output_types"
+ type: "list(type)"
+ has_minimum: true
+ minimum: 1
+ }
+ attr {
+ name: "output_shapes"
+ type: "list(shape)"
+ has_minimum: true
+ minimum: 1
+ }
+ is_stateful: true
+}
+op {
+ name: "ExperimentalIndexedDatasetMaterialize"
+ input_arg {
+ name: "dataset"
+ type: DT_VARIANT
+ }
+ input_arg {
+ name: "materialized"
+ type: DT_RESOURCE
+ }
+ is_stateful: true
+}
+op {
+ name: "ExperimentalIteratorGetDevice"
+ input_arg {
+ name: "resource"
+ type: DT_RESOURCE
+ }
+ output_arg {
+ name: "device"
+ type: DT_STRING
+ }
+ is_stateful: true
+}
+op {
+ name: "ExperimentalLMDBDataset"
+ input_arg {
+ name: "filenames"
+ type: DT_STRING
+ }
+ output_arg {
+ name: "handle"
+ type: DT_VARIANT
+ }
+ attr {
+ name: "output_types"
+ type: "list(type)"
+ has_minimum: true
+ minimum: 1
+ }
+ attr {
+ name: "output_shapes"
+ type: "list(shape)"
+ has_minimum: true
+ minimum: 1
+ }
+ is_stateful: true
+}
+op {
+ name: "ExperimentalMaterializedIndexDatasetHandle"
+ output_arg {
+ name: "handle"
+ type: DT_RESOURCE
+ }
+ attr {
+ name: "container"
+ type: "string"
+ }
+ attr {
+ name: "shared_name"
+ type: "string"
+ }
+ attr {
+ name: "output_types"
+ type: "list(type)"
+ has_minimum: true
+ minimum: 1
+ }
+ attr {
+ name: "output_shapes"
+ type: "list(shape)"
+ has_minimum: true
+ minimum: 1
+ }
+ is_stateful: true
+}
+op {
+ name: "ExperimentalThreadPoolDataset"
+ input_arg {
+ name: "input_dataset"
+ type: DT_VARIANT
+ }
+ input_arg {
+ name: "thread_pool"
+ type: DT_RESOURCE
+ }
+ output_arg {
+ name: "handle"
+ type: DT_VARIANT
+ }
+ attr {
+ name: "output_types"
+ type: "list(type)"
+ has_minimum: true
+ minimum: 1
+ }
+ attr {
+ name: "output_shapes"
+ type: "list(shape)"
+ has_minimum: true
+ minimum: 1
+ }
+ is_stateful: true
+}
+op {
+ name: "ExperimentalThreadPoolHandle"
+ output_arg {
+ name: "handle"
+ type: DT_RESOURCE
+ }
+ attr {
+ name: "num_threads"
+ type: "int"
+ }
+ attr {
+ name: "max_intra_op_parallelism"
+ type: "int"
+ default_value {
+ i: 1
+ }
+ }
+ attr {
+ name: "display_name"
+ type: "string"
+ }
+ attr {
+ name: "container"
+ type: "string"
+ default_value {
+ s: ""
+ }
+ }
+ attr {
+ name: "shared_name"
+ type: "string"
+ default_value {
+ s: ""
+ }
+ }
+ is_stateful: true
+}
+op {
+ name: "ExperimentalUniqueDataset"
+ input_arg {
+ name: "input_dataset"
+ type: DT_VARIANT
+ }
+ output_arg {
+ name: "handle"
+ type: DT_VARIANT
+ }
+ attr {
+ name: "output_types"
+ type: "list(type)"
+ has_minimum: true
+ minimum: 1
+ }
+ attr {
+ name: "output_shapes"
+ type: "list(shape)"
+ has_minimum: true
+ minimum: 1
+ }
+}
+op {
name: "Expm1"
input_arg {
name: "x"
@@ -24105,6 +24520,85 @@
}
}
op {
+ name: "FusedBatchNorm"
+ input_arg {
+ name: "x"
+ type_attr: "T"
+ }
+ input_arg {
+ name: "scale"
+ type_attr: "T"
+ }
+ input_arg {
+ name: "offset"
+ type_attr: "T"
+ }
+ input_arg {
+ name: "mean"
+ type_attr: "T"
+ }
+ input_arg {
+ name: "variance"
+ type_attr: "T"
+ }
+ output_arg {
+ name: "y"
+ type_attr: "T"
+ }
+ output_arg {
+ name: "batch_mean"
+ type_attr: "T"
+ }
+ output_arg {
+ name: "batch_variance"
+ type_attr: "T"
+ }
+ output_arg {
+ name: "reserve_space_1"
+ type_attr: "T"
+ }
+ output_arg {
+ name: "reserve_space_2"
+ type_attr: "T"
+ }
+ attr {
+ name: "T"
+ type: "type"
+ allowed_values {
+ list {
+ type: DT_FLOAT
+ }
+ }
+ }
+ attr {
+ name: "epsilon"
+ type: "float"
+ default_value {
+ f: 0.0001
+ }
+ }
+ attr {
+ name: "data_format"
+ type: "string"
+ default_value {
+ s: "NHWC"
+ }
+ allowed_values {
+ list {
+ s: "NHWC"
+ s: "NCHW"
+ }
+ }
+ }
+ attr {
+ name: "is_training"
+ type: "bool"
+ default_value {
+ b: true
+ }
+ }
+}
+op {
name: "FusedBatchNormGrad"
input_arg {
name: "y_backprop"
@@ -24178,6 +24672,85 @@
}
}
op {
+ name: "FusedBatchNormGrad"
+ input_arg {
+ name: "y_backprop"
+ type_attr: "T"
+ }
+ input_arg {
+ name: "x"
+ type_attr: "T"
+ }
+ input_arg {
+ name: "scale"
+ type_attr: "T"
+ }
+ input_arg {
+ name: "reserve_space_1"
+ type_attr: "T"
+ }
+ input_arg {
+ name: "reserve_space_2"
+ type_attr: "T"
+ }
+ output_arg {
+ name: "x_backprop"
+ type_attr: "T"
+ }
+ output_arg {
+ name: "scale_backprop"
+ type_attr: "T"
+ }
+ output_arg {
+ name: "offset_backprop"
+ type_attr: "T"
+ }
+ output_arg {
+ name: "reserve_space_3"
+ type_attr: "T"
+ }
+ output_arg {
+ name: "reserve_space_4"
+ type_attr: "T"
+ }
+ attr {
+ name: "T"
+ type: "type"
+ allowed_values {
+ list {
+ type: DT_FLOAT
+ }
+ }
+ }
+ attr {
+ name: "epsilon"
+ type: "float"
+ default_value {
+ f: 0.0001
+ }
+ }
+ attr {
+ name: "data_format"
+ type: "string"
+ default_value {
+ s: "NHWC"
+ }
+ allowed_values {
+ list {
+ s: "NHWC"
+ s: "NCHW"
+ }
+ }
+ }
+ attr {
+ name: "is_training"
+ type: "bool"
+ default_value {
+ b: true
+ }
+ }
+}
+op {
name: "FusedBatchNormGradV2"
input_arg {
name: "y_backprop"
@@ -24345,6 +24918,96 @@
}
}
op {
+ name: "FusedBatchNormGradV2"
+ input_arg {
+ name: "y_backprop"
+ type_attr: "T"
+ }
+ input_arg {
+ name: "x"
+ type_attr: "T"
+ }
+ input_arg {
+ name: "scale"
+ type: DT_FLOAT
+ }
+ input_arg {
+ name: "reserve_space_1"
+ type_attr: "U"
+ }
+ input_arg {
+ name: "reserve_space_2"
+ type_attr: "U"
+ }
+ output_arg {
+ name: "x_backprop"
+ type_attr: "T"
+ }
+ output_arg {
+ name: "scale_backprop"
+ type_attr: "U"
+ }
+ output_arg {
+ name: "offset_backprop"
+ type_attr: "U"
+ }
+ output_arg {
+ name: "reserve_space_3"
+ type_attr: "U"
+ }
+ output_arg {
+ name: "reserve_space_4"
+ type_attr: "U"
+ }
+ attr {
+ name: "T"
+ type: "type"
+ allowed_values {
+ list {
+ type: DT_HALF
+ type: DT_BFLOAT16
+ type: DT_FLOAT
+ }
+ }
+ }
+ attr {
+ name: "U"
+ type: "type"
+ allowed_values {
+ list {
+ type: DT_FLOAT
+ }
+ }
+ }
+ attr {
+ name: "epsilon"
+ type: "float"
+ default_value {
+ f: 0.0001
+ }
+ }
+ attr {
+ name: "data_format"
+ type: "string"
+ default_value {
+ s: "NHWC"
+ }
+ allowed_values {
+ list {
+ s: "NHWC"
+ s: "NCHW"
+ }
+ }
+ }
+ attr {
+ name: "is_training"
+ type: "bool"
+ default_value {
+ b: true
+ }
+ }
+}
+op {
name: "FusedBatchNormV2"
input_arg {
name: "x"
@@ -24512,6 +25175,96 @@
}
}
op {
+ name: "FusedBatchNormV2"
+ input_arg {
+ name: "x"
+ type_attr: "T"
+ }
+ input_arg {
+ name: "scale"
+ type_attr: "U"
+ }
+ input_arg {
+ name: "offset"
+ type_attr: "U"
+ }
+ input_arg {
+ name: "mean"
+ type_attr: "U"
+ }
+ input_arg {
+ name: "variance"
+ type_attr: "U"
+ }
+ output_arg {
+ name: "y"
+ type_attr: "T"
+ }
+ output_arg {
+ name: "batch_mean"
+ type_attr: "U"
+ }
+ output_arg {
+ name: "batch_variance"
+ type_attr: "U"
+ }
+ output_arg {
+ name: "reserve_space_1"
+ type_attr: "U"
+ }
+ output_arg {
+ name: "reserve_space_2"
+ type_attr: "U"
+ }
+ attr {
+ name: "T"
+ type: "type"
+ allowed_values {
+ list {
+ type: DT_HALF
+ type: DT_BFLOAT16
+ type: DT_FLOAT
+ }
+ }
+ }
+ attr {
+ name: "U"
+ type: "type"
+ allowed_values {
+ list {
+ type: DT_FLOAT
+ }
+ }
+ }
+ attr {
+ name: "epsilon"
+ type: "float"
+ default_value {
+ f: 0.0001
+ }
+ }
+ attr {
+ name: "data_format"
+ type: "string"
+ default_value {
+ s: "NHWC"
+ }
+ allowed_values {
+ list {
+ s: "NHWC"
+ s: "NCHW"
+ }
+ }
+ }
+ attr {
+ name: "is_training"
+ type: "bool"
+ default_value {
+ b: true
+ }
+ }
+}
+op {
name: "FusedPadConv2D"
input_arg {
name: "input"
@@ -44518,6 +45271,59 @@
is_stateful: true
}
op {
+ name: "ReduceDataset"
+ 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: "components"
+ type_list_attr: "output_types"
+ }
+ 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: "use_inter_op_parallelism"
+ type: "bool"
+ default_value {
+ b: true
+ }
+ }
+}
+op {
name: "ReduceJoin"
input_arg {
name: "inputs"
@@ -74573,6 +75379,17 @@
}
}
op {
+ name: "UnicodeScript"
+ input_arg {
+ name: "input"
+ type: DT_INT32
+ }
+ output_arg {
+ name: "output"
+ type: DT_INT32
+ }
+}
+op {
name: "UniformCandidateSampler"
input_arg {
name: "true_classes"
@@ -76283,6 +77100,62 @@
is_stateful: true
}
op {
+ name: "Xdivy"
+ input_arg {
+ name: "x"
+ type_attr: "T"
+ }
+ input_arg {
+ name: "y"
+ type_attr: "T"
+ }
+ output_arg {
+ name: "z"
+ type_attr: "T"
+ }
+ attr {
+ name: "T"
+ type: "type"
+ allowed_values {
+ list {
+ type: DT_HALF
+ type: DT_FLOAT
+ type: DT_DOUBLE
+ type: DT_COMPLEX64
+ type: DT_COMPLEX128
+ }
+ }
+ }
+}
+op {
+ name: "Xlogy"
+ input_arg {
+ name: "x"
+ type_attr: "T"
+ }
+ input_arg {
+ name: "y"
+ type_attr: "T"
+ }
+ output_arg {
+ name: "z"
+ type_attr: "T"
+ }
+ attr {
+ name: "T"
+ type: "type"
+ allowed_values {
+ list {
+ type: DT_HALF
+ type: DT_FLOAT
+ type: DT_DOUBLE
+ type: DT_COMPLEX64
+ type: DT_COMPLEX128
+ }
+ }
+ }
+}
+op {
name: "ZerosLike"
input_arg {
name: "x"
diff --git a/tensorflow/core/ops/dataset_ops.cc b/tensorflow/core/ops/dataset_ops.cc
index 1ada623..71f4cc3 100644
--- a/tensorflow/core/ops/dataset_ops.cc
+++ b/tensorflow/core/ops/dataset_ops.cc
@@ -756,6 +756,19 @@
.Attr("output_shapes: list(shape) >= 1")
.SetShapeFn(IteratorGetNextShapeFn);
+REGISTER_OP("ReduceDataset")
+ .Input("input_dataset: variant")
+ .Input("initial_state: Tstate")
+ .Input("other_arguments: Targuments")
+ .Output("components: output_types")
+ .Attr("f: func")
+ .Attr("Tstate: list(type) >= 1")
+ .Attr("Targuments: list(type) >= 0")
+ .Attr("output_types: list(type) >= 1")
+ .Attr("output_shapes: list(shape) >= 1")
+ .Attr("use_inter_op_parallelism: bool = true")
+ .SetShapeFn(IteratorGetNextShapeFn);
+
REGISTER_OP("IteratorToStringHandle")
.Input("resource_handle: resource")
.Output("string_handle: string")
diff --git a/tensorflow/contrib/data/ops/dataset_ops.cc b/tensorflow/core/ops/experimental_dataset_ops.cc
similarity index 62%
rename from tensorflow/contrib/data/ops/dataset_ops.cc
rename to tensorflow/core/ops/experimental_dataset_ops.cc
index d1a771f..f6bd5dc 100644
--- a/tensorflow/contrib/data/ops/dataset_ops.cc
+++ b/tensorflow/core/ops/experimental_dataset_ops.cc
@@ -17,24 +17,16 @@
namespace tensorflow {
-REGISTER_OP("DirectedInterleaveDataset")
+REGISTER_OP("ExperimentalDirectedInterleaveDataset")
.Input("selector_input_dataset: variant")
.Input("data_input_datasets: N * variant")
.Output("handle: variant")
.Attr("output_types: list(type) >= 1")
.Attr("output_shapes: list(shape) >= 1")
.Attr("N: int >= 1")
- .SetShapeFn(shape_inference::ScalarShape)
- .Doc(R"doc(
-A substitute for `InterleaveDataset` on a fixed list of `N` datasets.
+ .SetShapeFn(shape_inference::ScalarShape);
-selector_input_dataset: A dataset of scalar `DT_INT64` elements that determines
- which of the `N` data inputs should produce the next output element.
-data_input_datasets: `N` datasets with the same type that will be interleaved
- according to the values of `selector_input_dataset`.
-)doc");
-
-REGISTER_OP("CSVDataset")
+REGISTER_OP("ExperimentalCSVDataset")
.Input("filenames: string")
.Input("compression_type: string")
.Input("buffer_size: int64")
@@ -76,35 +68,26 @@
return shape_inference::ScalarShape(c);
});
-REGISTER_OP("IgnoreErrorsDataset")
+REGISTER_OP("ExperimentalIgnoreErrorsDataset")
.Input("input_dataset: variant")
.Output("handle: variant")
.Attr("output_types: list(type) >= 1")
.Attr("output_shapes: list(shape) >= 1")
- .SetShapeFn(shape_inference::ScalarShape)
- .Doc(R"doc(
-Creates a dataset that contains the elements of `input_dataset` ignoring errors.
-)doc");
+ .SetShapeFn(shape_inference::ScalarShape);
-REGISTER_OP("UniqueDataset")
+REGISTER_OP("ExperimentalUniqueDataset")
.Input("input_dataset: variant")
.Output("handle: variant")
.Attr("output_types: list(type) >= 1")
.Attr("output_shapes: list(shape) >= 1")
- .SetShapeFn(shape_inference::ScalarShape)
- .Doc(R"doc(
-Creates a dataset that contains the unique elements of `input_dataset`.
-)doc");
+ .SetShapeFn(shape_inference::ScalarShape);
-REGISTER_OP("IteratorGetDevice")
+REGISTER_OP("ExperimentalIteratorGetDevice")
.Input("resource: resource")
.Output("device: string")
- .SetShapeFn(shape_inference::ScalarShape)
- .Doc(R"doc(
-Returns the name of the device on which `resource` has been placed.
-)doc");
+ .SetShapeFn(shape_inference::ScalarShape);
-REGISTER_OP("FunctionBufferingResource")
+REGISTER_OP("ExperimentalFunctionBufferingResource")
.Input("string_arg: string")
.Input("target_device: string")
.Output("resource: resource")
@@ -113,77 +96,36 @@
.Attr("f: func")
.Attr("buffer_size: int")
.Attr("output_types: list(type)")
- .SetShapeFn(shape_inference::UnknownShape)
- .Doc(R"doc(
-Creates a resource that fills up a buffer by making function calls.
+ .SetShapeFn(shape_inference::UnknownShape);
-string_arg: String argument to the function call.
-target_device: Target device to execute the function on.
-resource: Handle to the resource created.
-f: Function to be executed.
-buffer_size: Size of the buffer.
-container: If non-empty, this resource is placed in the given container.
- Otherwise, a default container is used.
-shared_name: If non-empty, this resource will be shared under the given name
- across multiple sessions.
-output_types: The type list for the return values.
-)doc");
-
-REGISTER_OP("FunctionBufferingResourceGetNext")
+REGISTER_OP("ExperimentalFunctionBufferingResourceGetNext")
.Input("function_buffer_resource: resource")
.Attr("output_types: list(type)")
.Output("output: output_types")
- .SetShapeFn(shape_inference::UnknownShape)
- .Doc(R"doc(
-Gets the next element from a FunctionBufferingResource.
+ .SetShapeFn(shape_inference::UnknownShape);
-function_buffer_resource: The FunctionBufferingResource handle.
-output: A list of return values.
-output_types: The type list for the return values.
-)doc");
-
-REGISTER_OP("FunctionBufferingResourceReset")
+REGISTER_OP("ExperimentalFunctionBufferingResourceReset")
.Input("function_buffer_resource: resource")
- .SetShapeFn(shape_inference::UnknownShape)
- .Doc(R"doc(
-Resets the FunctionBufferingResource.
+ .SetShapeFn(shape_inference::UnknownShape);
-function_buffer_resource: The FunctionBufferingResource handle.
-)doc");
-
-REGISTER_OP("ThreadPoolDataset")
+REGISTER_OP("ExperimentalThreadPoolDataset")
.Input("input_dataset: variant")
.Input("thread_pool: resource")
.Output("handle: variant")
.Attr("output_types: list(type) >= 1")
.Attr("output_shapes: list(shape) >= 1")
- .SetShapeFn(shape_inference::ScalarShape)
- .Doc(R"doc(
-Creates a dataset that uses a custom thread pool to compute `input_dataset`.
+ .SetShapeFn(shape_inference::ScalarShape);
-handle: A resource produced by the ThreadPoolHandle op.
-)doc");
-
-REGISTER_OP("ThreadPoolHandle")
+REGISTER_OP("ExperimentalThreadPoolHandle")
.Output("handle: resource")
.SetShapeFn(shape_inference::ScalarShape)
.Attr("num_threads: int")
.Attr("max_intra_op_parallelism: int = 1")
.Attr("display_name: string")
.Attr("container: string = ''")
- .Attr("shared_name: string = ''")
- .Doc(R"doc(
-Creates a custom thread pool with the given number of threads.
+ .Attr("shared_name: string = ''");
-handle: A resource that can be consumed by one or more ThreadPoolDataset ops.
-num_threads: The number of threads in the thread pool.
-max_intra_op_parallelism: The maximum degree of parallelism to use within
- operations that execute on this threadpool.
-display_name: A human-readable name for the threads that may be visible in
- some visualizations.
-)doc");
-
-REGISTER_OP("AssertNextDataset")
+REGISTER_OP("ExperimentalAssertNextDataset")
.Input("input_dataset: variant")
.Input("transformations: string")
.Output("handle: variant")
@@ -196,7 +138,7 @@
return shape_inference::ScalarShape(c);
});
-REGISTER_OP("LMDBDataset")
+REGISTER_OP("ExperimentalLMDBDataset")
.Input("filenames: string")
.Output("handle: variant")
.Attr("output_types: list(type) >= 1")
@@ -205,4 +147,61 @@
// stateful to inhibit constant folding.
.SetShapeFn(shape_inference::ScalarShape);
+REGISTER_OP("ExperimentalIdentityIndexedDataset")
+ .Input("size: uint64")
+ .Output("handle: variant")
+ .SetIsStateful()
+ .SetShapeFn(
+ shape_inference::ScalarShape); // TODO(saeta): check input shapes.
+
+///////////////////////////////////////////////////////////////////////////////
+// IndexedDataset Internals
+///////////////////////////////////////////////////////////////////////////////
+
+// Creates the handle.
+REGISTER_OP("ExperimentalMaterializedIndexDatasetHandle")
+ .Output("handle: resource")
+ .Attr("container: string")
+ .Attr("shared_name: string")
+ .Attr("output_types: list(type) >= 1")
+ .Attr("output_shapes: list(shape) >= 1")
+ .SetShapeFn(shape_inference::ScalarShape);
+
+// Actually materialize the materialize handle.
+REGISTER_OP("ExperimentalIndexedDatasetMaterialize")
+ .Input("dataset: variant")
+ .Input("materialized: resource")
+ .SetShapeFn(shape_inference::NoOutputs);
+
+namespace {
+
+Status GetShapeFn(shape_inference::InferenceContext* c) {
+ shape_inference::ShapeHandle unused;
+ TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &unused));
+ std::vector<PartialTensorShape> output_shapes;
+ TF_RETURN_IF_ERROR(c->GetAttr("output_shapes", &output_shapes));
+ if (output_shapes.size() != c->num_outputs()) {
+ return errors::InvalidArgument(
+ "`output_shapes` must be the same length as `output_types` (",
+ output_shapes.size(), " vs. ", c->num_outputs());
+ }
+ for (size_t i = 0; i < output_shapes.size(); ++i) {
+ shape_inference::ShapeHandle output_shape_handle;
+ TF_RETURN_IF_ERROR(c->MakeShapeFromPartialTensorShape(
+ output_shapes[i], &output_shape_handle));
+ c->set_output(static_cast<int>(i), output_shape_handle);
+ }
+ return Status::OK();
+}
+
+} // namespace
+
+REGISTER_OP("ExperimentalIndexedDatasetGet")
+ .Input("materialized: resource")
+ .Input("index: uint64")
+ .Output("components: output_types")
+ .Attr("output_types: list(type) >= 1")
+ .Attr("output_shapes: list(shape) >= 1")
+ .SetShapeFn(GetShapeFn);
+
} // namespace tensorflow
diff --git a/tensorflow/core/ops/math_grad.cc b/tensorflow/core/ops/math_grad.cc
index 07f876c..55dcc50 100644
--- a/tensorflow/core/ops/math_grad.cc
+++ b/tensorflow/core/ops/math_grad.cc
@@ -549,6 +549,40 @@
}
REGISTER_OP_GRADIENT("Pow", PowGrad);
+Status XlogyGrad(const AttrSlice& attrs, FunctionDef* g) {
+ // clang-format off
+ return GradForBinaryCwise(g, {
+ {{"zeros"}, "ZerosLike", {"x"}},
+ {{"is_x_zero"}, "NotEqual", {"x", "zeros"}},
+ {{"is_zero_cast"}, "Cast", {"is_x_zero"},
+ {{"SrcT", DT_BOOL}, {"DstT", "$T"}}},
+ {{"safe_logy"}, "Xlogy", {"is_zero_cast", "y"}},
+ {{"xlogygrad"}, "Xdivy", {"x", "y"}},
+ {{"gx"}, "Mul", {"safe_logy", "dz"}},
+ {{"gy"}, "Mul", {"xlogygrad", "dz"}},
+ });
+ // clang-format on
+}
+REGISTER_OP_GRADIENT("Xlogy", XlogyGrad);
+
+Status XdivyGrad(const AttrSlice& attrs, FunctionDef* g) {
+ // clang-format off
+ return GradForBinaryCwise(g, {
+ {{"zeros"}, "ZerosLike", {"x"}},
+ {{"is_x_zero"}, "NotEqual", {"x", "zeros"}},
+ {{"is_zero_cast"}, "Cast", {"is_x_zero"},
+ {{"SrcT", DT_BOOL}, {"DstT", "$T"}}},
+ {{"safe_divy"}, "Xdivy", {"is_zero_cast", "y"}},
+ {{"y2"}, "Square", {"y"}},
+ {{"negy2"}, "Neg", {"y2"}},
+ {{"xdivygrad"}, "Xdivy", {"x", "negy2"}},
+ {{"gx"}, "Mul", {"safe_divy", "dz"}},
+ {{"gy"}, "Mul", {"xdivygrad", "dz"}},
+ });
+ // clang-format on
+}
+REGISTER_OP_GRADIENT("Xdivy", XdivyGrad);
+
Status MaximumMinimumGradHelper(const string& comparator,
const AttrSlice& attrs, FunctionDef* g) {
// clang-format off
diff --git a/tensorflow/core/ops/math_grad_test.cc b/tensorflow/core/ops/math_grad_test.cc
index 5ee7980..9fc6b34 100644
--- a/tensorflow/core/ops/math_grad_test.cc
+++ b/tensorflow/core/ops/math_grad_test.cc
@@ -909,6 +909,46 @@
}
#endif // TENSORFLOW_USE_SYCL
+TEST_F(MathGradTest, Xlogy) {
+ auto x = test::AsTensor<float>({0.f, 0.f, 2.f, 3.f, 4.f, 5.f},
+ TensorShape({2, 3}));
+ auto y = test::AsTensor<float>({.5f, 2.f}, TensorShape({2, 1}));
+ Tensor dx;
+ Tensor dy;
+ auto g = [](float x, float y) -> float { return x == 0. ? 0. : std::log(y); };
+ auto h = [](float x, float y) -> float { return x == 0. ? 0. : x / y; };
+ SymGrad("Xlogy", x, y, &dx, &dy);
+ test::ExpectClose(
+ dx, test::AsTensor<float>({g(0.f, .5f), g(0.f, 0.f), g(2.f, .5f),
+ g(3.f, 2.f), g(4.f, 2.f), g(5.f, 2.f)},
+ TensorShape({2, 3})));
+ test::ExpectClose(
+ dy, test::AsTensor<float>({h(0.f, .5f) + h(0.f, 0.f) + h(2.f, .5f),
+ h(3.f, 2.f) + h(4.f, 2.f) + h(5.f, 2.f)},
+ TensorShape({2, 1})));
+}
+
+TEST_F(MathGradTest, Xdivy) {
+ auto x = test::AsTensor<float>({0.f, 0.f, 2.f, 3.f, 4.f, 5.f},
+ TensorShape({2, 3}));
+ auto y = test::AsTensor<float>({.5f, 2.f}, TensorShape({2, 1}));
+ Tensor dx;
+ Tensor dy;
+ auto g = [](float x, float y) -> float { return x == 0. ? 0. : 1 / y; };
+ auto h = [](float x, float y) -> float {
+ return x == 0. ? 0. : -x / (y * y);
+ };
+ SymGrad("Xdivy", x, y, &dx, &dy);
+ test::ExpectClose(
+ dx, test::AsTensor<float>({g(0.f, .5f), g(0.f, 0.f), g(2.f, .5f),
+ g(3.f, 2.f), g(4.f, 2.f), g(5.f, 2.f)},
+ TensorShape({2, 3})));
+ test::ExpectClose(
+ dy, test::AsTensor<float>({h(0.f, .5f) + h(0.f, 0.f) + h(2.f, .5f),
+ h(3.f, 2.f) + h(4.f, 2.f) + h(5.f, 2.f)},
+ TensorShape({2, 1})));
+}
+
TEST_F(MathGradTest, Maximum) {
auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
TensorShape({2, 3}));
diff --git a/tensorflow/core/ops/math_ops.cc b/tensorflow/core/ops/math_ops.cc
index 717263a..3eff728 100644
--- a/tensorflow/core/ops/math_ops.cc
+++ b/tensorflow/core/ops/math_ops.cc
@@ -429,6 +429,20 @@
[here](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html)
)doc");
+REGISTER_OP("Xlogy")
+ .Input("x: T")
+ .Input("y: T")
+ .Output("z: T")
+ .Attr("T: {half, float, double, complex64, complex128}")
+ .SetShapeFn(shape_inference::BroadcastBinaryOpShapeFn);
+
+REGISTER_OP("Xdivy")
+ .Input("x: T")
+ .Input("y: T")
+ .Output("z: T")
+ .Attr("T: {half, float, double, complex64, complex128}")
+ .SetShapeFn(shape_inference::BroadcastBinaryOpShapeFn);
+
#undef BINARY_FEWER
#undef BINARY_MORE
diff --git a/tensorflow/core/ops/nn_ops.cc b/tensorflow/core/ops/nn_ops.cc
index dc39996..d1d81b2 100644
--- a/tensorflow/core/ops/nn_ops.cc
+++ b/tensorflow/core/ops/nn_ops.cc
@@ -178,7 +178,7 @@
.Output("reserve_space_2: T")
.Attr("T: {float}")
.Attr("epsilon: float = 0.0001")
- .Attr("data_format: string = 'NHWC'")
+ .Attr(GetConvnetDataFormatAttrString())
.Attr("is_training: bool = true")
.SetShapeFn(shape_inference::FusedBatchNormShape);
@@ -196,7 +196,7 @@
.Attr("T: {half, bfloat16, float}")
.Attr("U: {float}")
.Attr("epsilon: float = 0.0001")
- .Attr("data_format: string = 'NHWC'")
+ .Attr(GetConvnetDataFormatAttrString())
.Attr("is_training: bool = true")
.SetShapeFn(shape_inference::FusedBatchNormShape);
@@ -213,7 +213,7 @@
.Output("reserve_space_4: T")
.Attr("T: {float}")
.Attr("epsilon: float = 0.0001")
- .Attr("data_format: string = 'NHWC'")
+ .Attr(GetConvnetDataFormatAttrString())
.Attr("is_training: bool = true")
.SetShapeFn(shape_inference::FusedBatchNormGradShape);
@@ -231,7 +231,7 @@
.Attr("T: {half, bfloat16, float}")
.Attr("U: {float}")
.Attr("epsilon: float = 0.0001")
- .Attr("data_format: string = 'NHWC'")
+ .Attr(GetConvnetDataFormatAttrString())
.Attr("is_training: bool = true")
.SetShapeFn(shape_inference::FusedBatchNormGradShape);
diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index bdded2d..abee803 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -10039,6 +10039,421 @@
}
}
op {
+ name: "ExperimentalAssertNextDataset"
+ input_arg {
+ name: "input_dataset"
+ type: DT_VARIANT
+ }
+ input_arg {
+ name: "transformations"
+ type: DT_STRING
+ }
+ output_arg {
+ name: "handle"
+ type: DT_VARIANT
+ }
+ attr {
+ name: "output_types"
+ type: "list(type)"
+ has_minimum: true
+ minimum: 1
+ }
+ attr {
+ name: "output_shapes"
+ type: "list(shape)"
+ has_minimum: true
+ minimum: 1
+ }
+}
+op {
+ name: "ExperimentalCSVDataset"
+ input_arg {
+ name: "filenames"
+ type: DT_STRING
+ }
+ input_arg {
+ name: "compression_type"
+ type: DT_STRING
+ }
+ input_arg {
+ name: "buffer_size"
+ type: DT_INT64
+ }
+ input_arg {
+ name: "header"
+ type: DT_BOOL
+ }
+ input_arg {
+ name: "field_delim"
+ type: DT_STRING
+ }
+ input_arg {
+ name: "use_quote_delim"
+ type: DT_BOOL
+ }
+ input_arg {
+ name: "na_value"
+ type: DT_STRING
+ }
+ input_arg {
+ name: "select_cols"
+ type: DT_INT64
+ }
+ input_arg {
+ name: "record_defaults"
+ type_list_attr: "output_types"
+ }
+ output_arg {
+ name: "handle"
+ type: DT_VARIANT
+ }
+ attr {
+ name: "output_types"
+ type: "list(type)"
+ has_minimum: true
+ minimum: 1
+ allowed_values {
+ list {
+ type: DT_FLOAT
+ type: DT_DOUBLE
+ type: DT_INT32
+ type: DT_INT64
+ type: DT_STRING
+ }
+ }
+ }
+ attr {
+ name: "output_shapes"
+ type: "list(shape)"
+ has_minimum: true
+ minimum: 1
+ }
+ is_stateful: true
+}
+op {
+ name: "ExperimentalDirectedInterleaveDataset"
+ input_arg {
+ name: "selector_input_dataset"
+ type: DT_VARIANT
+ }
+ input_arg {
+ name: "data_input_datasets"
+ type: DT_VARIANT
+ number_attr: "N"
+ }
+ output_arg {
+ name: "handle"
+ type: DT_VARIANT
+ }
+ 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: "N"
+ type: "int"
+ has_minimum: true
+ minimum: 1
+ }
+}
+op {
+ name: "ExperimentalFunctionBufferingResource"
+ input_arg {
+ name: "string_arg"
+ type: DT_STRING
+ }
+ input_arg {
+ name: "target_device"
+ type: DT_STRING
+ }
+ output_arg {
+ name: "resource"
+ type: DT_RESOURCE
+ }
+ attr {
+ name: "shared_name"
+ type: "string"
+ }
+ attr {
+ name: "container"
+ type: "string"
+ }
+ attr {
+ name: "f"
+ type: "func"
+ }
+ attr {
+ name: "buffer_size"
+ type: "int"
+ }
+ attr {
+ name: "output_types"
+ type: "list(type)"
+ }
+ is_stateful: true
+}
+op {
+ name: "ExperimentalFunctionBufferingResourceGetNext"
+ input_arg {
+ name: "function_buffer_resource"
+ type: DT_RESOURCE
+ }
+ output_arg {
+ name: "output"
+ type_list_attr: "output_types"
+ }
+ attr {
+ name: "output_types"
+ type: "list(type)"
+ has_minimum: true
+ minimum: 1
+ }
+ is_stateful: true
+}
+op {
+ name: "ExperimentalFunctionBufferingResourceReset"
+ input_arg {
+ name: "function_buffer_resource"
+ type: DT_RESOURCE
+ }
+ is_stateful: true
+}
+op {
+ name: "ExperimentalIdentityIndexedDataset"
+ input_arg {
+ name: "size"
+ type: DT_UINT64
+ }
+ output_arg {
+ name: "handle"
+ type: DT_VARIANT
+ }
+ is_stateful: true
+}
+op {
+ name: "ExperimentalIgnoreErrorsDataset"
+ input_arg {
+ name: "input_dataset"
+ type: DT_VARIANT
+ }
+ output_arg {
+ name: "handle"
+ type: DT_VARIANT
+ }
+ attr {
+ name: "output_types"
+ type: "list(type)"
+ has_minimum: true
+ minimum: 1
+ }
+ attr {
+ name: "output_shapes"
+ type: "list(shape)"
+ has_minimum: true
+ minimum: 1
+ }
+}
+op {
+ name: "ExperimentalIndexedDatasetGet"
+ input_arg {
+ name: "materialized"
+ type: DT_RESOURCE
+ }
+ input_arg {
+ name: "index"
+ type: DT_UINT64
+ }
+ output_arg {
+ name: "components"
+ type_list_attr: "output_types"
+ }
+ attr {
+ name: "output_types"
+ type: "list(type)"
+ has_minimum: true
+ minimum: 1
+ }
+ attr {
+ name: "output_shapes"
+ type: "list(shape)"
+ has_minimum: true
+ minimum: 1
+ }
+ is_stateful: true
+}
+op {
+ name: "ExperimentalIndexedDatasetMaterialize"
+ input_arg {
+ name: "dataset"
+ type: DT_VARIANT
+ }
+ input_arg {
+ name: "materialized"
+ type: DT_RESOURCE
+ }
+ is_stateful: true
+}
+op {
+ name: "ExperimentalIteratorGetDevice"
+ input_arg {
+ name: "resource"
+ type: DT_RESOURCE
+ }
+ output_arg {
+ name: "device"
+ type: DT_STRING
+ }
+ is_stateful: true
+}
+op {
+ name: "ExperimentalLMDBDataset"
+ input_arg {
+ name: "filenames"
+ type: DT_STRING
+ }
+ output_arg {
+ name: "handle"
+ type: DT_VARIANT
+ }
+ attr {
+ name: "output_types"
+ type: "list(type)"
+ has_minimum: true
+ minimum: 1
+ }
+ attr {
+ name: "output_shapes"
+ type: "list(shape)"
+ has_minimum: true
+ minimum: 1
+ }
+ is_stateful: true
+}
+op {
+ name: "ExperimentalMaterializedIndexDatasetHandle"
+ output_arg {
+ name: "handle"
+ type: DT_RESOURCE
+ }
+ attr {
+ name: "container"
+ type: "string"
+ }
+ attr {
+ name: "shared_name"
+ type: "string"
+ }
+ attr {
+ name: "output_types"
+ type: "list(type)"
+ has_minimum: true
+ minimum: 1
+ }
+ attr {
+ name: "output_shapes"
+ type: "list(shape)"
+ has_minimum: true
+ minimum: 1
+ }
+ is_stateful: true
+}
+op {
+ name: "ExperimentalThreadPoolDataset"
+ input_arg {
+ name: "input_dataset"
+ type: DT_VARIANT
+ }
+ input_arg {
+ name: "thread_pool"
+ type: DT_RESOURCE
+ }
+ output_arg {
+ name: "handle"
+ type: DT_VARIANT
+ }
+ attr {
+ name: "output_types"
+ type: "list(type)"
+ has_minimum: true
+ minimum: 1
+ }
+ attr {
+ name: "output_shapes"
+ type: "list(shape)"
+ has_minimum: true
+ minimum: 1
+ }
+ is_stateful: true
+}
+op {
+ name: "ExperimentalThreadPoolHandle"
+ output_arg {
+ name: "handle"
+ type: DT_RESOURCE
+ }
+ attr {
+ name: "num_threads"
+ type: "int"
+ }
+ attr {
+ name: "max_intra_op_parallelism"
+ type: "int"
+ default_value {
+ i: 1
+ }
+ }
+ attr {
+ name: "display_name"
+ type: "string"
+ }
+ attr {
+ name: "container"
+ type: "string"
+ default_value {
+ s: ""
+ }
+ }
+ attr {
+ name: "shared_name"
+ type: "string"
+ default_value {
+ s: ""
+ }
+ }
+ is_stateful: true
+}
+op {
+ name: "ExperimentalUniqueDataset"
+ input_arg {
+ name: "input_dataset"
+ type: DT_VARIANT
+ }
+ output_arg {
+ name: "handle"
+ type: DT_VARIANT
+ }
+ attr {
+ name: "output_types"
+ type: "list(type)"
+ has_minimum: true
+ minimum: 1
+ }
+ attr {
+ name: "output_shapes"
+ type: "list(shape)"
+ has_minimum: true
+ minimum: 1
+ }
+}
+op {
name: "Expm1"
input_arg {
name: "x"
@@ -11459,6 +11874,12 @@
default_value {
s: "NHWC"
}
+ allowed_values {
+ list {
+ s: "NHWC"
+ s: "NCHW"
+ }
+ }
}
attr {
name: "is_training"
@@ -11532,6 +11953,12 @@
default_value {
s: "NHWC"
}
+ allowed_values {
+ list {
+ s: "NHWC"
+ s: "NCHW"
+ }
+ }
}
attr {
name: "is_training"
@@ -11616,6 +12043,12 @@
default_value {
s: "NHWC"
}
+ allowed_values {
+ list {
+ s: "NHWC"
+ s: "NCHW"
+ }
+ }
}
attr {
name: "is_training"
@@ -11700,6 +12133,12 @@
default_value {
s: "NHWC"
}
+ allowed_values {
+ list {
+ s: "NHWC"
+ s: "NCHW"
+ }
+ }
}
attr {
name: "is_training"
@@ -22845,6 +23284,59 @@
is_stateful: true
}
op {
+ name: "ReduceDataset"
+ 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: "components"
+ type_list_attr: "output_types"
+ }
+ 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: "use_inter_op_parallelism"
+ type: "bool"
+ default_value {
+ b: true
+ }
+ }
+}
+op {
name: "ReduceJoin"
input_arg {
name: "inputs"
@@ -35625,6 +36117,17 @@
}
}
op {
+ name: "UnicodeScript"
+ input_arg {
+ name: "input"
+ type: DT_INT32
+ }
+ output_arg {
+ name: "output"
+ type: DT_INT32
+ }
+}
+op {
name: "UniformCandidateSampler"
input_arg {
name: "true_classes"
@@ -36737,6 +37240,62 @@
is_stateful: true
}
op {
+ name: "Xdivy"
+ input_arg {
+ name: "x"
+ type_attr: "T"
+ }
+ input_arg {
+ name: "y"
+ type_attr: "T"
+ }
+ output_arg {
+ name: "z"
+ type_attr: "T"
+ }
+ attr {
+ name: "T"
+ type: "type"
+ allowed_values {
+ list {
+ type: DT_HALF
+ type: DT_FLOAT
+ type: DT_DOUBLE
+ type: DT_COMPLEX64
+ type: DT_COMPLEX128
+ }
+ }
+ }
+}
+op {
+ name: "Xlogy"
+ input_arg {
+ name: "x"
+ type_attr: "T"
+ }
+ input_arg {
+ name: "y"
+ type_attr: "T"
+ }
+ output_arg {
+ name: "z"
+ type_attr: "T"
+ }
+ attr {
+ name: "T"
+ type: "type"
+ allowed_values {
+ list {
+ type: DT_HALF
+ type: DT_FLOAT
+ type: DT_DOUBLE
+ type: DT_COMPLEX64
+ type: DT_COMPLEX128
+ }
+ }
+ }
+}
+op {
name: "ZerosLike"
input_arg {
name: "x"
diff --git a/tensorflow/core/ops/resource_variable_ops.cc b/tensorflow/core/ops/resource_variable_ops.cc
index 2649954..adc9cd1 100644
--- a/tensorflow/core/ops/resource_variable_ops.cc
+++ b/tensorflow/core/ops/resource_variable_ops.cc
@@ -19,6 +19,7 @@
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/framework/shape_inference.h"
+#include "tensorflow/core/lib/core/errors.h"
using ::tensorflow::shape_inference::InferenceContext;
using ::tensorflow::shape_inference::ShapeAndType;
@@ -56,6 +57,36 @@
return Status::OK();
}
+Status ReadVariablesShapeFn(InferenceContext* c) {
+ int n;
+ TF_RETURN_IF_ERROR(c->GetAttr("N", &n));
+ DataTypeVector value_dtypes;
+ TF_RETURN_IF_ERROR(c->GetAttr("dtypes", &value_dtypes));
+ if (n != value_dtypes.size()) {
+ return errors::InvalidArgument(
+ "Mismatched number of arguments to ReadVariablesOp");
+ }
+ for (int i = 0; i < n; ++i) {
+ ShapeAndType shape_and_type;
+ auto* handle_data = c->input_handle_shapes_and_types(i);
+ if (handle_data == nullptr || handle_data->empty()) {
+ shape_and_type.shape = c->UnknownShape();
+ shape_and_type.dtype = DT_INVALID;
+ } else {
+ shape_and_type = (*handle_data)[0];
+ if (shape_and_type.dtype != value_dtypes[i]) {
+ return errors::InvalidArgument(
+ "Trying to read variable with wrong dtype. "
+ "Expected ",
+ DataTypeString(shape_and_type.dtype), " got ",
+ DataTypeString(value_dtypes[i]));
+ }
+ }
+ c->set_output(i, shape_and_type.shape);
+ }
+ return Status::OK();
+}
+
} // namespace
REGISTER_OP("VarHandleOp")
@@ -79,12 +110,53 @@
return Status::OK();
});
+REGISTER_OP("_VarHandlesOp")
+ .Attr("containers: list(string)")
+ .Attr("shared_names: list(string)")
+ .Attr("N: int >= 0")
+ .Attr("dtypes: list(type)")
+ .Attr("shapes: list(shape)")
+ .Output("resources: N * resource")
+ .SetIsStateful()
+ .SetShapeFn([](InferenceContext* c) {
+ int n;
+ TF_RETURN_IF_ERROR(c->GetAttr("N", &n));
+ DataTypeVector dtypes;
+ TF_RETURN_IF_ERROR(c->GetAttr("dtypes", &dtypes));
+ std::vector<PartialTensorShape> shapes;
+ TF_RETURN_IF_ERROR(c->GetAttr("shapes", &shapes));
+ if (dtypes.size() != n) {
+ return errors::InvalidArgument("Mismatched number of dtypes (n=", n,
+ ", num dtypes=", dtypes.size(), ")");
+ }
+ if (shapes.size() != n) {
+ return errors::InvalidArgument("Mismatched number of shapes (n=", n,
+ ", num shapes=", shapes.size(), ")");
+ }
+ for (int i = 0; i < n; ++i) {
+ c->set_output(i, c->Scalar());
+ ShapeHandle s;
+ TF_RETURN_IF_ERROR(c->MakeShapeFromPartialTensorShape(shapes[i], &s));
+ c->set_output_handle_shapes_and_types(
+ i, std::vector<ShapeAndType>{{s, dtypes[i]}});
+ }
+
+ return Status::OK();
+ });
+
REGISTER_OP("ReadVariableOp")
.Input("resource: resource")
.Output("value: dtype")
.Attr("dtype: type")
.SetShapeFn(ReadVariableShapeFn);
+REGISTER_OP("_ReadVariablesOp")
+ .Attr("N: int >= 0")
+ .Input("resources: N * resource")
+ .Output("values: dtypes")
+ .Attr("dtypes: list(type)")
+ .SetShapeFn(ReadVariablesShapeFn);
+
Status ReadGrad(const AttrSlice& attrs, FunctionDef* g) {
// clang-format off
*g = FunctionDefHelper::Define(
diff --git a/tensorflow/core/ops/string_ops.cc b/tensorflow/core/ops/string_ops.cc
index da1d2a6..b4fbde5 100644
--- a/tensorflow/core/ops/string_ops.cc
+++ b/tensorflow/core/ops/string_ops.cc
@@ -244,4 +244,9 @@
return shape_inference::BroadcastBinaryOpShapeFn(c);
});
+REGISTER_OP("UnicodeScript")
+ .Input("input: int32")
+ .Output("output: int32")
+ .SetShapeFn(shape_inference::UnchangedShape);
+
} // namespace tensorflow
diff --git a/tensorflow/core/platform/cloud/compute_engine_metadata_client.cc b/tensorflow/core/platform/cloud/compute_engine_metadata_client.cc
index f41b83a..affb68e 100644
--- a/tensorflow/core/platform/cloud/compute_engine_metadata_client.cc
+++ b/tensorflow/core/platform/cloud/compute_engine_metadata_client.cc
@@ -17,7 +17,6 @@
#include <utility>
#include "tensorflow/core/platform/cloud/curl_http_request.h"
-#include "tensorflow/core/platform/cloud/retrying_utils.h"
namespace tensorflow {
@@ -25,21 +24,14 @@
// The URL to retrieve metadata when running in Google Compute Engine.
constexpr char kGceMetadataBaseUrl[] = "http://metadata/computeMetadata/v1/";
-// The default initial delay between retries with exponential backoff.
-constexpr int kInitialRetryDelayUsec = 500000; // 0.5 sec
} // namespace
ComputeEngineMetadataClient::ComputeEngineMetadataClient(
- std::shared_ptr<HttpRequest::Factory> http_request_factory)
- : ComputeEngineMetadataClient(std::move(http_request_factory),
- kInitialRetryDelayUsec) {}
-
-ComputeEngineMetadataClient::ComputeEngineMetadataClient(
std::shared_ptr<HttpRequest::Factory> http_request_factory,
- int64 initial_retry_delay_usec)
+ const RetryConfig& config)
: http_request_factory_(std::move(http_request_factory)),
- initial_retry_delay_usec_(initial_retry_delay_usec) {}
+ retry_config_(config) {}
Status ComputeEngineMetadataClient::GetMetadata(
const string& path, std::vector<char>* response_buffer) {
@@ -52,8 +44,7 @@
return Status::OK();
};
- return RetryingUtils::CallWithRetries(get_metadata_from_gce,
- initial_retry_delay_usec_);
+ return RetryingUtils::CallWithRetries(get_metadata_from_gce, retry_config_);
}
} // namespace tensorflow
diff --git a/tensorflow/core/platform/cloud/compute_engine_metadata_client.h b/tensorflow/core/platform/cloud/compute_engine_metadata_client.h
index 534ccf3..7f06032 100644
--- a/tensorflow/core/platform/cloud/compute_engine_metadata_client.h
+++ b/tensorflow/core/platform/cloud/compute_engine_metadata_client.h
@@ -18,6 +18,7 @@
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/cloud/http_request.h"
+#include "tensorflow/core/platform/cloud/retrying_utils.h"
namespace tensorflow {
@@ -31,10 +32,11 @@
class ComputeEngineMetadataClient {
public:
explicit ComputeEngineMetadataClient(
- std::shared_ptr<HttpRequest::Factory> http_request_factory);
- ComputeEngineMetadataClient(
std::shared_ptr<HttpRequest::Factory> http_request_factory,
- int64 initial_retry_delay_usec);
+ const RetryConfig& config = RetryConfig(
+ 10000, /* init_delay_time_us = 1 ms */
+ 1000000 /* max_delay_time_us = 1 s */
+ ));
virtual ~ComputeEngineMetadataClient() {}
/// \brief Get the metadata value for a given attribute of the metadata
@@ -54,7 +56,7 @@
private:
std::shared_ptr<HttpRequest::Factory> http_request_factory_;
- const int64 initial_retry_delay_usec_;
+ const RetryConfig retry_config_;
TF_DISALLOW_COPY_AND_ASSIGN(ComputeEngineMetadataClient);
};
diff --git a/tensorflow/core/platform/cloud/compute_engine_metadata_client_test.cc b/tensorflow/core/platform/cloud/compute_engine_metadata_client_test.cc
index 4c41cca..e891b4a 100644
--- a/tensorflow/core/platform/cloud/compute_engine_metadata_client_test.cc
+++ b/tensorflow/core/platform/cloud/compute_engine_metadata_client_test.cc
@@ -30,7 +30,8 @@
std::shared_ptr<HttpRequest::Factory> http_factory =
std::make_shared<FakeHttpRequestFactory>(&requests);
- ComputeEngineMetadataClient client(http_factory, 0);
+ ComputeEngineMetadataClient client(http_factory,
+ RetryConfig(0 /* init_delay_time_us */));
std::vector<char> result;
TF_EXPECT_OK(
@@ -56,7 +57,8 @@
std::shared_ptr<HttpRequest::Factory> http_factory =
std::make_shared<FakeHttpRequestFactory>(&requests);
- ComputeEngineMetadataClient client(http_factory, 0);
+ ComputeEngineMetadataClient client(http_factory,
+ RetryConfig(0 /* init_delay_time_us */));
std::vector<char> result;
TF_EXPECT_OK(
diff --git a/tensorflow/core/platform/cloud/compute_engine_zone_provider_test.cc b/tensorflow/core/platform/cloud/compute_engine_zone_provider_test.cc
index f7477ec..476e4f9 100644
--- a/tensorflow/core/platform/cloud/compute_engine_zone_provider_test.cc
+++ b/tensorflow/core/platform/cloud/compute_engine_zone_provider_test.cc
@@ -34,8 +34,8 @@
auto httpRequestFactory = std::make_shared<FakeHttpRequestFactory>(&requests);
- auto metadata_client =
- std::make_shared<ComputeEngineMetadataClient>(httpRequestFactory, 0);
+ auto metadata_client = std::make_shared<ComputeEngineMetadataClient>(
+ httpRequestFactory, RetryConfig(0 /* init_delay_time_us */));
ComputeEngineZoneProvider provider(metadata_client);
@@ -55,8 +55,8 @@
auto httpRequestFactory = std::make_shared<FakeHttpRequestFactory>(&requests);
- auto metadata_client =
- std::make_shared<ComputeEngineMetadataClient>(httpRequestFactory, 0);
+ auto metadata_client = std::make_shared<ComputeEngineMetadataClient>(
+ httpRequestFactory, RetryConfig(0 /* init_delay_time_us */));
ComputeEngineZoneProvider provider(metadata_client);
diff --git a/tensorflow/core/platform/cloud/gcs_file_system.cc b/tensorflow/core/platform/cloud/gcs_file_system.cc
index 83ea853..c61b68a 100644
--- a/tensorflow/core/platform/cloud/gcs_file_system.cc
+++ b/tensorflow/core/platform/cloud/gcs_file_system.cc
@@ -333,14 +333,14 @@
GcsFileSystem* filesystem,
GcsFileSystem::TimeoutConfig* timeouts,
std::function<void()> file_cache_erase,
- int64 initial_retry_delay_usec)
+ RetryConfig retry_config)
: bucket_(bucket),
object_(object),
filesystem_(filesystem),
timeouts_(timeouts),
file_cache_erase_(std::move(file_cache_erase)),
sync_needed_(true),
- initial_retry_delay_usec_(initial_retry_delay_usec) {
+ retry_config_(retry_config) {
// TODO: to make it safer, outfile_ should be constructed from an FD
if (GetTmpFilename(&tmp_content_filename_).ok()) {
outfile_.open(tmp_content_filename_,
@@ -357,14 +357,14 @@
GcsFileSystem* filesystem, const string& tmp_content_filename,
GcsFileSystem::TimeoutConfig* timeouts,
std::function<void()> file_cache_erase,
- int64 initial_retry_delay_usec)
+ RetryConfig retry_config)
: bucket_(bucket),
object_(object),
filesystem_(filesystem),
timeouts_(timeouts),
file_cache_erase_(std::move(file_cache_erase)),
sync_needed_(true),
- initial_retry_delay_usec_(initial_retry_delay_usec) {
+ retry_config_(retry_config) {
tmp_content_filename_ = tmp_content_filename;
outfile_.open(tmp_content_filename_,
std::ofstream::binary | std::ofstream::app);
@@ -441,7 +441,7 @@
first_attempt = false;
return UploadToSession(session_uri, already_uploaded);
},
- initial_retry_delay_usec_);
+ retry_config_);
if (upload_status.code() == errors::Code::NOT_FOUND) {
// GCS docs recommend retrying the whole upload. We're relying on the
// RetryingFileSystem to retry the Sync() call.
@@ -586,7 +586,7 @@
GcsFileSystem::TimeoutConfig* timeouts_;
std::function<void()> file_cache_erase_;
bool sync_needed_; // whether there is buffered data that needs to be synced
- int64 initial_retry_delay_usec_;
+ RetryConfig retry_config_;
};
class GcsReadOnlyMemoryRegion : public ReadOnlyMemoryRegion {
@@ -791,7 +791,7 @@
std::unique_ptr<ZoneProvider> zone_provider, size_t block_size,
size_t max_bytes, uint64 max_staleness, uint64 stat_cache_max_age,
size_t stat_cache_max_entries, uint64 matching_paths_cache_max_age,
- size_t matching_paths_cache_max_entries, int64 initial_retry_delay_usec,
+ size_t matching_paths_cache_max_entries, RetryConfig retry_config,
TimeoutConfig timeouts, const std::unordered_set<string>& allowed_locations,
std::pair<const string, const string>* additional_header)
: auth_provider_(std::move(auth_provider)),
@@ -806,7 +806,7 @@
kCacheNeverExpire, kBucketLocationCacheMaxEntries)),
allowed_locations_(allowed_locations),
timeouts_(timeouts),
- initial_retry_delay_usec_(initial_retry_delay_usec),
+ retry_config_(retry_config),
additional_header_(additional_header) {}
Status GcsFileSystem::NewRandomAccessFile(
@@ -941,7 +941,7 @@
TF_RETURN_IF_ERROR(ParseGcsPath(fname, false, &bucket, &object));
result->reset(new GcsWritableFile(bucket, object, this, &timeouts_,
[this, fname]() { ClearFileCaches(fname); },
- initial_retry_delay_usec_));
+ retry_config_));
return Status::OK();
}
@@ -981,7 +981,7 @@
TF_RETURN_IF_ERROR(ParseGcsPath(fname, false, &bucket, &object));
result->reset(new GcsWritableFile(
bucket, object, this, old_content_filename, &timeouts_,
- [this, fname]() { ClearFileCaches(fname); }, initial_retry_delay_usec_));
+ [this, fname]() { ClearFileCaches(fname); }, retry_config_));
return Status::OK();
}
@@ -1534,7 +1534,7 @@
// on the server side, we can't just retry the whole RenameFile operation
// because the source object is already gone.
return RetryingUtils::DeleteWithRetries(
- [this, &src]() { return DeleteFile(src); }, initial_retry_delay_usec_);
+ [this, &src]() { return DeleteFile(src); }, retry_config_);
}
Status GcsFileSystem::IsDirectory(const string& fname) {
@@ -1590,8 +1590,7 @@
// and therefore RetryingFileSystem won't pay attention to the failures,
// we need to make sure these failures are properly retried.
const auto& delete_file_status = RetryingUtils::DeleteWithRetries(
- [this, &full_path]() { return DeleteFile(full_path); },
- initial_retry_delay_usec_);
+ [this, &full_path]() { return DeleteFile(full_path); }, retry_config_);
if (!delete_file_status.ok()) {
if (IsDirectory(full_path).ok()) {
// The object is a directory marker.
diff --git a/tensorflow/core/platform/cloud/gcs_file_system.h b/tensorflow/core/platform/cloud/gcs_file_system.h
index 71db707..d0840a3 100644
--- a/tensorflow/core/platform/cloud/gcs_file_system.h
+++ b/tensorflow/core/platform/cloud/gcs_file_system.h
@@ -93,7 +93,7 @@
uint64 stat_cache_max_age, size_t stat_cache_max_entries,
uint64 matching_paths_cache_max_age,
size_t matching_paths_cache_max_entries,
- int64 initial_retry_delay_usec, TimeoutConfig timeouts,
+ RetryConfig retry_config, TimeoutConfig timeouts,
const std::unordered_set<string>& allowed_locations,
std::pair<const string, const string>* additional_header);
@@ -332,7 +332,7 @@
GcsStatsInterface* stats_ = nullptr; // Not owned.
/// The initial delay for exponential backoffs when retrying failed calls.
- const int64 initial_retry_delay_usec_ = 1000000L;
+ RetryConfig retry_config_;
// Additional header material to be transmitted with all GCS requests
std::unique_ptr<std::pair<const string, const string>> additional_header_;
@@ -344,7 +344,8 @@
class RetryingGcsFileSystem : public RetryingFileSystem<GcsFileSystem> {
public:
RetryingGcsFileSystem()
- : RetryingFileSystem(std::unique_ptr<GcsFileSystem>(new GcsFileSystem)) {}
+ : RetryingFileSystem(std::unique_ptr<GcsFileSystem>(new GcsFileSystem),
+ RetryConfig(100000 /* init_delay_time_us */)) {}
};
} // namespace tensorflow
diff --git a/tensorflow/core/platform/cloud/gcs_file_system_test.cc b/tensorflow/core/platform/cloud/gcs_file_system_test.cc
index 14376ad..702802b 100644
--- a/tensorflow/core/platform/cloud/gcs_file_system_test.cc
+++ b/tensorflow/core/platform/cloud/gcs_file_system_test.cc
@@ -24,6 +24,8 @@
namespace {
static GcsFileSystem::TimeoutConfig kTestTimeoutConfig(5, 1, 10, 20, 30);
+static RetryConfig kTestRetryConfig(0 /* init_delay_time_us */);
+
// Default (empty) constraint config
static std::unordered_set<string>* kAllowedLocationsDefault =
new std::unordered_set<string>();
@@ -62,16 +64,16 @@
"Range: 6-11\n"
"Timeouts: 5 1 20\n",
"6789")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
std::unique_ptr<RandomAccessFile> file;
TF_EXPECT_OK(fs.NewRandomAccessFile("gs://bucket/random_access.txt", &file));
@@ -108,9 +110,9 @@
0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
0 /* stat cache max age */, 0 /* stat cache max entries */,
0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */,
- 0 /* initial retry delay */, kTestTimeoutConfig,
- *kAllowedLocationsAuto, nullptr /* gcs additional header */);
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsAuto,
+ nullptr /* gcs additional header */);
std::unique_ptr<RandomAccessFile> file;
TF_EXPECT_OK(fs.NewRandomAccessFile("gs://bucket/random_access.txt", &file));
@@ -150,9 +152,9 @@
0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
0 /* stat cache max age */, 0 /* stat cache max entries */,
0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */,
- 0 /* initial retry delay */, kTestTimeoutConfig,
- *kAllowedLocationsAuto, nullptr /* gcs additional header */);
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsAuto,
+ nullptr /* gcs additional header */);
std::unique_ptr<RandomAccessFile> file;
@@ -191,9 +193,9 @@
0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
0 /* stat cache max age */, 0 /* stat cache max entries */,
0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */,
- 0 /* initial retry delay */, kTestTimeoutConfig,
- *kAllowedLocationsAuto, nullptr /* gcs additional header */);
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsAuto,
+ nullptr /* gcs additional header */);
std::unique_ptr<RandomAccessFile> file;
EXPECT_EQ(tensorflow::errors::FailedPrecondition(
@@ -216,16 +218,16 @@
"Range: 3-12\n"
"Timeouts: 5 1 20\n",
"3456789")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
std::unique_ptr<RandomAccessFile> file;
TF_EXPECT_OK(fs.NewRandomAccessFile("gs://bucket/random_access.txt", &file));
@@ -283,7 +285,7 @@
std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 9 /* block size */,
18 /* max bytes */, 0 /* max staleness */, 3600 /* stat cache max age */,
0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
kTestTimeoutConfig, *kAllowedLocationsDefault,
nullptr /* gcs additional header */);
@@ -372,7 +374,7 @@
std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 9 /* block size */,
18 /* max bytes */, 0 /* max staleness */, 3600 /* stat cache max age */,
0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
kTestTimeoutConfig, *kAllowedLocationsDefault,
nullptr /* gcs additional header */);
@@ -414,17 +416,17 @@
"Range: 8-15\n"
"Timeouts: 5 1 20\n",
"89abcdef")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 8 /* block size */,
- 16 /* max bytes */, 3600 /* max staleness */,
- 3600 /* stat cache max age */, 0 /* stat cache max entries */,
- 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 8 /* block size */, 16 /* max bytes */,
+ 3600 /* max staleness */, 3600 /* stat cache max age */,
+ 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
char scratch[100];
StringPiece result;
// There should only be two HTTP requests issued to GCS even though we iterate
@@ -492,7 +494,7 @@
std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 9 /* block size */,
18 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
kTestTimeoutConfig, *kAllowedLocationsDefault,
nullptr /* gcs additional header */);
@@ -513,17 +515,17 @@
TEST(GcsFileSystemTest, NewRandomAccessFile_NoObjectName) {
std::vector<HttpRequest*> requests;
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
- 0 /* read ahead bytes */, 0 /* max bytes */, 0 /* max staleness */,
- 0 /* stat cache max age */, 0 /* stat cache max entries */,
- 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* read ahead bytes */, 0 /* max bytes */,
+ 0 /* max staleness */, 0 /* stat cache max age */,
+ 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
std::unique_ptr<RandomAccessFile> file;
EXPECT_EQ(errors::Code::INVALID_ARGUMENT,
@@ -547,16 +549,16 @@
"012")});
// Set stat_cache_max_age to 1000s so that StatCache could work.
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 1e3 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 1e3 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
// Stat the file first so that the file stats are cached.
FileStatistics stat;
@@ -621,7 +623,7 @@
std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 8 /* block size */,
8 /* max bytes */, 0 /* max staleness */, 3600 /* stat cache max age */,
0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
kTestTimeoutConfig, *kAllowedLocationsDefault,
nullptr /* gcs additional header */);
@@ -703,16 +705,16 @@
"Timeouts: 5 1 30\n"
"Put body: t2\n",
"")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
std::unique_ptr<WritableFile> file;
TF_EXPECT_OK(fs.NewWritableFile("gs://bucket/path/writeable.txt", &file));
@@ -773,17 +775,17 @@
"Range: 0-7\n"
"Timeouts: 5 1 20\n",
"01234567")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 8 /* block size */,
- 8 /* max bytes */, 3600 /* max staleness */,
- 3600 /* stat cache max age */, 0 /* stat cache max entries */,
- 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 8 /* block size */, 8 /* max bytes */,
+ 3600 /* max staleness */, 3600 /* stat cache max age */,
+ 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
// Pull the file's first block into the cache. This will trigger the first
// HTTP request to GCS.
std::unique_ptr<RandomAccessFile> rfile;
@@ -867,9 +869,9 @@
std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 2 /* initial retry delay */,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ 0 /* matching paths cache max entries */,
+ RetryConfig(2 /* .init_delay_time_us */), kTestTimeoutConfig,
+ *kAllowedLocationsDefault, nullptr /* gcs additional header */);
std::unique_ptr<WritableFile> file;
TF_EXPECT_OK(fs.NewWritableFile("gs://bucket/path/writeable.txt", &file));
@@ -918,16 +920,16 @@
"Timeouts: 5 1 30\n"
"Put body: content1,content2\n",
"")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
std::unique_ptr<WritableFile> file;
TF_EXPECT_OK(fs.NewWritableFile("gs://bucket/path/writeable.txt", &file));
@@ -948,16 +950,16 @@
TEST(GcsFileSystemTest, NewWritableFile_NoObjectName) {
std::vector<HttpRequest*> requests;
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
std::unique_ptr<WritableFile> file;
EXPECT_EQ(errors::Code::INVALID_ARGUMENT,
@@ -1013,7 +1015,7 @@
std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 32 /* block size */,
32 /* max bytes */, 0 /* max staleness */, 3600 /* stat cache max age */,
0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
kTestTimeoutConfig, *kAllowedLocationsDefault,
nullptr /* gcs additional header */);
@@ -1041,16 +1043,16 @@
TEST(GcsFileSystemTest, NewAppendableFile_NoObjectName) {
std::vector<HttpRequest*> requests;
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
std::unique_ptr<WritableFile> file;
EXPECT_EQ(errors::Code::INVALID_ARGUMENT,
@@ -1075,16 +1077,16 @@
"Range: 0-",
content.size() - 1, "\n", "Timeouts: 5 1 20\n"),
content)});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
std::unique_ptr<ReadOnlyMemoryRegion> region;
TF_EXPECT_OK(fs.NewReadOnlyMemoryRegionFromFile(
@@ -1096,16 +1098,16 @@
TEST(GcsFileSystemTest, NewReadOnlyMemoryRegionFromFile_NoObjectName) {
std::vector<HttpRequest*> requests;
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
std::unique_ptr<ReadOnlyMemoryRegion> region;
EXPECT_EQ(errors::Code::INVALID_ARGUMENT,
@@ -1120,16 +1122,16 @@
"Timeouts: 5 1 10\n",
strings::StrCat("{\"size\": \"1010\",\"generation\": \"1\","
"\"updated\": \"2016-04-29T23:15:24.896Z\"}"))});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
TF_EXPECT_OK(fs.FileExists("gs://bucket/path/file1.txt"));
}
@@ -1150,16 +1152,16 @@
"Timeouts: 5 1 10\n",
"{\"items\": [ "
" { \"name\": \"path/subfolder/\" }]}")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
TF_EXPECT_OK(fs.FileExists("gs://bucket/path/subfolder"));
}
@@ -1176,16 +1178,16 @@
"Auth Token: fake_token\n"
"Timeouts: 5 1 10\n",
"{\"size\": \"100\"}")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
TF_EXPECT_OK(fs.FileExists("gs://bucket1"));
TF_EXPECT_OK(fs.FileExists("gs://bucket1/"));
@@ -1206,16 +1208,16 @@
"Auth Token: fake_token\n"
"Timeouts: 5 1 10\n",
"{\"items\": []}")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
EXPECT_EQ(errors::Code::NOT_FOUND,
fs.FileExists("gs://bucket/path/file1.txt").code());
@@ -1233,16 +1235,16 @@
"Auth Token: fake_token\n"
"Timeouts: 5 1 10\n",
"", errors::NotFound("404"), 404)});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
EXPECT_EQ(errors::Code::INVALID_ARGUMENT,
fs.FileExists("gs://bucket2/").code());
EXPECT_EQ(errors::Code::INVALID_ARGUMENT,
@@ -1279,7 +1281,7 @@
std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
0 /* max bytes */, 0 /* max staleness */, 3600 /* stat cache max age */,
0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
kTestTimeoutConfig, *kAllowedLocationsDefault,
nullptr /* gcs additional header */);
@@ -1306,7 +1308,7 @@
std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
0 /* max bytes */, 0 /* max staleness */, 3600 /* stat cache max age */,
0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
kTestTimeoutConfig, *kAllowedLocationsDefault,
nullptr /* gcs additional header */);
@@ -1322,16 +1324,16 @@
"Auth Token: fake_token\n"
"Timeouts: 5 1 10\n",
"{\"prefixes\": [\"path/subpath/\"]}")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
std::vector<string> children;
TF_EXPECT_OK(fs.GetChildren("gs://bucket/path/", &children));
@@ -1350,16 +1352,16 @@
" { \"name\": \"path/file1.txt\" },"
" { \"name\": \"path/file3.txt\" }],"
"\"prefixes\": [\"path/subpath/\"]}")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
std::vector<string> children;
TF_EXPECT_OK(fs.GetChildren("gs://bucket/path/", &children));
@@ -1379,16 +1381,16 @@
" { \"name\": \"path/\" },"
" { \"name\": \"path/file3.txt\" }],"
"\"prefixes\": [\"path/subpath/\"]}")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
std::vector<string> children;
TF_EXPECT_OK(fs.GetChildren("gs://bucket/path/", &children));
@@ -1407,16 +1409,16 @@
" { \"name\": \"path/file1.txt\" },"
" { \"name\": \"path/file3.txt\" }],"
"\"prefixes\": [\"path/subpath/\"]}")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
std::vector<string> children;
TF_EXPECT_OK(fs.GetChildren("gs://bucket/path", &children));
@@ -1432,16 +1434,16 @@
"Auth Token: fake_token\n"
"Timeouts: 5 1 10\n",
"{}")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
std::vector<string> children;
TF_EXPECT_OK(fs.GetChildren("gs://bucket-a-b-c", &children));
@@ -1457,16 +1459,16 @@
"Auth Token: fake_token\n"
"Timeouts: 5 1 10\n",
"{}")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
std::vector<string> children;
TF_EXPECT_OK(fs.GetChildren("gs://bucket/path/", &children));
@@ -1498,16 +1500,16 @@
" { \"name\": \"path/file4.txt\" },"
" { \"name\": \"path/file5.txt\" }]}")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
std::vector<string> children;
TF_EXPECT_OK(fs.GetChildren("gs://bucket/path", &children));
@@ -1525,16 +1527,16 @@
"Timeouts: 5 1 10\n",
"{\"items\": [ "
" { \"name\": \"path/subpath/file2.txt\" }]}")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
std::vector<string> result;
TF_EXPECT_OK(
@@ -1553,16 +1555,16 @@
" { \"name\": \"path/file1.txt\" },"
" { \"name\": \"path/subpath/file2.txt\" },"
" { \"name\": \"path/file3.txt\" }]}")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
std::vector<string> result;
TF_EXPECT_OK(fs.GetMatchingPaths("gs://bucket/*/*", &result));
@@ -1582,16 +1584,16 @@
" { \"name\": \"path/file1.txt\" },"
" { \"name\": \"path/subpath/file2.txt\" },"
" { \"name\": \"path/file3.txt\" }]}")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
std::vector<string> result;
TF_EXPECT_OK(fs.GetMatchingPaths("gs://bucket/path/*/file2.txt", &result));
@@ -1608,16 +1610,16 @@
"{\"items\": [ "
" { \"name\": \"path/\" },"
" { \"name\": \"path/file3.txt\" }]}")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
std::vector<string> result;
TF_EXPECT_OK(fs.GetMatchingPaths("gs://bucket/path/*", &result));
@@ -1634,16 +1636,16 @@
" { \"name\": \"path/file1.txt\" },"
" { \"name\": \"path/subpath/file2.txt\" },"
" { \"name\": \"path/file3.txt\" }]}")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
std::vector<string> result;
TF_EXPECT_OK(fs.GetMatchingPaths("gs://bucket/path/*/file3.txt", &result));
@@ -1652,16 +1654,16 @@
TEST(GcsFileSystemTest, GetMatchingPaths_OnlyWildcard) {
std::vector<HttpRequest*> requests;
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
std::vector<string> result;
EXPECT_EQ(errors::Code::INVALID_ARGUMENT,
@@ -1686,16 +1688,16 @@
" { \"name\": \"path/file1.txt\" },"
" { \"name\": \"path/subpath/file2.txt\" },"
" { \"name\": \"path/file3.txt\" }]}")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 3600 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 3600 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
// Repeated calls to fs.GetMatchingPaths on these patterns should not lead to
// any additional HTTP requests to GCS.
@@ -1729,16 +1731,16 @@
"Timeouts: 5 1 10\n",
"{\"items\": [ "
" { \"name\": \"path/subpath/file2.txt\" }]}")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 3600 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 3600 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
// This loop should trigger the first HTTP request to GCS.
for (int i = 0; i < 10; i++) {
@@ -1800,7 +1802,7 @@
std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 16 /* block size */,
16 /* max bytes */, 0 /* max staleness */, 3600 /* stat cache max age */,
0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
kTestTimeoutConfig, *kAllowedLocationsDefault,
nullptr /* gcs additional header */);
@@ -1821,16 +1823,16 @@
TEST(GcsFileSystemTest, DeleteFile_NoObjectName) {
std::vector<HttpRequest*> requests;
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
EXPECT_EQ(errors::Code::INVALID_ARGUMENT,
fs.DeleteFile("gs://bucket/").code());
@@ -1871,7 +1873,7 @@
std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 16 /* block size */,
16 /* max bytes */, 0 /* max staleness */, 3600 /* stat cache max age */,
0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
kTestTimeoutConfig, *kAllowedLocationsDefault,
nullptr /* gcs additional header */);
@@ -1894,16 +1896,16 @@
"Auth Token: fake_token\n"
"Timeouts: 5 1 10\n",
"{}")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
TF_EXPECT_OK(fs.DeleteDir("gs://bucket/path/"));
}
@@ -1923,16 +1925,16 @@
"Timeouts: 5 1 10\n"
"Delete: yes\n",
"")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
TF_EXPECT_OK(fs.DeleteDir("gs://bucket/path/"));
}
@@ -1943,16 +1945,16 @@
"name%2CnextPageToken&maxResults=2\nAuth Token: fake_token\n"
"Timeouts: 5 1 10\n",
"{}")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
TF_EXPECT_OK(fs.DeleteDir("gs://bucket"));
}
@@ -1965,16 +1967,16 @@
"Timeouts: 5 1 10\n",
"{\"items\": [ "
" { \"name\": \"path/file1.txt\" }]}")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
EXPECT_EQ(error::Code::FAILED_PRECONDITION,
fs.DeleteDir("gs://bucket/path/").code());
@@ -1988,16 +1990,16 @@
"Timeouts: 5 1 10\n",
strings::StrCat("{\"size\": \"1010\",\"generation\": \"1\","
"\"updated\": \"2016-04-29T23:15:24.896Z\"}"))});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
uint64 size;
TF_EXPECT_OK(fs.GetFileSize("gs://bucket/file.txt", &size));
@@ -2006,16 +2008,16 @@
TEST(GcsFileSystemTest, GetFileSize_NoObjectName) {
std::vector<HttpRequest*> requests;
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
uint64 size;
EXPECT_EQ(errors::Code::INVALID_ARGUMENT,
@@ -2092,16 +2094,16 @@
"Timeouts: 5 1 10\n"
"Delete: yes\n",
"")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
TF_EXPECT_OK(fs.RenameFile("gs://bucket/path1", "gs://bucket/path2/"));
}
@@ -2191,7 +2193,7 @@
std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 16 /* block size */,
64 /* max bytes */, 0 /* max staleness */, 3600 /* stat cache max age */,
0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
kTestTimeoutConfig, *kAllowedLocationsDefault,
nullptr /* gcs additional header */);
// Do an initial read of the source and destination files to load their
@@ -2272,7 +2274,7 @@
std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
0 /* max bytes */, 0 /* max staleness */, 3600 /* stat cache max age */,
0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
kTestTimeoutConfig, *kAllowedLocationsDefault,
nullptr /* gcs additional header */);
// Do an initial stat of the destination file to load their contents into the
@@ -2332,16 +2334,16 @@
"Timeouts: 5 1 10\n"
"Delete: yes\n",
"", errors::NotFound("404"), 404)});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
TF_EXPECT_OK(
fs.RenameFile("gs://bucket/path/src.txt", "gs://bucket/path/dst.txt"));
@@ -2374,16 +2376,16 @@
"Post: yes\n"
"Timeouts: 5 1 10\n",
"{\"done\": false}")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
EXPECT_EQ(
errors::Code::UNIMPLEMENTED,
@@ -2399,16 +2401,16 @@
"Timeouts: 5 1 10\n",
strings::StrCat("{\"size\": \"1010\",\"generation\": \"1\","
"\"updated\": \"2016-04-29T23:15:24.896Z\"}"))});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
FileStatistics stat;
TF_EXPECT_OK(fs.Stat("gs://bucket/file.txt", &stat));
@@ -2433,16 +2435,16 @@
"Timeouts: 5 1 10\n",
"{\"items\": [ "
" { \"name\": \"subfolder/\" }]}")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
FileStatistics stat;
TF_EXPECT_OK(fs.Stat("gs://bucket/subfolder", &stat));
@@ -2466,16 +2468,16 @@
"Auth Token: fake_token\n"
"Timeouts: 5 1 10\n",
"{}")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
FileStatistics stat;
EXPECT_EQ(error::Code::NOT_FOUND, fs.Stat("gs://bucket/path", &stat).code());
@@ -2487,16 +2489,16 @@
"Auth Token: fake_token\n"
"Timeouts: 5 1 10\n",
"{}")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
FileStatistics stat;
TF_EXPECT_OK(fs.Stat("gs://bucket/", &stat));
@@ -2511,16 +2513,16 @@
"Auth Token: fake_token\n"
"Timeouts: 5 1 10\n",
"", errors::NotFound("404"), 404)});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
FileStatistics stat;
EXPECT_EQ(error::Code::NOT_FOUND, fs.Stat("gs://bucket/", &stat).code());
@@ -2556,7 +2558,7 @@
std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
0 /* max bytes */, 0 /* max staleness */, 3600 /* stat cache max age */,
0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
kTestTimeoutConfig, *kAllowedLocationsDefault,
nullptr /* gcs additional header */);
@@ -2598,7 +2600,7 @@
std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
0 /* max bytes */, 0 /* max staleness */, 3600 /* stat cache max age */,
0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
kTestTimeoutConfig, *kAllowedLocationsDefault,
nullptr /* gcs additional header */);
// There should be a single HTTP request to GCS for fs.Stat in this loop.
@@ -2628,16 +2630,16 @@
"Timeouts: 5 1 10\n",
strings::StrCat("{\"size\": \"5\",\"generation\": \"1\","
"\"updated\": \"2016-04-29T23:15:24.896Z\"}"))});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
FileStatistics stat;
TF_EXPECT_OK(fs.Stat("gs://bucket/dir/", &stat));
@@ -2660,16 +2662,16 @@
"Auth Token: fake_token\n"
"Timeouts: 5 1 10\n",
"", errors::NotFound("404"), 404)});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
EXPECT_EQ(error::Code::NOT_FOUND,
fs.IsDirectory("gs://bucket/file.txt").code());
@@ -2691,16 +2693,16 @@
"Timeouts: 5 1 10\n",
strings::StrCat("{\"size\": \"1010\",\"generation\": \"1\","
"\"updated\": \"2016-04-29T23:15:24.896Z\"}"))});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
EXPECT_EQ(error::Code::FAILED_PRECONDITION,
fs.IsDirectory("gs://bucket/file.txt").code());
@@ -2722,16 +2724,16 @@
"Auth Token: fake_token\n"
"Timeouts: 5 1 10\n",
"{\"items\": [{\"name\": \"subfolder/\"}]}")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
TF_EXPECT_OK(fs.IsDirectory("gs://bucket/subfolder"));
TF_EXPECT_OK(fs.IsDirectory("gs://bucket/subfolder/"));
@@ -2749,16 +2751,16 @@
"Auth Token: fake_token\n"
"Timeouts: 5 1 10\n",
"{}")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
TF_EXPECT_OK(fs.IsDirectory("gs://bucket"));
TF_EXPECT_OK(fs.IsDirectory("gs://bucket/"));
@@ -2770,16 +2772,16 @@
"Auth Token: fake_token\n"
"Timeouts: 5 1 10\n",
"", errors::NotFound("404"), 404)});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
EXPECT_EQ(error::Code::NOT_FOUND, fs.IsDirectory("gs://bucket/").code());
}
@@ -2812,16 +2814,16 @@
"Timeouts: 5 1 30\n"
"Put body: \n",
"")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
TF_EXPECT_OK(fs.CreateDir("gs://bucket/subpath"));
TF_EXPECT_OK(fs.CreateDir("gs://bucket/subpath/"));
@@ -2839,16 +2841,16 @@
"Auth Token: fake_token\n"
"Timeouts: 5 1 10\n",
"")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
TF_EXPECT_OK(fs.CreateDir("gs://bucket/"));
TF_EXPECT_OK(fs.CreateDir("gs://bucket"));
@@ -2911,16 +2913,16 @@
"Timeouts: 5 1 10\n"
"Delete: yes\n",
"")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
int64 undeleted_files, undeleted_dirs;
TF_EXPECT_OK(fs.DeleteRecursively("gs://bucket/path", &undeleted_files,
@@ -3004,16 +3006,16 @@
"Timeouts: 5 1 10\n",
"", errors::NotFound("404"), 404)});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
int64 undeleted_files, undeleted_dirs;
TF_EXPECT_OK(fs.DeleteRecursively("gs://bucket/path", &undeleted_files,
@@ -3039,16 +3041,16 @@
"Auth Token: fake_token\n"
"Timeouts: 5 1 10\n",
"", errors::NotFound("404"), 404)});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
int64 undeleted_files, undeleted_dirs;
EXPECT_EQ(error::Code::NOT_FOUND,
@@ -3130,7 +3132,7 @@
std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
kTestTimeoutConfig, *kAllowedLocationsDefault,
add_header /* gcs additional header */);
@@ -3199,16 +3201,16 @@
"Auth Token: fake_token\n"
"Header Hello: world\n",
"{}")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
std::unique_ptr<HttpRequest> request;
TF_EXPECT_OK(fs.CreateHttpRequest(&request));
@@ -3262,16 +3264,16 @@
"Timeouts: 5 1 10\n",
strings::StrCat("{\"size\": \"1010\",\"generation\": \"1\","
"\"updated\": \"2016-04-29T23:15:24.896Z\"}"))});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
TestGcsStats stats;
fs.SetStats(&stats);
@@ -3289,16 +3291,16 @@
"Range: 0-5\n"
"Timeouts: 5 1 20\n",
"012345")});
- GcsFileSystem fs(
- std::unique_ptr<AuthProvider>(new FakeAuthProvider),
- std::unique_ptr<HttpRequest::Factory>(
- new FakeHttpRequestFactory(&requests)),
- std::unique_ptr<ZoneProvider>(new FakeZoneProvider), 0 /* block size */,
- 0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
- 0 /* stat cache max entries */, 0 /* matching paths cache max age */,
- 0 /* matching paths cache max entries */, 0 /* initial retry delay */,
- kTestTimeoutConfig, *kAllowedLocationsDefault,
- nullptr /* gcs additional header */);
+ GcsFileSystem fs(std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+ std::unique_ptr<HttpRequest::Factory>(
+ new FakeHttpRequestFactory(&requests)),
+ std::unique_ptr<ZoneProvider>(new FakeZoneProvider),
+ 0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+ 0 /* stat cache max age */, 0 /* stat cache max entries */,
+ 0 /* matching paths cache max age */,
+ 0 /* matching paths cache max entries */, kTestRetryConfig,
+ kTestTimeoutConfig, *kAllowedLocationsDefault,
+ nullptr /* gcs additional header */);
TestGcsStats stats;
fs.SetStats(&stats);
diff --git a/tensorflow/core/platform/cloud/google_auth_provider_test.cc b/tensorflow/core/platform/cloud/google_auth_provider_test.cc
index 07b88a8..ec31c5e 100644
--- a/tensorflow/core/platform/cloud/google_auth_provider_test.cc
+++ b/tensorflow/core/platform/cloud/google_auth_provider_test.cc
@@ -93,8 +93,8 @@
std::shared_ptr<HttpRequest::Factory> fakeHttpRequestFactory =
std::make_shared<FakeHttpRequestFactory>(&requests);
- auto metadataClient =
- std::make_shared<ComputeEngineMetadataClient>(fakeHttpRequestFactory, 0);
+ auto metadataClient = std::make_shared<ComputeEngineMetadataClient>(
+ fakeHttpRequestFactory, RetryConfig(0 /* init_delay_time_us */));
GoogleAuthProvider provider(std::unique_ptr<OAuthClient>(oauth_client),
metadataClient, &env);
oauth_client->return_token = "fake-token";
@@ -129,8 +129,8 @@
FakeEnv env;
std::shared_ptr<HttpRequest::Factory> fakeHttpRequestFactory =
std::make_shared<FakeHttpRequestFactory>(&requests);
- auto metadataClient =
- std::make_shared<ComputeEngineMetadataClient>(fakeHttpRequestFactory, 0);
+ auto metadataClient = std::make_shared<ComputeEngineMetadataClient>(
+ fakeHttpRequestFactory, RetryConfig(0 /* init_delay_time_us */));
GoogleAuthProvider provider(std::unique_ptr<OAuthClient>(oauth_client),
metadataClient, &env);
@@ -178,8 +178,8 @@
FakeEnv env;
std::shared_ptr<HttpRequest::Factory> fakeHttpRequestFactory =
std::make_shared<FakeHttpRequestFactory>(&requests);
- auto metadataClient =
- std::make_shared<ComputeEngineMetadataClient>(fakeHttpRequestFactory, 0);
+ auto metadataClient = std::make_shared<ComputeEngineMetadataClient>(
+ fakeHttpRequestFactory, RetryConfig(0 /* init_delay_time_us */));
GoogleAuthProvider provider(std::unique_ptr<OAuthClient>(oauth_client),
metadataClient, &env);
@@ -206,8 +206,8 @@
FakeEnv env;
std::shared_ptr<HttpRequest::Factory> fakeHttpRequestFactory =
std::make_shared<FakeHttpRequestFactory>(&empty_requests);
- auto metadataClient =
- std::make_shared<ComputeEngineMetadataClient>(fakeHttpRequestFactory, 0);
+ auto metadataClient = std::make_shared<ComputeEngineMetadataClient>(
+ fakeHttpRequestFactory, RetryConfig(0 /* init_delay_time_us */));
GoogleAuthProvider provider(std::unique_ptr<OAuthClient>(oauth_client),
metadataClient, &env);
@@ -228,8 +228,8 @@
FakeEnv env;
std::shared_ptr<HttpRequest::Factory> fakeHttpRequestFactory =
std::make_shared<FakeHttpRequestFactory>(&requests);
- auto metadataClient =
- std::make_shared<ComputeEngineMetadataClient>(fakeHttpRequestFactory, 0);
+ auto metadataClient = std::make_shared<ComputeEngineMetadataClient>(
+ fakeHttpRequestFactory, RetryConfig(0 /* init_delay_time_us */));
GoogleAuthProvider provider(std::unique_ptr<OAuthClient>(oauth_client),
metadataClient, &env);
diff --git a/tensorflow/core/platform/cloud/retrying_file_system.h b/tensorflow/core/platform/cloud/retrying_file_system.h
index 941ab7a..5ce6670 100644
--- a/tensorflow/core/platform/cloud/retrying_file_system.h
+++ b/tensorflow/core/platform/cloud/retrying_file_system.h
@@ -34,9 +34,9 @@
class RetryingFileSystem : public FileSystem {
public:
RetryingFileSystem(std::unique_ptr<Underlying> base_file_system,
- int64 delay_microseconds = 1000000)
+ const RetryConfig& retry_config)
: base_file_system_(std::move(base_file_system)),
- initial_delay_microseconds_(delay_microseconds) {}
+ retry_config_(retry_config) {}
Status NewRandomAccessFile(
const string& filename,
@@ -55,7 +55,7 @@
Status FileExists(const string& fname) override {
return RetryingUtils::CallWithRetries(
[this, &fname]() { return base_file_system_->FileExists(fname); },
- initial_delay_microseconds_);
+ retry_config_);
}
Status GetChildren(const string& dir, std::vector<string>* result) override {
@@ -63,7 +63,7 @@
[this, &dir, result]() {
return base_file_system_->GetChildren(dir, result);
},
- initial_delay_microseconds_);
+ retry_config_);
}
Status GetMatchingPaths(const string& pattern,
@@ -72,31 +72,31 @@
[this, &pattern, result]() {
return base_file_system_->GetMatchingPaths(pattern, result);
},
- initial_delay_microseconds_);
+ retry_config_);
}
Status Stat(const string& fname, FileStatistics* stat) override {
return RetryingUtils::CallWithRetries(
[this, &fname, stat]() { return base_file_system_->Stat(fname, stat); },
- initial_delay_microseconds_);
+ retry_config_);
}
Status DeleteFile(const string& fname) override {
return RetryingUtils::DeleteWithRetries(
[this, &fname]() { return base_file_system_->DeleteFile(fname); },
- initial_delay_microseconds_);
+ retry_config_);
}
Status CreateDir(const string& dirname) override {
return RetryingUtils::CallWithRetries(
[this, &dirname]() { return base_file_system_->CreateDir(dirname); },
- initial_delay_microseconds_);
+ retry_config_);
}
Status DeleteDir(const string& dirname) override {
return RetryingUtils::DeleteWithRetries(
[this, &dirname]() { return base_file_system_->DeleteDir(dirname); },
- initial_delay_microseconds_);
+ retry_config_);
}
Status GetFileSize(const string& fname, uint64* file_size) override {
@@ -104,7 +104,7 @@
[this, &fname, file_size]() {
return base_file_system_->GetFileSize(fname, file_size);
},
- initial_delay_microseconds_);
+ retry_config_);
}
Status RenameFile(const string& src, const string& target) override {
@@ -112,13 +112,13 @@
[this, &src, &target]() {
return base_file_system_->RenameFile(src, target);
},
- initial_delay_microseconds_);
+ retry_config_);
}
Status IsDirectory(const string& dirname) override {
return RetryingUtils::CallWithRetries(
[this, &dirname]() { return base_file_system_->IsDirectory(dirname); },
- initial_delay_microseconds_);
+ retry_config_);
}
Status DeleteRecursively(const string& dirname, int64* undeleted_files,
@@ -128,7 +128,7 @@
return base_file_system_->DeleteRecursively(dirname, undeleted_files,
undeleted_dirs);
},
- initial_delay_microseconds_);
+ retry_config_);
}
void FlushCaches() override { base_file_system_->FlushCaches(); }
@@ -137,7 +137,7 @@
private:
std::unique_ptr<Underlying> base_file_system_;
- const int64 initial_delay_microseconds_;
+ const RetryConfig retry_config_;
TF_DISALLOW_COPY_AND_ASSIGN(RetryingFileSystem);
};
@@ -147,9 +147,8 @@
class RetryingRandomAccessFile : public RandomAccessFile {
public:
RetryingRandomAccessFile(std::unique_ptr<RandomAccessFile> base_file,
- int64 delay_microseconds)
- : base_file_(std::move(base_file)),
- initial_delay_microseconds_(delay_microseconds) {}
+ const RetryConfig& retry_config)
+ : base_file_(std::move(base_file)), retry_config_(retry_config) {}
Status Read(uint64 offset, size_t n, StringPiece* result,
char* scratch) const override {
@@ -157,20 +156,19 @@
[this, offset, n, result, scratch]() {
return base_file_->Read(offset, n, result, scratch);
},
- initial_delay_microseconds_);
+ retry_config_);
}
private:
std::unique_ptr<RandomAccessFile> base_file_;
- const int64 initial_delay_microseconds_;
+ const RetryConfig retry_config_;
};
class RetryingWritableFile : public WritableFile {
public:
RetryingWritableFile(std::unique_ptr<WritableFile> base_file,
- int64 delay_microseconds)
- : base_file_(std::move(base_file)),
- initial_delay_microseconds_(delay_microseconds) {}
+ const RetryConfig& retry_config)
+ : base_file_(std::move(base_file)), retry_config_(retry_config) {}
~RetryingWritableFile() override {
// Makes sure the retrying version of Close() is called in the destructor.
@@ -179,25 +177,24 @@
Status Append(StringPiece data) override {
return RetryingUtils::CallWithRetries(
- [this, &data]() { return base_file_->Append(data); },
- initial_delay_microseconds_);
+ [this, &data]() { return base_file_->Append(data); }, retry_config_);
}
Status Close() override {
return RetryingUtils::CallWithRetries(
- [this]() { return base_file_->Close(); }, initial_delay_microseconds_);
+ [this]() { return base_file_->Close(); }, retry_config_);
}
Status Flush() override {
return RetryingUtils::CallWithRetries(
- [this]() { return base_file_->Flush(); }, initial_delay_microseconds_);
+ [this]() { return base_file_->Flush(); }, retry_config_);
}
Status Sync() override {
return RetryingUtils::CallWithRetries(
- [this]() { return base_file_->Sync(); }, initial_delay_microseconds_);
+ [this]() { return base_file_->Sync(); }, retry_config_);
}
private:
std::unique_ptr<WritableFile> base_file_;
- const int64 initial_delay_microseconds_;
+ const RetryConfig retry_config_;
};
} // namespace retrying_internals
@@ -210,9 +207,9 @@
[this, &filename, &base_file]() {
return base_file_system_->NewRandomAccessFile(filename, &base_file);
},
- initial_delay_microseconds_));
+ retry_config_));
result->reset(new retrying_internals::RetryingRandomAccessFile(
- std::move(base_file), initial_delay_microseconds_));
+ std::move(base_file), retry_config_));
return Status::OK();
}
@@ -224,9 +221,9 @@
[this, &filename, &base_file]() {
return base_file_system_->NewWritableFile(filename, &base_file);
},
- initial_delay_microseconds_));
+ retry_config_));
result->reset(new retrying_internals::RetryingWritableFile(
- std::move(base_file), initial_delay_microseconds_));
+ std::move(base_file), retry_config_));
return Status::OK();
}
@@ -238,9 +235,9 @@
[this, &filename, &base_file]() {
return base_file_system_->NewAppendableFile(filename, &base_file);
},
- initial_delay_microseconds_));
+ retry_config_));
result->reset(new retrying_internals::RetryingWritableFile(
- std::move(base_file), initial_delay_microseconds_));
+ std::move(base_file), retry_config_));
return Status::OK();
}
@@ -252,7 +249,7 @@
return base_file_system_->NewReadOnlyMemoryRegionFromFile(filename,
result);
},
- initial_delay_microseconds_);
+ retry_config_);
}
} // namespace tensorflow
diff --git a/tensorflow/core/platform/cloud/retrying_file_system_test.cc b/tensorflow/core/platform/cloud/retrying_file_system_test.cc
index 5910fef..868eea0 100644
--- a/tensorflow/core/platform/cloud/retrying_file_system_test.cc
+++ b/tensorflow/core/platform/cloud/retrying_file_system_test.cc
@@ -184,7 +184,8 @@
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
base_fs->random_access_file_to_return = std::move(base_file);
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
// Retrieve the wrapped random access file.
std::unique_ptr<RandomAccessFile> random_access_file;
@@ -211,7 +212,8 @@
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
base_fs->random_access_file_to_return = std::move(base_file);
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
// Retrieve the wrapped random access file.
std::unique_ptr<RandomAccessFile> random_access_file;
@@ -235,7 +237,8 @@
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
base_fs->random_access_file_to_return = std::move(base_file);
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
// Retrieve the wrapped random access file.
std::unique_ptr<RandomAccessFile> random_access_file;
@@ -265,7 +268,8 @@
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
base_fs->random_access_file_to_return = std::move(base_file);
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
// Retrieve the wrapped random access file.
std::unique_ptr<RandomAccessFile> random_access_file;
@@ -291,7 +295,8 @@
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
base_fs->writable_file_to_return = std::move(base_file);
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
// Retrieve the wrapped writable file.
std::unique_ptr<WritableFile> writable_file;
@@ -317,7 +322,8 @@
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
base_fs->writable_file_to_return = std::move(base_file);
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
// Retrieve the wrapped writable file.
std::unique_ptr<WritableFile> writable_file;
@@ -343,7 +349,8 @@
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
base_fs->writable_file_to_return = std::move(base_file);
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
// Retrieve the wrapped writable file.
std::unique_ptr<WritableFile> writable_file;
@@ -368,7 +375,8 @@
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
base_fs->writable_file_to_return = std::move(base_file);
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
// Retrieve the wrapped appendable file.
std::unique_ptr<WritableFile> writable_file;
@@ -391,7 +399,8 @@
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
base_fs->writable_file_to_return = std::move(base_file);
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
// Retrieve the wrapped writable file.
std::unique_ptr<WritableFile> writable_file;
@@ -412,7 +421,8 @@
std::make_tuple("NewReadOnlyMemoryRegionFromFile", Status::OK())});
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
std::unique_ptr<ReadOnlyMemoryRegion> result;
TF_EXPECT_OK(fs.NewReadOnlyMemoryRegionFromFile("filename.txt", &result));
@@ -423,7 +433,8 @@
CreateRetriableErrors("NewReadOnlyMemoryRegionFromFile", 11);
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
std::unique_ptr<ReadOnlyMemoryRegion> result;
const auto& status =
@@ -440,7 +451,8 @@
std::make_tuple("GetChildren", Status::OK())});
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
std::vector<string> result;
TF_EXPECT_OK(fs.GetChildren("gs://path", &result));
@@ -450,7 +462,8 @@
ExpectedCalls expected_fs_calls = CreateRetriableErrors("GetChildren", 11);
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
std::vector<string> result;
const auto& status = fs.GetChildren("gs://path", &result);
@@ -466,7 +479,8 @@
std::make_tuple("GetMatchingPaths", Status::OK())});
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
std::vector<string> result;
TF_EXPECT_OK(fs.GetMatchingPaths("gs://path/dir", &result));
@@ -477,7 +491,8 @@
CreateRetriableErrors("GetMatchingPaths", 11);
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
std::vector<string> result;
const auto& status = fs.GetMatchingPaths("gs://path/dir", &result);
@@ -492,7 +507,8 @@
std::make_tuple("DeleteFile", Status::OK())});
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
std::vector<string> result;
TF_EXPECT_OK(fs.DeleteFile("gs://path/file.txt"));
@@ -502,7 +518,8 @@
ExpectedCalls expected_fs_calls = CreateRetriableErrors("DeleteFile", 11);
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
std::vector<string> result;
const auto& status = fs.DeleteFile("gs://path/file.txt");
@@ -517,7 +534,8 @@
std::make_tuple("CreateDir", Status::OK())});
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
std::vector<string> result;
TF_EXPECT_OK(fs.CreateDir("gs://path/newdir"));
@@ -527,7 +545,8 @@
ExpectedCalls expected_fs_calls = CreateRetriableErrors("CreateDir", 11);
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
std::vector<string> result;
const auto& status = fs.CreateDir("gs://path/newdir");
@@ -542,7 +561,8 @@
std::make_tuple("DeleteDir", Status::OK())});
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
std::vector<string> result;
TF_EXPECT_OK(fs.DeleteDir("gs://path/dir"));
@@ -552,7 +572,8 @@
ExpectedCalls expected_fs_calls = CreateRetriableErrors("DeleteDir", 11);
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
std::vector<string> result;
const auto& status = fs.DeleteDir("gs://path/dir");
@@ -568,7 +589,8 @@
std::make_tuple("GetFileSize", Status::OK())});
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
uint64 size;
TF_EXPECT_OK(fs.GetFileSize("gs://path/file.txt", &size));
@@ -578,7 +600,8 @@
ExpectedCalls expected_fs_calls = CreateRetriableErrors("GetFileSize", 11);
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
uint64 size;
const auto& status = fs.GetFileSize("gs://path/file.txt", &size);
@@ -593,7 +616,8 @@
std::make_tuple("RenameFile", Status::OK())});
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
TF_EXPECT_OK(fs.RenameFile("old_name", "new_name"));
}
@@ -602,7 +626,8 @@
ExpectedCalls expected_fs_calls = CreateRetriableErrors("RenameFile", 11);
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
const auto& status = fs.RenameFile("old_name", "new_name");
EXPECT_TRUE(
@@ -616,7 +641,8 @@
std::make_tuple("Stat", Status::OK())});
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
FileStatistics stat;
TF_EXPECT_OK(fs.Stat("file_name", &stat));
@@ -626,7 +652,8 @@
ExpectedCalls expected_fs_calls = CreateRetriableErrors("Stat", 11);
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
FileStatistics stat;
const auto& status = fs.Stat("file_name", &stat);
@@ -639,7 +666,8 @@
ExpectedCalls expected_fs_calls = CreateRetriableErrors("FileExists", 11);
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
const auto& status = fs.FileExists("file_name");
EXPECT_TRUE(
@@ -653,7 +681,8 @@
std::make_tuple("FileExists", Status::OK())});
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
TF_EXPECT_OK(fs.FileExists("gs://path/dir"));
}
@@ -665,7 +694,8 @@
std::make_tuple("IsDirectory", Status::OK())});
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
TF_EXPECT_OK(fs.IsDirectory("gs://path/dir"));
}
@@ -674,7 +704,8 @@
ExpectedCalls expected_fs_calls = CreateRetriableErrors("IsDirectory", 11);
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
const auto& status = fs.IsDirectory("gs://path/dir");
EXPECT_TRUE(
@@ -689,7 +720,8 @@
std::make_tuple("DeleteRecursively", Status::OK())});
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
int64 undeleted_files, undeleted_dirs;
TF_EXPECT_OK(
@@ -701,7 +733,8 @@
CreateRetriableErrors("DeleteRecursively", 11);
std::unique_ptr<MockFileSystem> base_fs(
new MockFileSystem(expected_fs_calls));
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
int64 undeleted_files, undeleted_dirs;
const auto& status =
@@ -715,7 +748,8 @@
ExpectedCalls none;
bool flushed = false;
std::unique_ptr<MockFileSystem> base_fs(new MockFileSystem(none, &flushed));
- RetryingFileSystem<MockFileSystem> fs(std::move(base_fs), 0);
+ RetryingFileSystem<MockFileSystem> fs(
+ std::move(base_fs), RetryConfig(0 /* init_delay_time_us */));
fs.FlushCaches();
EXPECT_TRUE(flushed);
}
diff --git a/tensorflow/core/platform/cloud/retrying_utils.cc b/tensorflow/core/platform/cloud/retrying_utils.cc
index d2df422..cb0aecd 100644
--- a/tensorflow/core/platform/cloud/retrying_utils.cc
+++ b/tensorflow/core/platform/cloud/retrying_utils.cc
@@ -23,11 +23,6 @@
namespace {
-// In case of failure, every call will be retried kMaxRetries times.
-constexpr int kMaxRetries = 10;
-// Maximum backoff time in microseconds.
-constexpr int64 kMaximumBackoffMicroseconds = 32000000; // 32 seconds.
-
bool IsRetriable(error::Code code) {
switch (code) {
case error::UNAVAILABLE:
@@ -43,40 +38,41 @@
} // namespace
Status RetryingUtils::CallWithRetries(const std::function<Status()>& f,
- const int64 initial_delay_microseconds) {
- return CallWithRetries(f, initial_delay_microseconds, [](int64 micros) {
- return Env::Default()->SleepForMicroseconds(micros);
- });
+ const RetryConfig& config) {
+ return CallWithRetries(
+ f,
+ [](int64 micros) { return Env::Default()->SleepForMicroseconds(micros); },
+ config);
}
Status RetryingUtils::CallWithRetries(
- const std::function<Status()>& f, const int64 initial_delay_microseconds,
- const std::function<void(int64)>& sleep_usec) {
+ const std::function<Status()>& f,
+ const std::function<void(int64)>& sleep_usec, const RetryConfig& config) {
int retries = 0;
while (true) {
auto status = f();
if (!IsRetriable(status.code())) {
return status;
}
- if (retries >= kMaxRetries) {
+ if (retries >= config.max_retries) {
// Return AbortedError, so that it doesn't get retried again somewhere
// at a higher level.
return Status(
error::ABORTED,
strings::StrCat(
- "All ", kMaxRetries,
+ "All ", config.max_retries,
" retry attempts failed. The last failure: ", status.ToString()));
}
int64 delay_micros = 0;
- if (initial_delay_microseconds > 0) {
+ if (config.init_delay_time_us > 0) {
const int64 random_micros = random::New64() % 1000000;
- delay_micros = std::min(initial_delay_microseconds << retries,
- kMaximumBackoffMicroseconds) +
+ delay_micros = std::min(config.init_delay_time_us << retries,
+ config.max_delay_time_us) +
random_micros;
}
LOG(INFO) << "The operation failed and will be automatically retried in "
<< (delay_micros / 1000000.0) << " seconds (attempt "
- << (retries + 1) << " out of " << kMaxRetries
+ << (retries + 1) << " out of " << config.max_retries
<< "), caused by: " << status.ToString();
sleep_usec(delay_micros);
retries++;
@@ -84,8 +80,7 @@
}
Status RetryingUtils::DeleteWithRetries(
- const std::function<Status()>& delete_func,
- const int64 initial_delay_microseconds) {
+ const std::function<Status()>& delete_func, const RetryConfig& config) {
bool is_retried = false;
return RetryingUtils::CallWithRetries(
[delete_func, &is_retried]() {
@@ -96,7 +91,7 @@
is_retried = true;
return status;
},
- initial_delay_microseconds);
+ config);
}
} // namespace tensorflow
diff --git a/tensorflow/core/platform/cloud/retrying_utils.h b/tensorflow/core/platform/cloud/retrying_utils.h
index 546b8d1..1a7ce1b 100644
--- a/tensorflow/core/platform/cloud/retrying_utils.h
+++ b/tensorflow/core/platform/cloud/retrying_utils.h
@@ -21,6 +21,26 @@
namespace tensorflow {
+// Default time before reporting failure: ~100 seconds.
+struct RetryConfig {
+ RetryConfig(int64 init_delay_time_us = 100 * 1000,
+ int64 max_delay_time_us = 32 * 1000 * 1000,
+ int max_retries = 10) {
+ this->init_delay_time_us = init_delay_time_us;
+ this->max_delay_time_us = max_delay_time_us;
+ this->max_retries = max_retries;
+ }
+
+ // In case of failure, every call will be retried max_retries times.
+ int max_retries;
+
+ // Initial backoff time
+ int64 init_delay_time_us;
+
+ // Maximum backoff time in microseconds.
+ int64 max_delay_time_us;
+};
+
class RetryingUtils {
public:
/// \brief Retries the function in case of failure with exponential backoff.
@@ -31,18 +51,19 @@
/// retries.
/// If all retries failed, returns the last error status.
static Status CallWithRetries(const std::function<Status()>& f,
- const int64 initial_delay_microseconds);
+ const RetryConfig& config);
+
/// sleep_usec is a function that sleeps for the given number of microseconds.
static Status CallWithRetries(const std::function<Status()>& f,
- const int64 initial_delay_microseconds,
- const std::function<void(int64)>& sleep_usec);
+ const std::function<void(int64)>& sleep_usec,
+ const RetryConfig& config);
/// \brief A retrying wrapper for a function that deletes a resource.
///
/// The function takes care of the scenario when a delete operation
/// returns a failure but succeeds under the hood: if a retry returns
/// NOT_FOUND, the whole operation is considered a success.
static Status DeleteWithRetries(const std::function<Status()>& delete_func,
- const int64 initial_delay_microseconds);
+ const RetryConfig& config);
};
} // namespace tensorflow
diff --git a/tensorflow/core/platform/cloud/retrying_utils_test.cc b/tensorflow/core/platform/cloud/retrying_utils_test.cc
index 1b65276..75fe8a9 100644
--- a/tensorflow/core/platform/cloud/retrying_utils_test.cc
+++ b/tensorflow/core/platform/cloud/retrying_utils_test.cc
@@ -30,7 +30,8 @@
};
std::function<Status()> f = []() { return errors::Unavailable("Failed."); };
- const auto& status = RetryingUtils::CallWithRetries(f, 500000L, sleep);
+ const auto& status = RetryingUtils::CallWithRetries(
+ f, sleep, RetryConfig(500000 /* init_delay_time_us */));
EXPECT_EQ(errors::Code::ABORTED, status.code());
EXPECT_TRUE(str_util::StrContains(
status.error_message(),
@@ -60,8 +61,10 @@
results.erase(results.begin());
return result;
};
- EXPECT_EQ(errors::Code::NOT_FOUND,
- RetryingUtils::CallWithRetries(f, 0).code());
+ EXPECT_EQ(
+ errors::Code::NOT_FOUND,
+ RetryingUtils::CallWithRetries(f, RetryConfig(0 /* init_delay_time_us */))
+ .code());
}
TEST(RetryingUtilsTest, CallWithRetries_ImmediateSuccess) {
@@ -74,7 +77,8 @@
results.erase(results.begin());
return result;
};
- TF_EXPECT_OK(RetryingUtils::CallWithRetries(f, 1.0, sleep));
+ TF_EXPECT_OK(RetryingUtils::CallWithRetries(
+ f, sleep, RetryConfig(1L /* init_delay_time_us */)));
}
TEST(RetryingUtilsTest, CallWithRetries_EventualSuccess) {
@@ -86,7 +90,8 @@
results.erase(results.begin());
return result;
};
- TF_EXPECT_OK(RetryingUtils::CallWithRetries(f, 0));
+ TF_EXPECT_OK(RetryingUtils::CallWithRetries(
+ f, RetryConfig(0 /* init_delay_time_us */)));
}
TEST(RetryingUtilsTest, DeleteWithRetries_ImmediateSuccess) {
@@ -96,7 +101,8 @@
delete_results.erase(delete_results.begin());
return result;
};
- TF_EXPECT_OK(RetryingUtils::DeleteWithRetries(delete_func, 0));
+ TF_EXPECT_OK(RetryingUtils::DeleteWithRetries(
+ delete_func, RetryConfig(0 /* init_delay_time_us */)));
}
TEST(RetryingUtilsTest, DeleteWithRetries_EventualSuccess) {
@@ -106,7 +112,8 @@
delete_results.erase(delete_results.begin());
return result;
};
- TF_EXPECT_OK(RetryingUtils::DeleteWithRetries(delete_func, 0));
+ TF_EXPECT_OK(RetryingUtils::DeleteWithRetries(
+ delete_func, RetryConfig(0 /* init_delay_time_us */)));
}
TEST(RetryingUtilsTest, DeleteWithRetries_PermissionDeniedNotRetried) {
@@ -118,7 +125,9 @@
return result;
};
EXPECT_EQ(errors::Code::PERMISSION_DENIED,
- RetryingUtils::DeleteWithRetries(delete_func, 0).code());
+ RetryingUtils::DeleteWithRetries(
+ delete_func, RetryConfig(0 /* init_delay_time_us */))
+ .code());
}
TEST(RetryingUtilsTest, DeleteWithRetries_SuccessThroughFileNotFound) {
@@ -129,7 +138,8 @@
delete_results.erase(delete_results.begin());
return result;
};
- TF_EXPECT_OK(RetryingUtils::DeleteWithRetries(delete_func, 0));
+ TF_EXPECT_OK(RetryingUtils::DeleteWithRetries(
+ delete_func, RetryConfig(0 /* init_delay_time_us */)));
}
TEST(RetryingUtilsTest, DeleteWithRetries_FirstNotFoundReturnedAsIs) {
@@ -140,7 +150,9 @@
return result;
};
EXPECT_EQ(error::NOT_FOUND,
- RetryingUtils::DeleteWithRetries(delete_func, 0).code());
+ RetryingUtils::DeleteWithRetries(
+ delete_func, RetryConfig(0 /* init_delay_time_us */))
+ .code());
}
} // namespace
diff --git a/tensorflow/core/platform/default/build_config.bzl b/tensorflow/core/platform/default/build_config.bzl
index bb841ae..d884c1a 100644
--- a/tensorflow/core/platform/default/build_config.bzl
+++ b/tensorflow/core/platform/default/build_config.bzl
@@ -615,11 +615,7 @@
def tf_additional_lib_defines():
"""Additional defines needed to build TF libraries."""
- return select({
- "//tensorflow:with_jemalloc_linux_x86_64": ["TENSORFLOW_USE_JEMALLOC"],
- "//tensorflow:with_jemalloc_linux_ppc64le": ["TENSORFLOW_USE_JEMALLOC"],
- "//conditions:default": [],
- })
+ return []
def tf_additional_lib_deps():
"""Additional dependencies needed to build TF libraries."""
@@ -631,64 +627,45 @@
] + if_static(
["@nsync//:nsync_cpp"],
["@nsync//:nsync_headers"],
- ) + select({
- "//tensorflow:with_jemalloc_linux_x86_64_dynamic": ["@jemalloc//:jemalloc_headers"],
- "//tensorflow:with_jemalloc_linux_ppc64le_dynamic": ["@jemalloc//:jemalloc_headers"],
- "//tensorflow:with_jemalloc_linux_x86_64": ["@jemalloc//:jemalloc_impl"],
- "//tensorflow:with_jemalloc_linux_ppc64le": ["@jemalloc//:jemalloc_impl"],
- "//conditions:default": [],
- })
+ )
def tf_additional_core_deps():
return select({
- "//tensorflow:with_gcp_support_android_override": [],
- "//tensorflow:with_gcp_support_ios_override": [],
- "//tensorflow:with_gcp_support": [
+ "//tensorflow:android": [],
+ "//tensorflow:windows": [],
+ "//tensorflow:ios": [],
+ "//tensorflow:linux_s390x": [],
+ "//conditions:default": [
"//tensorflow/core/platform/cloud:gcs_file_system",
- ],
- "//conditions:default": [],
- }) + select({
- "//tensorflow:with_hdfs_support_windows_override": [],
- "//tensorflow:with_hdfs_support_android_override": [],
- "//tensorflow:with_hdfs_support_ios_override": [],
- "//tensorflow:with_hdfs_support": [
+ "//tensorflow/core/platform/s3:s3_file_system",
"//tensorflow/core/platform/hadoop:hadoop_file_system",
],
- "//conditions:default": [],
- }) + select({
- "//tensorflow:with_aws_support_windows_override": [],
- "//tensorflow:with_aws_support_android_override": [],
- "//tensorflow:with_aws_support_ios_override": [],
- "//tensorflow:with_aws_support": [
- "//tensorflow/core/platform/s3:s3_file_system",
- ],
- "//conditions:default": [],
})
# TODO(jart, jhseu): Delete when GCP is default on.
def tf_additional_cloud_op_deps():
return select({
- "//tensorflow:with_gcp_support_windows_override": [],
- "//tensorflow:with_gcp_support_android_override": [],
- "//tensorflow:with_gcp_support_ios_override": [],
- "//tensorflow:with_gcp_support": [
+ "//tensorflow:android": [],
+ "//tensorflow:windows": [],
+ "//tensorflow:ios": [],
+ "//tensorflow:linux_s390x": [],
+ "//conditions:default": [
"//tensorflow/contrib/cloud:bigquery_reader_ops_op_lib",
"//tensorflow/contrib/cloud:gcs_config_ops_op_lib",
],
- "//conditions:default": [],
})
# TODO(jart, jhseu): Delete when GCP is default on.
def tf_additional_cloud_kernel_deps():
return select({
- "//tensorflow:with_gcp_support_windows_override": [],
- "//tensorflow:with_gcp_support_android_override": [],
- "//tensorflow:with_gcp_support_ios_override": [],
- "//tensorflow:with_gcp_support": [
+ "//tensorflow:android": [],
+ "//tensorflow:windows": [],
+ "//tensorflow:ios": [],
+ "//tensorflow:linux_s390x": [],
+ "//conditions:default": [
"//tensorflow/contrib/cloud/kernels:bigquery_reader_ops",
"//tensorflow/contrib/cloud/kernels:gcs_config_ops",
],
- "//conditions:default": [],
})
def tf_lib_proto_parsing_deps():
@@ -738,11 +715,7 @@
"//tensorflow/stream_executor:cuda_platform",
"//tensorflow/core/platform/default/build_config:cuda",
],
- ) + select({
- "//tensorflow:with_jemalloc_linux_x86_64": ["@jemalloc//:jemalloc_impl"],
- "//tensorflow:with_jemalloc_linux_ppc64le": ["@jemalloc//:jemalloc_impl"],
- "//conditions:default": [],
- }) + [
+ ) + [
# TODO(allenl): Split these out into their own shared objects (they are
# here because they are shared between contrib/ op shared objects and
# core).
diff --git a/tensorflow/core/platform/posix/port.cc b/tensorflow/core/platform/posix/port.cc
index b46b992..acdd779 100644
--- a/tensorflow/core/platform/posix/port.cc
+++ b/tensorflow/core/platform/posix/port.cc
@@ -13,10 +13,6 @@
limitations under the License.
==============================================================================*/
-#ifdef TENSORFLOW_USE_JEMALLOC
-#include "jemalloc/jemalloc.h"
-#endif
-
#include "absl/base/internal/sysinfo.h"
#include "tensorflow/core/platform/cpu_info.h"
@@ -101,11 +97,7 @@
// memory aligned to at least the size of a pointer.
const int required_alignment = sizeof(void*);
if (minimum_alignment < required_alignment) return Malloc(size);
-#ifdef TENSORFLOW_USE_JEMALLOC
- int err = jemalloc_posix_memalign(&ptr, minimum_alignment, size);
-#else
int err = posix_memalign(&ptr, minimum_alignment, size);
-#endif
if (err != 0) {
return nullptr;
} else {
@@ -116,29 +108,11 @@
void AlignedFree(void* aligned_memory) { Free(aligned_memory); }
-void* Malloc(size_t size) {
-#ifdef TENSORFLOW_USE_JEMALLOC
- return jemalloc_malloc(size);
-#else
- return malloc(size);
-#endif
-}
+void* Malloc(size_t size) { return malloc(size); }
-void* Realloc(void* ptr, size_t size) {
-#ifdef TENSORFLOW_USE_JEMALLOC
- return jemalloc_realloc(ptr, size);
-#else
- return realloc(ptr, size);
-#endif
-}
+void* Realloc(void* ptr, size_t size) { return realloc(ptr, size); }
-void Free(void* ptr) {
-#ifdef TENSORFLOW_USE_JEMALLOC
- jemalloc_free(ptr);
-#else
- free(ptr);
-#endif
-}
+void Free(void* ptr) { free(ptr); }
void* NUMAMalloc(int node, size_t size, int minimum_alignment) {
return AlignedMalloc(size, minimum_alignment);
@@ -146,9 +120,7 @@
void NUMAFree(void* ptr, size_t size) { Free(ptr); }
-int NUMAGetMemAffinity(const void* addr) {
- return kNUMANoAffinity;
-}
+int NUMAGetMemAffinity(const void* addr) { return kNUMANoAffinity; }
void MallocExtension_ReleaseToSystem(std::size_t num_bytes) {
// No-op.
diff --git a/tensorflow/core/platform/windows/port.cc b/tensorflow/core/platform/windows/port.cc
index 5375f56..911ea19 100644
--- a/tensorflow/core/platform/windows/port.cc
+++ b/tensorflow/core/platform/windows/port.cc
@@ -13,10 +13,6 @@
limitations under the License.
==============================================================================*/
-#ifdef TENSORFLOW_USE_JEMALLOC
-#include "jemalloc/jemalloc.h"
-#endif
-
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
@@ -70,55 +66,16 @@
int NUMAGetThreadNodeAffinity() { return kNUMANoAffinity; }
void* AlignedMalloc(size_t size, int minimum_alignment) {
-#ifdef TENSORFLOW_USE_JEMALLOC
- void* ptr = NULL;
- // posix_memalign requires that the requested alignment be at least
- // sizeof(void*). In this case, fall back on malloc which should return
- // memory aligned to at least the size of a pointer.
- const int required_alignment = sizeof(void*);
- if (minimum_alignment < required_alignment) return Malloc(size);
- int err = jemalloc_posix_memalign(&ptr, minimum_alignment, size);
- if (err != 0) {
- return NULL;
- } else {
- return ptr;
- }
-#else
return _aligned_malloc(size, minimum_alignment);
-#endif
}
-void AlignedFree(void* aligned_memory) {
-#ifdef TENSORFLOW_USE_JEMALLOC
- jemalloc_free(aligned_memory);
-#else
- _aligned_free(aligned_memory);
-#endif
-}
+void AlignedFree(void* aligned_memory) { _aligned_free(aligned_memory); }
-void* Malloc(size_t size) {
-#ifdef TENSORFLOW_USE_JEMALLOC
- return jemalloc_malloc(size);
-#else
- return malloc(size);
-#endif
-}
+void* Malloc(size_t size) { return malloc(size); }
-void* Realloc(void* ptr, size_t size) {
-#ifdef TENSORFLOW_USE_JEMALLOC
- return jemalloc_realloc(ptr, size);
-#else
- return realloc(ptr, size);
-#endif
-}
+void* Realloc(void* ptr, size_t size) { return realloc(ptr, size); }
-void Free(void* ptr) {
-#ifdef TENSORFLOW_USE_JEMALLOC
- return jemalloc_free(ptr);
-#else
- return free(ptr);
-#endif
-}
+void Free(void* ptr) { return free(ptr); }
void* NUMAMalloc(int node, size_t size, int minimum_alignment) {
return AlignedMalloc(size, minimum_alignment);
diff --git a/tensorflow/core/profiler/BUILD b/tensorflow/core/profiler/BUILD
index af034bd..2bf3712 100644
--- a/tensorflow/core/profiler/BUILD
+++ b/tensorflow/core/profiler/BUILD
@@ -40,7 +40,6 @@
name = "protos_all",
srcs = glob(["**/*.proto"]),
cc_api_version = 2,
- java_api_version = 2,
protodeps = tf_additional_all_protos(),
visibility = ["//visibility:public"],
)
diff --git a/tensorflow/core/protobuf/config.proto b/tensorflow/core/protobuf/config.proto
index 85cd023..104ab03 100644
--- a/tensorflow/core/protobuf/config.proto
+++ b/tensorflow/core/protobuf/config.proto
@@ -453,6 +453,11 @@
// same group_key value (in a distributed computation where tasks
// run disjoint graphs).
int64 collective_graph_key = 1;
+ // If true, then operations (using the inter-op pool) across all
+ // session::run() calls will be centrally scheduled, optimizing for (median
+ // and tail) latency.
+ // Consider using this option for CPU-bound workloads like inference.
+ bool use_run_handler_pool = 2;
};
Experimental experimental = 8;
diff --git a/tensorflow/core/protobuf/rewriter_config.proto b/tensorflow/core/protobuf/rewriter_config.proto
index bb8f883..8c31468 100644
--- a/tensorflow/core/protobuf/rewriter_config.proto
+++ b/tensorflow/core/protobuf/rewriter_config.proto
@@ -75,8 +75,10 @@
// Try to allocate some independent Op outputs contiguously in order to
// merge or eliminate downstream Ops (off by default).
Toggle scoped_allocator_optimization = 15;
- // Force small ops onto the CPU (default is ON).
+ // Force small ops onto the CPU (default is OFF).
Toggle pin_to_host_optimization = 18;
+ // Disable the entire meta optimizer (off by default).
+ bool disable_meta_optimizer = 19;
// Controls how many times we run the optimizers in meta optimizer (default
// is once).
@@ -143,8 +145,8 @@
// not configurable (in contrast to memory optimization passes through the
// meta-optimizer) and act only on manual op annotations.
//
- // Custom registered optimizers will be run after the base optimizers, in
- // the order that they are specified.
+ // Custom optimizers (see custom_optimizers) that are not part of this
+ // schedule will be run after - in the order that they were specified.
repeated string optimizers = 100;
// Message to describe custom graph optimizer and its parameters
diff --git a/tensorflow/core/util/mkl_util.h b/tensorflow/core/util/mkl_util.h
index cf7ffd8..04aaea4 100644
--- a/tensorflow/core/util/mkl_util.h
+++ b/tensorflow/core/util/mkl_util.h
@@ -2039,8 +2039,8 @@
/// Fuction to check whether primitive memory optimization is enabled
static inline bool IsPrimitiveMemOptEnabled() {
bool is_primitive_mem_opt_enabled = true;
- TF_CHECK_OK(ReadBoolFromEnvVar("TF_MKL_OPTIMIZE_PRIMITVE_MEMUSE", true,
- &is_primitive_mem_opt_enabled));
+ TF_CHECK_OK(ReadBoolFromEnvVar("TF_MKL_OPTIMIZE_PRIMITIVE_MEMUSE", true,
+ &is_primitive_mem_opt_enabled));
return is_primitive_mem_opt_enabled;
}
@@ -2095,9 +2095,8 @@
fmt_desired = is_2d ? memory::format::nChw16c : memory::format::nCdhw16c;
} else if (port::TestCPUFeature(port::CPUFeature::AVX2) &&
(channel % 8) == 0) {
- fmt_desired = is_2d
- ? memory::format::nChw8c
- : memory::format::ncdhw; //not support avx2 for 3d yet.
+ fmt_desired = is_2d ? memory::format::nChw8c
+ : memory::format::ncdhw; // no avx2 support for 3d yet.
} else {
fmt_desired = is_2d ? memory::format::nchw : memory::format::ncdhw;
}
@@ -2209,7 +2208,8 @@
// utility function to determine if it is conv 1x1 and stride != 1
// for purpose of temporarily disabling primitive reuse
-inline bool IsConv1x1StrideNot1(memory::dims filter_dims, memory::dims strides) {
+inline bool IsConv1x1StrideNot1(memory::dims filter_dims,
+ memory::dims strides) {
if (filter_dims.size() != 4 || strides.size() != 2) return false;
return ((filter_dims[2] == 1) && (filter_dims[3] == 1) &&
diff --git a/tensorflow/core/util/tensor_bundle/BUILD b/tensorflow/core/util/tensor_bundle/BUILD
index 6483586..f40ec9b 100644
--- a/tensorflow/core/util/tensor_bundle/BUILD
+++ b/tensorflow/core/util/tensor_bundle/BUILD
@@ -64,6 +64,11 @@
tf_cc_test(
name = "tensor_bundle_test",
srcs = ["tensor_bundle_test.cc"],
+ data = glob(["testdata/**"]),
+ tags = [
+ "nomsan",
+ "notsan",
+ ],
deps = [
":tensor_bundle",
"//tensorflow/core:framework",
diff --git a/tensorflow/core/util/tensor_bundle/tensor_bundle.cc b/tensorflow/core/util/tensor_bundle/tensor_bundle.cc
index ea8a259..2dcb57a 100644
--- a/tensorflow/core/util/tensor_bundle/tensor_bundle.cc
+++ b/tensorflow/core/util/tensor_bundle/tensor_bundle.cc
@@ -64,27 +64,36 @@
// Reads "num_elements" string elements from file[offset, offset+size) into the
// length-N "destination". Discards the original content of "destination".
//
-// Checksums the string lengths (as restored uint32, not varint32 bytes) and
-// string bytes, and stores it into "actual_crc32c".
+// Checksums the string lengths (as restored uint32 or uint64, not varint64
+// bytes) and string bytes, and stores it into "actual_crc32c".
Status ReadStringTensor(io::InputBuffer* buffered_file, size_t num_elements,
size_t offset, size_t size, string* destination,
uint32* actual_crc32c) {
if (size == 0) return Status::OK();
CHECK_GT(size, 0);
- // Reads "num_elements" varint32's from "buffered_file".
+ // Reads "num_elements" varint64's from "buffered_file".
TF_RETURN_IF_ERROR(buffered_file->Seek(offset));
- std::vector<uint32> string_lengths(num_elements);
+ std::vector<uint64> string_lengths(num_elements);
for (size_t i = 0; i < num_elements; ++i) {
- TF_RETURN_IF_ERROR(buffered_file->ReadVarint32(&string_lengths[i]));
+ TF_RETURN_IF_ERROR(buffered_file->ReadVarint64(&string_lengths[i]));
+ if (string_lengths[i] <= UINT32_MAX) {
+ // We need to do this because older checkpoints only used uint32s and we
+ // should still support them.
+ const uint32 elem_size_uint32 = static_cast<uint32>(string_lengths[i]);
+ *actual_crc32c = crc32c::Extend(
+ *actual_crc32c, reinterpret_cast<const char*>(&elem_size_uint32),
+ sizeof(uint32));
+ } else {
+ *actual_crc32c = crc32c::Extend(
+ *actual_crc32c, reinterpret_cast<const char*>(&string_lengths[i]),
+ sizeof(uint64));
+ }
}
if (offset + size < buffered_file->Tell()) {
return errors::DataLoss("String lengths longer than expected offset ",
offset + size);
}
- *actual_crc32c =
- crc32c::Value(reinterpret_cast<const char*>(string_lengths.data()),
- sizeof(uint32) * num_elements);
// Reads the length-checksum.
uint32 length_checksum = 0;
@@ -104,7 +113,7 @@
// Reads the actual string bytes.
for (size_t i = 0; i < num_elements; ++i) {
- const uint32 string_length = string_lengths[i];
+ const uint64 string_length = string_lengths[i];
string* buffer = &destination[i];
buffer->resize(string_length);
@@ -218,8 +227,8 @@
Status WriteStringTensor(const Tensor& val, FileOutputBuffer* out,
size_t* bytes_written, uint32* crc32c) {
// On-disk format:
- // [varint32 len0]..[varint32 lenL][4 byte cksum on lengths][string bytes]
- // Var "crc32c" checksums the string lengths (as uint32, not varint32 bytes),
+ // [varint64 len0]..[varint64 lenL][4 byte cksum on lengths][string bytes]
+ // Var "crc32c" checksums the string lengths (as uint64, not varint64 bytes),
// the length-checksum, and all the string bytes.
DCHECK_EQ(val.dtype(), DT_STRING);
const string* strings = GetStringBackingBuffer(val);
@@ -230,12 +239,21 @@
*crc32c = 0;
for (int64 i = 0; i < val.NumElements(); ++i) {
const string* elem = &strings[i];
- DCHECK_EQ(elem->size(), static_cast<uint32>(elem->size()));
- const uint32 elem_size = static_cast<uint32>(elem->size());
+ DCHECK_EQ(elem->size(), static_cast<uint64>(elem->size()));
+ const uint64 elem_size = static_cast<uint64>(elem->size());
- core::PutVarint32(&lengths, elem_size);
- *crc32c = crc32c::Extend(*crc32c, reinterpret_cast<const char*>(&elem_size),
- sizeof(uint32));
+ core::PutVarint64(&lengths, elem_size);
+ if (elem_size <= UINT32_MAX) {
+ // We need to do this because older checkpoints only used uint32s and we
+ // should still support them.
+ const uint32 elem_size_uint32 = static_cast<uint32>(elem_size);
+ *crc32c = crc32c::Extend(*crc32c,
+ reinterpret_cast<const char*>(&elem_size_uint32),
+ sizeof(uint32));
+ } else {
+ *crc32c = crc32c::Extend(
+ *crc32c, reinterpret_cast<const char*>(&elem_size), sizeof(uint64));
+ }
}
TF_RETURN_IF_ERROR(out->Append(lengths));
*bytes_written = lengths.size();
diff --git a/tensorflow/core/util/tensor_bundle/tensor_bundle_test.cc b/tensorflow/core/util/tensor_bundle/tensor_bundle_test.cc
index 59c42ba..9567e47 100644
--- a/tensorflow/core/util/tensor_bundle/tensor_bundle_test.cc
+++ b/tensorflow/core/util/tensor_bundle/tensor_bundle_test.cc
@@ -39,6 +39,11 @@
return strings::StrCat(testing::TmpDir(), "/", prefix);
}
+string TestdataPrefix(const string& prefix) {
+ return strings::StrCat(testing::TensorFlowSrcRoot(),
+ "/core/util/tensor_bundle/testdata/", prefix);
+}
+
template <typename T>
Tensor Constant(T v, TensorShape shape) {
Tensor ret(DataTypeToEnum<T>::value, shape);
@@ -458,7 +463,26 @@
TestNonStandardShapes<qint8>();
}
+TEST(TensorBundleTest, StringTensorsOldFormat) {
+ // Test string tensor bundle made with previous version of code that use
+ // varint32s to store string lengths (we now use varint64s).
+ BundleReader reader(Env::Default(), TestdataPrefix("old_string_tensors/foo"));
+ TF_ASSERT_OK(reader.status());
+ EXPECT_EQ(AllTensorKeys(&reader),
+ std::vector<string>({"floats", "scalar", "string_tensor", "strs"}));
+
+ Expect<string>(&reader, "string_tensor", Tensor(DT_STRING, TensorShape({1})));
+ Expect<string>(&reader, "scalar", test::AsTensor<string>({"hello"}));
+ Expect<string>(
+ &reader, "strs",
+ test::AsTensor<string>({"hello", "", "x01", string(1 << 10, 'c')}));
+ Expect<float>(&reader, "floats", Constant_2x3<float>(16.18));
+}
+
TEST(TensorBundleTest, StringTensors) {
+ constexpr size_t kLongLength = static_cast<size_t>(UINT32_MAX) + 1;
+ Tensor long_string_tensor(DT_STRING, TensorShape({1}));
+
{
BundleWriter writer(Env::Default(), Prefix("foo"));
TF_EXPECT_OK(writer.Add("string_tensor",
@@ -467,6 +491,12 @@
TF_EXPECT_OK(writer.Add(
"strs",
test::AsTensor<string>({"hello", "", "x01", string(1 << 25, 'c')})));
+
+ // Requires a 64-bit length.
+ string* backing_string = long_string_tensor.flat<string>().data();
+ backing_string->assign(kLongLength, 'd');
+ TF_EXPECT_OK(writer.Add("long_scalar", long_string_tensor));
+
// Mixes in some floats.
TF_EXPECT_OK(writer.Add("floats", Constant_2x3<float>(16.18)));
TF_ASSERT_OK(writer.Finish());
@@ -474,9 +504,9 @@
{
BundleReader reader(Env::Default(), Prefix("foo"));
TF_ASSERT_OK(reader.status());
- EXPECT_EQ(
- AllTensorKeys(&reader),
- std::vector<string>({"floats", "scalar", "string_tensor", "strs"}));
+ EXPECT_EQ(AllTensorKeys(&reader),
+ std::vector<string>({"floats", "long_scalar", "scalar",
+ "string_tensor", "strs"}));
Expect<string>(&reader, "string_tensor",
Tensor(DT_STRING, TensorShape({1})));
@@ -484,7 +514,35 @@
Expect<string>(
&reader, "strs",
test::AsTensor<string>({"hello", "", "x01", string(1 << 25, 'c')}));
+
Expect<float>(&reader, "floats", Constant_2x3<float>(16.18));
+
+ // We don't use the Expect function so we can re-use the
+ // `long_string_tensor` buffer for reading out long_scalar to keep memory
+ // usage reasonable.
+ EXPECT_TRUE(reader.Contains("long_scalar"));
+ DataType dtype;
+ TensorShape shape;
+ TF_ASSERT_OK(reader.LookupDtypeAndShape("long_scalar", &dtype, &shape));
+ EXPECT_EQ(DT_STRING, dtype);
+ EXPECT_EQ(TensorShape({1}), shape);
+
+ // Zero-out the string so that we can be sure the new one is read in.
+ string* backing_string = long_string_tensor.flat<string>().data();
+ backing_string->assign("");
+
+ // Read long_scalar and check it contains kLongLength 'd's.
+ TF_ASSERT_OK(reader.Lookup("long_scalar", &long_string_tensor));
+ ASSERT_EQ(backing_string, long_string_tensor.flat<string>().data());
+ EXPECT_EQ(kLongLength, backing_string->length());
+ for (char c : *backing_string) {
+ // Not using ASSERT_EQ('d', c) because this way is twice as fast due to
+ // compiler optimizations.
+ if (c != 'd') {
+ FAIL() << "long_scalar is not full of 'd's as expected.";
+ break;
+ }
+ }
}
}
diff --git a/tensorflow/core/util/tensor_bundle/testdata/old_string_tensors/README b/tensorflow/core/util/tensor_bundle/testdata/old_string_tensors/README
new file mode 100644
index 0000000..428d3ef
--- /dev/null
+++ b/tensorflow/core/util/tensor_bundle/testdata/old_string_tensors/README
@@ -0,0 +1,3 @@
+This tensor bundle was generated from cl/214343133, before string tensor
+lengths were written as varint64s. This is here to check backwards
+compatibility between the new code and old checkpoints.
diff --git a/tensorflow/core/util/tensor_bundle/testdata/old_string_tensors/foo.data-00000-of-00001 b/tensorflow/core/util/tensor_bundle/testdata/old_string_tensors/foo.data-00000-of-00001
new file mode 100644
index 0000000..23b488e
--- /dev/null
+++ b/tensorflow/core/util/tensor_bundle/testdata/old_string_tensors/foo.data-00000-of-00001
Binary files differ
diff --git a/tensorflow/core/util/tensor_bundle/testdata/old_string_tensors/foo.index b/tensorflow/core/util/tensor_bundle/testdata/old_string_tensors/foo.index
new file mode 100644
index 0000000..a22a69e
--- /dev/null
+++ b/tensorflow/core/util/tensor_bundle/testdata/old_string_tensors/foo.index
Binary files differ
diff --git a/tensorflow/docs_src/BUILD b/tensorflow/docs_src/BUILD
deleted file mode 100644
index 34bf7b6..0000000
--- a/tensorflow/docs_src/BUILD
+++ /dev/null
@@ -1,14 +0,0 @@
-# Files used to generate TensorFlow docs.
-
-licenses(["notice"]) # Apache 2.0
-
-package(
- default_visibility = ["//tensorflow:internal"],
-)
-
-exports_files(["LICENSE"])
-
-filegroup(
- name = "docs_src",
- data = glob(["**/*.md"]),
-)
diff --git a/tensorflow/docs_src/performance/xla/operation_semantics.md b/tensorflow/docs_src/performance/xla/operation_semantics.md
new file mode 100644
index 0000000..96d269b
--- /dev/null
+++ b/tensorflow/docs_src/performance/xla/operation_semantics.md
@@ -0,0 +1,2426 @@
+# Operation Semantics
+
+The following describes the semantics of operations defined in the
+[`XlaBuilder`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h)
+interface. Typically, these operations map one-to-one to operations defined in
+the RPC interface in
+[`xla_data.proto`](https://www.tensorflow.org/code/tensorflow/compiler/xla/xla_data.proto).
+
+A note on nomenclature: the generalized data type XLA deals with is an
+N-dimensional array holding elements of some uniform type (such as 32-bit
+float). Throughout the documentation, *array* is used to denote an
+arbitrary-dimensional array. For convenience, special cases have more specific
+and familiar names; for example a *vector* is a 1-dimensional array and a
+*matrix* is a 2-dimensional array.
+
+## AllToAll
+
+See also
+[`XlaBuilder::AllToAll`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+Alltoall is a collective operation that sends data from all cores to all cores.
+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.
+2. the gather phase. Each core concatenates the received blocks along the
+ `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.
+
+Prerequisites:
+
+- The dimension size of the operand on the split_dimension is divisible by
+ split_count.
+- The operand's shape is not tuple.
+
+<b> `AllToAll(operand, split_dimension, concat_dimension, split_count,
+replica_groups)` </b>
+
+
+| Arguments | Type | Semantics |
+| ------------------ | --------------------- | ------------------------------- |
+| `operand` | `XlaOp` | n dimensional input array |
+| `split_dimension` | `int64` | A value in the interval `[0, |
+: : : n)` that names the dimension :
+: : : along which the operand is :
+: : : split :
+| `concat_dimension` | `int64` | a value in the interval `[0, |
+: : : n)` that names the dimension :
+: : : along which the split blocks :
+: : : are concatenated :
+| `split_count` | `int64` | the number of cores that |
+: : : participate this operation. If :
+: : : `replica_groups` is empty, this :
+: : : should be the number of :
+: : : replicas; otherwise, this :
+: : : should be equal to the number :
+: : : of replicas in each group. :
+| `replica_groups` | `ReplicaGroup` vector | each group contains a list of |
+: : : replica id. :
+
+Below shows an example of Alltoall.
+
+```
+XlaBuilder b("alltoall");
+auto x = Parameter(&b, 0, ShapeUtil::MakeShape(F32, {4, 16}), "x");
+AllToAll(x, /*split_dimension=*/1, /*concat_dimension=*/0, /*split_count=*/4);
+```
+
+<div style="width:95%; margin:auto; margin-bottom:10px; margin-top:20px;">
+ <img style="width:100%" src="../../images/xla/ops_alltoall.png">
+</div>
+
+In this example, there are 4 cores participating the Alltoall. On each core, the
+operand is split into 4 parts along dimension 0, so each part has shape
+f32[4,4]. The 4 parts are scattered to all cores. Then each core concatenates
+the received parts along dimension 1, in the order or core 0-4. So the output on
+each core has shape f32[16,4].
+
+## BatchNormGrad
+
+See also
+[`XlaBuilder::BatchNormGrad`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h)
+and [the original batch normalization paper](https://arxiv.org/abs/1502.03167)
+for a detailed description of the algorithm.
+
+Calculates gradients of batch norm.
+
+<b> `BatchNormGrad(operand, scale, mean, variance, grad_output, epsilon, feature_index)` </b>
+
+| Arguments | Type | Semantics |
+| --------------- | ----------------------- | -------------------------------- |
+| `operand` | `XlaOp` | n dimensional array to be |
+: : : normalized (x) :
+| `scale` | `XlaOp` | 1 dimensional array |
+: : : (\\(\gamma\\)) :
+| `mean` | `XlaOp` | 1 dimensional array (\\(\mu\\)) |
+| `variance` | `XlaOp` | 1 dimensional array |
+: : : (\\(\sigma^2\\)) :
+| `grad_output` | `XlaOp` | Gradients passed to |
+: : : `BatchNormTraining` :
+: : : (\\( \nabla y\\)) :
+| `epsilon` | `float` | Epsilon value (\\(\epsilon\\)) |
+| `feature_index` | `int64` | Index to feature dimension in |
+: : : `operand` :
+
+For each feature in the feature dimension (`feature_index` is the index for the
+feature dimension in `operand`), the operation calculates the gradients with
+respect to `operand`, `offset` and `scale` across all the other dimensions. The
+`feature_index` must be a valid index for the feature dimension in `operand`.
+
+The three gradients are defined by the following formulas (assuming a
+4-dimensional tensor as `operand` and with feature dimension index \\(l\\),
+batch size `m` and spatial sizes `w` and `h`):
+
+\\[ \begin{split} c_l&=
+\frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h
+\left( \nabla y_{ijkl} \frac{x_{ijkl} - \mu_l}{\sigma^2_l+\epsilon} \right)
+\\\\
+\nabla x_{ijkl} &= \frac{\gamma_{l}}{\sqrt{\sigma^2_{l}+\epsilon}}
+\left( \nabla y_{ijkl} - \mathrm{mean}(\nabla y) - c_l (x_{ijkl} - \mu_{l})
+\right)
+\\\\
+\nabla \gamma_l &= \sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h \left( \nabla y_{ijkl}
+\frac{x_{ijkl} - \mu_l}{\sqrt{\sigma^2_{l}+\epsilon}} \right)
+\\\\\
+\nabla \beta_l &= \sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h \nabla y_{ijkl}
+\end{split} \\]
+
+The inputs `mean` and `variance` represent moments value
+across batch and spatial dimensions.
+
+The output type is a tuple of three handles:
+
+| Outputs | Type | Semantics |
+| ------------- | ----------------------- | --------------------------------- |
+| `grad_operand` | `XlaOp` | gradient with respect to input |
+: : : `operand` (\\( \nabla x\\)) :
+| `grad_scale` | `XlaOp` | gradient with respect to input |
+: : : `scale` (\\( \nabla \gamma\\)) :
+| `grad_offset` | `XlaOp` | gradient with respect to input |
+: : : `offset`(\\( \nabla \beta\\)) :
+
+## BatchNormInference
+
+See also
+[`XlaBuilder::BatchNormInference`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h)
+and [the original batch normalization paper](https://arxiv.org/abs/1502.03167)
+for a detailed description of the algorithm.
+
+Normalizes an array across batch and spatial dimensions.
+
+<b> `BatchNormInference(operand, scale, offset, mean, variance, epsilon, feature_index)` </b>
+
+Arguments | Type | Semantics
+--------------- | ------- | ---------------------------------------
+`operand` | `XlaOp` | n dimensional array to be normalized
+`scale` | `XlaOp` | 1 dimensional array
+`offset` | `XlaOp` | 1 dimensional array
+`mean` | `XlaOp` | 1 dimensional array
+`variance` | `XlaOp` | 1 dimensional array
+`epsilon` | `float` | Epsilon value
+`feature_index` | `int64` | Index to feature dimension in `operand`
+
+For each feature in the feature dimension (`feature_index` is the index for the
+feature dimension in `operand`), the operation calculates the mean and variance
+across all the other dimensions and uses the mean and variance to normalize each
+element in `operand`. The `feature_index` must be a valid index for the feature
+dimension in `operand`.
+
+`BatchNormInference` is equivalent to calling `BatchNormTraining` without
+computing `mean` and `variance` for each batch. It uses the input `mean` and
+`variance` instead as estimated values. The purpose of this op is to reduce
+latency in inference, hence the name `BatchNormInference`.
+
+The output is an n-dimensional, normalized array with the same shape as input
+`operand`.
+
+## BatchNormTraining
+
+See also
+[`XlaBuilder::BatchNormTraining`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h)
+and [`the original batch normalization paper`](https://arxiv.org/abs/1502.03167)
+for a detailed description of the algorithm.
+
+Normalizes an array across batch and spatial dimensions.
+
+<b> `BatchNormTraining(operand, scale, offset, epsilon, feature_index)` </b>
+
+Arguments | Type | Semantics
+--------------- | ------- | ----------------------------------------
+`operand` | `XlaOp` | n dimensional array to be normalized (x)
+`scale` | `XlaOp` | 1 dimensional array (\\(\gamma\\))
+`offset` | `XlaOp` | 1 dimensional array (\\(\beta\\))
+`epsilon` | `float` | Epsilon value (\\(\epsilon\\))
+`feature_index` | `int64` | Index to feature dimension in `operand`
+
+For each feature in the feature dimension (`feature_index` is the index for the
+feature dimension in `operand`), the operation calculates the mean and variance
+across all the other dimensions and uses the mean and variance to normalize each
+element in `operand`. The `feature_index` must be a valid index for the feature
+dimension in `operand`.
+
+The algorithm goes as follows for each batch in `operand` \\(x\\) that
+contains `m` elements with `w` and `h` as the size of spatial dimensions
+(assuming `operand` is an 4 dimensional array):
+
+- Calculates batch mean \\(\mu_l\\) for each feature `l` in feature dimension:
+\\(\mu_l=\frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h x_{ijkl}\\)
+
+- Calculates batch variance \\(\sigma^2_l\\):
+\\(\sigma^2_l=\frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h (x_{ijkl} - \mu_l)^2\\)
+
+- Normalizes, scales and shifts:
+\\(y_{ijkl}=\frac{\gamma_l(x_{ijkl}-\mu_l)}{\sqrt[2]{\sigma^2_l+\epsilon}}+\beta_l\\)
+
+The epsilon value, usually a small number, is added to avoid divide-by-zero errors.
+
+The output type is a tuple of three `XlaOp`s:
+
+| Outputs | Type | Semantics |
+| ------------ | ----------------------- | -------------------------------------|
+| `output` | `XlaOp` | n dimensional array with the same |
+: : : shape as input `operand` (y) :
+| `batch_mean` | `XlaOp` | 1 dimensional array (\\(\mu\\)) |
+| `batch_var` | `XlaOp` | 1 dimensional array (\\(\sigma^2\\)) |
+
+The `batch_mean` and `batch_var` are moments calculated across the batch and
+spatial dimensions using the formulas above.
+
+## BitcastConvertType
+
+See also
+[`XlaBuilder::BitcastConvertType`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+Similar to a `tf.bitcast` in TensorFlow, performs an element-wise bitcast
+operation from a data shape to a target shape. The dimensions must match, and
+the conversion is an element-wise one; e.g. `s32` elements become `f32` elements
+via bitcast routine. Bitcast is implemented as a low-level cast, so machines
+with different floating-point representations will give different results.
+
+<b> `BitcastConvertType(operand, new_element_type)` </b>
+
+Arguments | Type | Semantics
+------------------ | --------------- | ---------------------------
+`operand` | `XlaOp` | array of type T with dims D
+`new_element_type` | `PrimitiveType` | type U
+
+The dimensions of the operand and the target shape must match. The bit-width of
+the source and destination element types must be equal. The source
+and destination element types must not be tuples.
+
+## Broadcast
+
+See also
+[`XlaBuilder::Broadcast`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+Adds dimensions to an array by duplicating the data in the array.
+
+<b> `Broadcast(operand, broadcast_sizes)` </b>
+
+Arguments | Type | Semantics
+----------------- | ------------------- | -------------------------------
+`operand` | `XlaOp` | The array to duplicate
+`broadcast_sizes` | `ArraySlice<int64>` | The sizes of the new dimensions
+
+The new dimensions are inserted on the left, i.e. if `broadcast_sizes` has
+values `{a0, ..., aN}` and the operand shape has dimensions `{b0, ..., bM}` then
+the shape of the output has dimensions `{a0, ..., aN, b0, ..., bM}`.
+
+The new dimensions index into copies of the operand, i.e.
+
+```
+output[i0, ..., iN, j0, ..., jM] = operand[j0, ..., jM]
+```
+
+For example, if `operand` is a scalar `f32` with value `2.0f`, and
+`broadcast_sizes` is `{2, 3}`, then the result will be an array with shape
+`f32[2, 3]` and all the values in the result will be `2.0f`.
+
+## Call
+
+See also
+[`XlaBuilder::Call`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+Invokes a computation with the given arguments.
+
+<b> `Call(computation, args...)` </b>
+
+| Arguments | Type | Semantics |
+| ------------- | ---------------------- | ----------------------------------- |
+| `computation` | `XlaComputation` | computation of type `T_0, T_1, ..., |
+: : : T_N -> S` with N parameters of :
+: : : arbitrary type :
+| `args` | sequence of N `XlaOp`s | N arguments of arbitrary type |
+
+The arity and types of the `args` must match the parameters of the
+`computation`. It is allowed to have no `args`.
+
+## Clamp
+
+See also
+[`XlaBuilder::Clamp`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+Clamps an operand to within the range between a minimum and maximum value.
+
+<b> `Clamp(min, operand, max)` </b>
+
+Arguments | Type | Semantics
+--------- | ------- | ---------------
+`min` | `XlaOp` | array of type T
+`operand` | `XlaOp` | array of type T
+`max` | `XlaOp` | array of type T
+
+Given an operand and minimum and maximum values, returns the operand if it is in
+the range between the minimum and maximum, else returns the minimum value if the
+operand is below this range or the maximum value if the operand is above this
+range. That is, `clamp(a, x, b) = min(max(a, x), b)`.
+
+All three arrays must be the same shape. Alternatively, as a restricted form of
+[broadcasting](broadcasting.md), `min` and/or `max` can be a scalar of type `T`.
+
+Example with scalar `min` and `max`:
+
+```
+let operand: s32[3] = {-1, 5, 9};
+let min: s32 = 0;
+let max: s32 = 6;
+==>
+Clamp(min, operand, max) = s32[3]{0, 5, 6};
+```
+
+## Collapse
+
+See also
+[`XlaBuilder::Collapse`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h)
+and the `tf.reshape` operation.
+
+Collapses dimensions of an array into one dimension.
+
+<b> `Collapse(operand, dimensions)` </b>
+
+Arguments | Type | Semantics
+------------ | -------------- | -----------------------------------------------
+`operand` | `XlaOp` | array of type T
+`dimensions` | `int64` vector | in-order, consecutive subset of T's dimensions.
+
+Collapse replaces the given subset of the operand's dimensions by a single
+dimension. The input arguments are an arbitrary array of type T and a
+compile-time-constant vector of dimension indices. The dimension indices must be
+an in-order (low to high dimension numbers), consecutive subset of T's
+dimensions. Thus, {0, 1, 2}, {0, 1}, or {1, 2} are all valid dimension sets, but
+{1, 0} or {0, 2} are not. They are replaced by a single new dimension, in the
+same position in the dimension sequence as those they replace, with the new
+dimension size equal to the product of original dimension sizes. The lowest
+dimension number in `dimensions` is the slowest varying dimension (most major)
+in the loop nest which collapses these dimension, and the highest dimension
+number is fastest varying (most minor). See the `tf.reshape` operator
+if more general collapse ordering is needed.
+
+For example, let v be an array of 24 elements:
+
+```
+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}}};
+
+// 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};
+
+// 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}};
+
+// 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}};
+
+```
+
+## Concatenate
+
+See also
+[`XlaBuilder::ConcatInDim`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+Concatenate composes an array from multiple array operands. The array is of the
+same rank as each of the input array operands (which must be of the same rank as
+each other) and contains the arguments in the order that they were specified.
+
+<b> `Concatenate(operands..., dimension)` </b>
+
+| Arguments | Type | Semantics |
+| ----------- | --------------------- | -------------------------------------- |
+| `operands` | sequence of N `XlaOp` | N arrays of type T with dimensions |
+: : : [L0, L1, ...]. Requires N >= 1. :
+| `dimension` | `int64` | A value in the interval `[0, N)` that |
+: : : names the dimension to be concatenated :
+: : : between the `operands`. :
+
+With the exception of `dimension` all dimensions must be the same. This is
+because XLA does not support "ragged" arrays. Also note that rank-0 values
+cannot be concatenated (as it's impossible to name the dimension along which the
+concatenation occurs).
+
+1-dimensional example:
+
+```
+Concat({{2, 3}, {4, 5}, {6, 7}}, 0)
+>>> {2, 3, 4, 5, 6, 7}
+```
+
+2-dimensional example:
+
+```
+let a = {
+ {1, 2},
+ {3, 4},
+ {5, 6},
+};
+let b = {
+ {7, 8},
+};
+Concat({a, b}, 0)
+>>> {
+ {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="https://www.tensorflow.org/images/ops_concatenate.png">
+</div>
+
+## Conditional
+
+See also
+[`XlaBuilder::Conditional`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+<b> `Conditional(pred, true_operand, true_computation, false_operand,
+false_computation)` </b>
+
+Arguments | Type | Semantics
+------------------- | ---------------- | ---------------------------------
+`pred` | `XlaOp` | Scalar of type `PRED`
+`true_operand` | `XlaOp` | Argument of type `T_0`
+`true_computation` | `XlaComputation` | XlaComputation of type `T_0 -> S`
+`false_operand` | `XlaOp` | Argument of type `T_1`
+`false_computation` | `XlaComputation` | XlaComputation of type `T_1 -> S`
+
+Executes `true_computation` if `pred` is `true`, `false_computation` if `pred`
+is `false`, and returns the result.
+
+The `true_computation` must take in a single argument of type `T_0` and will be
+invoked with `true_operand` which must be of the same type. The
+`false_computation` must take in a single argument of type `T_1` and will be
+invoked with `false_operand` which must be of the same type. The type of the
+returned value of `true_computation` and `false_computation` must be the same.
+
+Note that only one of `true_computation` and `false_computation` will be
+executed depending on the value of `pred`.
+
+## Conv (convolution)
+
+See also
+[`XlaBuilder::Conv`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+As ConvWithGeneralPadding, but the padding is specified in a short-hand way as
+either SAME or VALID. SAME padding pads the input (`lhs`) with zeroes so that
+the output has the same shape as the input when not taking striding into
+account. VALID padding simply means no padding.
+
+## ConvWithGeneralPadding (convolution)
+
+See also
+[`XlaBuilder::ConvWithGeneralPadding`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+Computes a convolution of the kind used in neural networks. Here, a convolution
+can be thought of as a n-dimensional window moving across a n-dimensional base
+area and a computation is performed for each possible position of the window.
+
+| Arguments | Type | Semantics |
+| --------------------- | -------------------- | ----------------------------- |
+| `lhs` | `XlaOp` | rank n+2 array of inputs |
+| `rhs` | `XlaOp` | rank n+2 array of kernel |
+: : : weights :
+| `window_strides` | `ArraySlice<int64>` | n-d array of kernel strides |
+| `padding` | `ArraySlice< | n-d array of (low, high) |
+: : pair<int64, int64>>` : padding :
+| `lhs_dilation` | `ArraySlice<int64>` | n-d lhs dilation factor array |
+| `rhs_dilation` | `ArraySlice<int64>` | n-d rhs dilation factor array |
+| `feature_group_count` | int64 | the number of feature groups |
+
+Let n be the number of spatial dimensions. The `lhs` argument is a rank n+2
+array describing the base area. This is called the input, even though of course
+the rhs is also an input. In a neural network, these are the input activations.
+The n+2 dimensions are, in this order:
+
+* `batch`: Each coordinate in this dimension represents an independent input
+ 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.
+* `spatial_dims`: Describes the `n` spatial dimensions that define the base
+ 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.
+* `spatial_dims`: Describes the `n` spatial dimensions that define the n-d
+ 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
+dimension is 3, then the window can only be placed at coordinates where the
+first spatial index is divisible by 3.
+
+The `padding` argument specifies the amount of zero padding to be applied to the
+base area. The amount of padding can be negative -- the absolute value of
+negative padding indicates the number of elements to remove from the specified
+dimension before doing the convolution. `padding[0]` specifies the padding for
+dimension `y` and `padding[1]` specifies the padding for dimension `x`. Each
+pair has the low padding as the first element and the high padding as the second
+element. The low padding is applied in the direction of lower indices while the
+high padding is applied in the direction of higher indices. For example, if
+`padding[1]` is `(2,3)` then there will be a padding by 2 zeroes on the left and
+by 3 zeroes on the right in the second spatial dimension. Using padding is
+equivalent to inserting those same zero values into the input (`lhs`) before
+doing the convolution.
+
+The `lhs_dilation` and `rhs_dilation` arguments specify the dilation factor to
+be applied to the lhs and rhs, respectively, in each spatial dimension. If the
+dilation factor in a spatial dimension is d, then d-1 holes are implicitly
+placed between each of the entries in that dimension, increasing the size of the
+array. The holes are filled with a no-op value, which for convolution means
+zeroes.
+
+Dilation of the rhs is also called atrous convolution. For more details, see
+`tf.nn.atrous_conv2d`. Dilation of the lhs is also called transposed
+convolution. For more details, see `tf.nn.conv2d_transpose`.
+
+The `feature_group_count` argument (default value 1) can be used for grouped
+convolutions. `feature_group_count` needs to be a divisor of both the input and
+the output feature dimension. If `feature_group_count` is greater than 1, it
+means that conceptually the input and output feature dimension and the `rhs`
+output feature dimension are split evenly into `feature_group_count` many
+groups, each group consisting of a consecutive subsequence of features. The
+input feature dimension of `rhs` needs to be equal to the `lhs` input feature
+dimension divided by `feature_group_count` (so it already has the size of a
+group of input features). The i-th groups are used together to compute
+`feature_group_count` many separate convolutions. The results of these
+convolutions are concatenated together in the output feature dimension.
+
+For depthwise convolution the `feature_group_count` argument would be set to the
+input feature dimension, and the filter would be reshaped from
+`[filter_height, filter_width, in_channels, channel_multiplier]` to
+`[filter_height, filter_width, 1, in_channels * channel_multiplier]`. For more
+details, see `tf.nn.depthwise_conv2d`.
+
+The output shape has these dimensions, in this order:
+
+* `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.
+
+The valid placements of the convolutional window are determined by the strides
+and the size of the base area after padding.
+
+To describe what a convolution does, consider a 2d convolution, and pick some
+fixed `batch`, `z`, `y`, `x` coordinates in the output. Then `(y,x)` is a
+position of a corner of the window within the base area (e.g. the upper left
+corner, depending on how you interpret the spatial dimensions). We now have a 2d
+window, taken from the base area, where each 2d point is associated to a 1d
+vector, so we get a 3d box. From the convolutional kernel, since we fixed the
+output coordinate `z`, we also have a 3d box. The two boxes have the same
+dimensions, so we can take the sum of the element-wise products between the two
+boxes (similar to a dot product). That is the output value.
+
+Note that if `output-z` is e.g., 5, then each position of the window produces 5
+values in the output into the `z` dimension of the output. These values differ
+in what part of the convolutional kernel is used - there is a separate 3d box of
+values used for each `output-z` coordinate. So you could think of it as 5
+separate convolutions with a different filter for each of them.
+
+Here is pseudo-code for a 2d convolution with padding and striding:
+
+```
+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;
+}
+```
+
+## ConvertElementType
+
+See also
+[`XlaBuilder::ConvertElementType`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+Similar to an element-wise `static_cast` in C++, performs an element-wise
+conversion operation from a data shape to a target shape. The dimensions must
+match, and the conversion is an element-wise one; e.g. `s32` elements become
+`f32` elements via an `s32`-to-`f32` conversion routine.
+
+<b> `ConvertElementType(operand, new_element_type)` </b>
+
+Arguments | Type | Semantics
+------------------ | --------------- | ---------------------------
+`operand` | `XlaOp` | array of type T with dims D
+`new_element_type` | `PrimitiveType` | type U
+
+The dimensions of the operand and the target shape must match. The source and
+destination element types must not be tuples.
+
+A conversion such as `T=s32` to `U=f32` will perform a normalizing int-to-float
+conversion routine such as round-to-nearest-even.
+
+> Note: The precise float-to-int and visa-versa conversions are currently
+> unspecified, but may become additional arguments to the convert operation in
+> the future. Not all possible conversions have been implemented for all
+>targets.
+
+```
+let a: s32[3] = {0, 1, 2};
+let b: f32[3] = convert(a, f32);
+then b == f32[3]{0.0, 1.0, 2.0}
+```
+
+## CrossReplicaSum
+
+See also
+[`XlaBuilder::CrossReplicaSum`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+Computes a sum across replicas.
+
+<b> `CrossReplicaSum(operand)` </b>
+
+Arguments | Type | Semantics
+--------- | ------- | -----------------------------
+`operand` | `XlaOp` | Array to sum across replicas.
+| `replica_group_ids` | `int64` vector | Group ID for each replica. |
+
+The output shape is the same as the input shape. For example, if there are two
+replicas and the operand has the value `(1.0, 2.5)` and `(3.0, 5.25)`
+respectively on the two replicas, then the output value from this op will be
+`(4.0, 7.75)` on both replicas.
+
+`replica_group_ids` identifies the group ID of each replica. The group ID must
+either be empty (all replicas belong to a single group), or contain the same
+number of elements as the number of replicas. For example, if
+`replica_group_ids` = {0, 1, 2, 3, 0, 1, 2, 3} has eight replicas, there are
+four subgroups of replica IDs: {0, 4}, {1, 5}, {2, 6}, and {3, 7}. The size of
+each subgroup *must* be identical, so, for example, using:
+`replica_group_ids` = {0, 1, 2, 0} for four replicas is invalid.
+
+Computing the result of CrossReplicaSum requires having one input from each
+replica, so if one replica executes a CrossReplicaSum node more times than
+another, then the former replica will wait forever. Since the replicas are all
+running the same program, there are not a lot of ways for that to happen, but it
+is possible when a while loop's condition depends on data from infeed and the
+data that is infed causes the while loop to iterate more times on one replica
+than another.
+
+## CustomCall
+
+See also
+[`XlaBuilder::CustomCall`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+Call a user-provided function within a computation.
+
+<b> `CustomCall(target_name, args..., shape)` </b>
+
+| Arguments | Type | Semantics |
+| ------------- | ---------------------- | --------------------------------- |
+| `target_name` | `string` | Name of the function. A call |
+: : : instruction will be emitted which :
+: : : targets this symbol name. :
+| `args` | sequence of N `XlaOp`s | N arguments of arbitrary type, |
+: : : which will be passed to the :
+: : : function. :
+| `shape` | `Shape` | Output shape of the function |
+
+The function signature is the same, regardless of the arity or type of args:
+
+```
+extern "C" void target_name(void* out, void** in);
+```
+
+For example, if CustomCall is used as follows:
+
+```
+let x = f32[2] {1,2};
+let y = f32[2x3] {{10, 20, 30}, {40, 50, 60}};
+
+CustomCall("myfunc", {x, y}, f32[3x3])
+```
+
+Here is an example of an implementation of `myfunc`:
+
+```
+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];
+ // ...
+}
+```
+
+The user-provided function must not have side-effects and its execution must be
+idempotent.
+
+> Note: The opaque nature of the user-provided function restricts optimization
+> opportunities for the compiler. Try to express your computation in terms of
+> native XLA ops whenever possible; only use CustomCall as a last resort.
+
+## Dot
+
+See also
+[`XlaBuilder::Dot`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+<b> `Dot(lhs, rhs)` </b>
+
+Arguments | Type | Semantics
+--------- | ------- | ---------------
+`lhs` | `XlaOp` | array of type T
+`rhs` | `XlaOp` | array of type T
+
+The exact semantics of this operation depend on the ranks of the operands:
+
+| Input | Output | Semantics |
+| ----------------------- | --------------------- | ----------------------- |
+| vector [n] `dot` vector | scalar | vector dot product |
+: [n] : : :
+| matrix [m x k] `dot` | vector [m] | matrix-vector |
+: vector [k] : : multiplication :
+| matrix [m x k] `dot` | matrix [m x n] | matrix-matrix |
+: matrix [k x n] : : multiplication :
+
+The operation performs sum of products over the last dimension of `lhs` and the
+one-before-last dimension of `rhs`. These are the "contracted" dimensions. The
+contracted dimensions of `lhs` and `rhs` must be of the same size. In practice,
+it can be used to perform dot products between vectors, vector/matrix
+multiplications or matrix/matrix multiplications.
+
+## DotGeneral
+
+See also
+[`XlaBuilder::DotGeneral`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+<b> `DotGeneral(lhs, rhs, dimension_numbers)` </b>
+
+Arguments | Type | Semantics
+------------------- | --------------------- | ---------------
+`lhs` | `XlaOp` | array of type T
+`rhs` | `XlaOp` | array of type T
+`dimension_numbers` | `DotDimensionNumbers` | array of type T
+
+As Dot, but allows contracting and batch dimension numbers to be specified for
+both the 'lhs' and 'rhs'.
+
+| DotDimensionNumbers Fields | Type | Semantics
+| --------- | ----------------------- | ---------------
+| 'lhs_contracting_dimensions' | repeated int64 | 'lhs' contracting dimension numbers |
+| 'rhs_contracting_dimensions' | repeated int64 | 'rhs' contracting dimension numbers |
+| 'lhs_batch_dimensions' | repeated int64 | 'lhs' batch dimension numbers |
+| 'rhs_batch_dimensions' | repeated int64 | 'rhs' batch dimension numbers |
+
+DotGeneral performs the sum of products over contracting dimensions specified
+in 'dimension_numbers'.
+
+Associated contracting dimension numbers from the 'lhs' and 'rhs' do not need
+to be the same, but must be listed in the same order in both
+'lhs/rhs_contracting_dimensions' arrays and have the same dimension sizes.
+There must be exactly one contracting dimension on both 'lhs' and 'rhs'.
+
+Example with contracting dimension numbers:
+
+```
+lhs = { {1.0, 2.0, 3.0},
+ {4.0, 5.0, 6.0} }
+
+rhs = { {1.0, 1.0, 1.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} }
+```
+
+Associated batch dimension numbers from the 'lhs' and 'rhs' must have the same
+dimension number, must be listed in the same order in both arrays, must
+have the same dimension sizes, and must be ordered before contracting and
+non-contracting/non-batch dimension numbers.
+
+Example with batch dimension numbers (batch size 2, 2x2 matrices):
+
+```
+lhs = { { {1.0, 2.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} } }
+
+DotDimensionNumbers dnums;
+dnums.add_lhs_contracting_dimensions(2);
+dnums.add_rhs_contracting_dimensions(1);
+dnums.add_lhs_batch_dimensions(0);
+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} } }
+```
+
+| Input | Output | Semantics |
+| ----------------------------------- | ----------------- | ---------------- |
+| [b0, m, k] `dot` [b0, k, n] | [b0, m, n] | batch matmul |
+| [b0, b1, m, k] `dot` [b0, b1, k, n] | [b0, b1, m, n] | batch matmul |
+
+It follows that the resulting dimension number starts with the batch dimension,
+then the 'lhs' non-contracting/non-batch dimension, and finally the 'rhs'
+non-contracting/non-batch dimension.
+
+## DynamicSlice
+
+See also
+[`XlaBuilder::DynamicSlice`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+DynamicSlice extracts a sub-array from the input array at dynamic
+`start_indices`. The size of the slice in each dimension is passed in
+`size_indices`, which specify the end point of exclusive slice intervals in each
+dimension: [start, start + size). The shape of `start_indices` must be rank ==
+1, with dimension size equal to the rank of `operand`.
+
+<b> `DynamicSlice(operand, start_indices, size_indices)` </b>
+
+| Arguments | Type | Semantics |
+| --------------- | ------------------- | ----------------------------------- |
+| `operand` | `XlaOp` | N dimensional array of type T |
+| `start_indices` | `XlaOp` | Rank 1 array of N integers |
+: : : containing the starting indices of :
+: : : the slice for each dimension. Value :
+: : : must be greater than or equal to :
+: : : zero. :
+| `size_indices` | `ArraySlice<int64>` | List of N integers containing the |
+: : : slice size for each dimension. Each :
+: : : value must be strictly greater than :
+: : : zero, and start + size must be less :
+: : : than or equal to the size of the :
+: : : dimension to avoid wrapping modulo :
+: : : dimension size. :
+
+The effective slice indices are computed by applying the following
+transformation for each index `i` in `[1, N)` before performing the slice:
+
+```
+start_indices[i] = clamp(start_indices[i], 0, operand.dimension_size[i] - size_indices[i])
+```
+
+This ensures that the extracted slice is always in-bounds with respect to the
+operand array. If the slice is in-bounds before the transformation is applied,
+the transformation has no effect.
+
+1-dimensional example:
+
+```
+let a = {0.0, 1.0, 2.0, 3.0, 4.0}
+let s = {2}
+
+DynamicSlice(a, s, {2}) produces:
+ {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} }
+let s = {2, 1}
+
+DynamicSlice(b, s, {2, 2}) produces:
+ { { 7.0, 8.0},
+ {10.0, 11.0} }
+```
+## DynamicUpdateSlice
+
+See also
+[`XlaBuilder::DynamicUpdateSlice`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+DynamicUpdateSlice generates a result which is the value of the input array
+`operand`, with a slice `update` overwritten at `start_indices`.
+The shape of `update` determines the shape of the sub-array of the result which
+is updated.
+The shape of `start_indices` must be rank == 1, with dimension size equal to
+the rank of `operand`.
+
+<b> `DynamicUpdateSlice(operand, update, start_indices)` </b>
+
+| Arguments | Type | Semantics |
+| --------------- | ------- | ------------------------------------------------ |
+| `operand` | `XlaOp` | N dimensional array of type T |
+| `update` | `XlaOp` | N dimensional array of type T containing the |
+: : : slice update. Each dimension of update shape :
+: : : must be strictly greater than zero, and start + :
+: : : update must be less than or equal to the operand :
+: : : size for each dimension to avoid generating :
+: : : out-of-bounds update indices. :
+| `start_indices` | `XlaOp` | Rank 1 array of N integers containing the |
+: : : starting indices of the slice for each :
+: : : dimension. Value must be greater than or equal :
+: : : to zero. :
+
+The effective slice indices are computed by applying the following
+transformation for each index `i` in `[1, N)` before performing the slice:
+
+```
+start_indices[i] = clamp(start_indices[i], 0, operand.dimension_size[i] - update.dimension_size[i])
+```
+
+This ensures that the updated slice is always in-bounds with respect to the
+operand array. If the slice is in-bounds before the transformation is applied,
+the transformation has no effect.
+
+1-dimensional example:
+
+```
+let a = {0.0, 1.0, 2.0, 3.0, 4.0}
+let u = {5.0, 6.0}
+let s = {2}
+
+DynamicUpdateSlice(a, u, s) produces:
+ {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} }
+let u =
+ { {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} }
+```
+
+## Element-wise binary arithmetic operations
+
+See also
+[`XlaBuilder::Add`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+A set of element-wise binary arithmetic operations is supported.
+
+<b> `Op(lhs, rhs)` </b>
+
+Where `Op` is one of `Add` (addition), `Sub` (subtraction), `Mul`
+(multiplication), `Div` (division), `Rem` (remainder), `Max` (maximum), `Min`
+(minimum), `LogicalAnd` (logical AND), or `LogicalOr` (logical OR).
+
+Arguments | Type | Semantics
+--------- | ------- | ----------------------------------------
+`lhs` | `XlaOp` | left-hand-side operand: array of type T
+`rhs` | `XlaOp` | right-hand-side operand: array of type T
+
+The arguments' shapes have to be either similar or compatible. See the
+[broadcasting](../../performance/xla/broadcasting.md) documentation about what it means for shapes to
+be compatible. The result of an operation has a shape which is the result of
+broadcasting the two input arrays. In this variant, operations between arrays of
+different ranks are *not* supported, unless one of the operands is a scalar.
+
+When `Op` is `Rem`, the sign of the result is taken from the dividend, and the
+absolute value of the result is always less than the divisor's absolute value.
+
+Integer division overflow (signed/unsigned division/remainder by zero or signed
+divison/remainder of `INT_SMIN` with `-1`) produces an implementation defined
+value.
+
+An alternative variant with different-rank broadcasting support exists for these
+operations:
+
+<b> `Op(lhs, rhs, broadcast_dimensions)` </b>
+
+Where `Op` is the same as above. This variant of the operation should be used
+for arithmetic operations between arrays of different ranks (such as adding a
+matrix to a vector).
+
+The additional `broadcast_dimensions` operand is a slice of integers used to
+expand the rank of the lower-rank operand up to the rank of the higher-rank
+operand. `broadcast_dimensions` maps the dimensions of the lower-rank shape to
+the dimensions of the higher-rank shape. The unmapped dimensions of the expanded
+shape are filled with dimensions of size one. Degenerate-dimension broadcasting
+then broadcasts the shapes along these degenerate dimensions to equalize the
+shapes of both operands. The semantics are described in detail on the
+[broadcasting page](../../performance/xla/broadcasting.md).
+
+## Element-wise comparison operations
+
+See also
+[`XlaBuilder::Eq`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+A set of standard element-wise binary comparison operations is supported. Note
+that standard IEEE 754 floating-point comparison semantics apply when comparing
+floating-point types.
+
+<b> `Op(lhs, rhs)` </b>
+
+Where `Op` is one of `Eq` (equal-to), `Ne` (not equal-to), `Ge`
+(greater-or-equal-than), `Gt` (greater-than), `Le` (less-or-equal-than), `Lt`
+(less-than).
+
+Arguments | Type | Semantics
+--------- | ------- | ----------------------------------------
+`lhs` | `XlaOp` | left-hand-side operand: array of type T
+`rhs` | `XlaOp` | right-hand-side operand: array of type T
+
+The arguments' shapes have to be either similar or compatible. See the
+[broadcasting](../../performance/xla/broadcasting.md) documentation about what it means for shapes to
+be compatible. The result of an operation has a shape which is the result of
+broadcasting the two input arrays with the element type `PRED`. In this variant,
+operations between arrays of different ranks are *not* supported, unless one of
+the operands is a scalar.
+
+An alternative variant with different-rank broadcasting support exists for these
+operations:
+
+<b> `Op(lhs, rhs, broadcast_dimensions)` </b>
+
+Where `Op` is the same as above. This variant of the operation should be used
+for comparison operations between arrays of different ranks (such as adding a
+matrix to a vector).
+
+The additional `broadcast_dimensions` operand is a slice of integers specifying
+the dimensions to use for broadcasting the operands. The semantics are described
+in detail on the [broadcasting page](../../performance/xla/broadcasting.md).
+
+## Element-wise unary functions
+
+XlaBuilder supports these element-wise unary functions:
+
+<b>`Abs(operand)`</b> Element-wise abs `x -> |x|`.
+
+<b>`Ceil(operand)`</b> Element-wise ceil `x -> ⌈x⌉`.
+
+<b>`Cos(operand)`</b> Element-wise cosine `x -> cos(x)`.
+
+<b>`Exp(operand)`</b> Element-wise natural exponential `x -> e^x`.
+
+<b>`Floor(operand)`</b> Element-wise floor `x -> ⌊x⌋`.
+
+<b>`IsFinite(operand)`</b> Tests whether each element of `operand` is finite,
+i.e., is not positive or negative infinity, and is not `NaN`. Returns an array
+of `PRED` values with the same shape as the input, where each element is `true`
+if and only if the corresponding input element is finite.
+
+<b>`Log(operand)`</b> Element-wise natural logarithm `x -> ln(x)`.
+
+<b>`LogicalNot(operand)`</b> Element-wise logical not `x -> !(x)`.
+
+<b>`Neg(operand)`</b> Element-wise negation `x -> -x`.
+
+<b>`Sign(operand)`</b> Element-wise sign operation `x -> sgn(x)` where
+
+$$\text{sgn}(x) = \begin{cases} -1 & x < 0\\ 0 & x = 0\\ 1 & x > 0 \end{cases}$$
+
+using the comparison operator of the element type of `operand`.
+
+<b>`Tanh(operand)`</b> Element-wise hyperbolic tangent `x -> tanh(x)`.
+
+
+Arguments | Type | Semantics
+--------- | ------- | ---------------------------
+`operand` | `XlaOp` | The operand to the function
+
+The function is applied to each element in the `operand` array, resulting in an
+array with the same shape. It is allowed for `operand` to be a scalar (rank 0).
+
+## Gather
+
+The XLA gather operation stitches together several slices (each slice at a
+potentially different runtime offset) of an input array.
+
+### General Semantics
+
+See also
+[`XlaBuilder::Gather`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+For a more intuitive description, see the "Informal Description" section below.
+
+<b> `gather(operand, start_indices, offset_dims, collapsed_slice_dims, slice_sizes, start_index_map)` </b>
+
+|Arguments | Type | Semantics |
+|----------------- | ----------------------- | --------------------------------|
+|`operand` | `XlaOp` | The array we’re gathering |
+: : : from. :
+|`start_indices` | `XlaOp` | Array containing the starting |
+: : : indices of the slices we gather.:
+|`index_vector_dim` | `int64` | The dimension in |
+: : : `start_indices` that "contains" :
+: : : the starting indices. See :
+: : : below for a detailed :
+: : : description. :
+|`offset_dims` | `ArraySlice<int64>` | The set of dimensions in the :
+: : : output shape that offset into a :
+: : : array sliced from operand. :
+|`slice_sizes` | `ArraySlice<int64>` | `slice_sizes[i]` is the bounds |
+: : : for the slice on dimension `i`.:
+|`collapsed_slice_dims` | `ArraySlice<int64>` | The set of dimensions in each :
+| : | slice that are collapsed away. :
+| : | These dimensions must have size:
+| : | 1. |
+|`start_index_map` | `ArraySlice<int64>` | A map that describes how to map|
+: : : indices in `start_indices` to :
+: : : to legal indices into operand. :
+
+For convenience, we label dimensions in the output array not in `offset_dims`
+as `batch_dims`.
+
+The output is an array of rank `batch_dims.size` + `operand.rank` -
+`collapsed_slice_dims`.size.
+
+If `index_vector_dim` is equal to `start_indices.rank` we implicitly consider
+`start_indices` to have a trailing `1` dimension (i.e. if `start_indices` was of
+shape `[6,7]` and `index_vector_dim` is `2` then we implicitly consider the
+shape of `start_indices` to be `[6,7,1]`).
+
+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).
+
+ 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`.
+
+ 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.
+
+`expand_offset_dims` is the monotonic function with domain [`0`, `offset.size`)
+and range [`0`, `operand.rank`) \ `collapsed_slice_dims`. So if, e.g.,
+`offset.size` is `4`, `operand.rank` is `6` and `collapsed_slice_dims` is {`0`,
+`2`} then `expand_offset_dims` is {`0`→`1`, `1`→`3`, `2`→`4`, `3`→`5`}.
+
+### Informal Description and Examples
+
+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 `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 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`.
+
+`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
+change the operation fundamentally, but makes the visual representation more
+cumbersome.
+
+To get an intuition on how all of the above fits together, let's look at an
+example that gathers 5 slices of shape `[8,6]` from a `[16,11]` array. The
+position of a slice into the `[16,11]` array can be represented as an index
+vector of shape `S64[2]`, so the set of 5 positions can be represented as a
+`S64[5,2]` array.
+
+The behavior of the gather operation can then be depicted as an index
+transformation that takes [`G`,`O`<sub>`0`</sub>,`O`<sub>`1`</sub>], an index in
+the output shape, and maps it to an element in the input array in the following
+way:
+
+<div style="width:95%; margin:auto; margin-bottom:10px; margin-top:20px;">
+ <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`.
+The element in the output array at index
+[`G`,`O`<sub>`0`</sub>,`O`<sub>`1`</sub>] is then the element in the input
+array at index [`X`+`O`<sub>`0`</sub>,`Y`+`O`<sub>`1`</sub>].
+
+`slice_sizes` is `[8,6]`, which decides the range of W<sub>`0`</sub> and
+W<sub>`1`</sub>, and this in turn decides the bounds of the slice.
+
+This gather operation acts as a batch dynamic slice with `G` as the batch
+dimension.
+
+The gather indices may be multidimensional. For instance, a more general
+version of the example above using a "gather indices" array of shape `[4,5,2]`
+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">
+</div>
+
+Again, this acts as a batch dynamic slice `G`<sub>`0`</sub> and
+`G`<sub>`1`</sub> as the batch dimensions. The slice size is still `[8,6]`.
+
+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.
+
+ 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.
+
+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">
+</div>
+
+`G`<sub>`0`</sub> and `G`<sub>`1`</sub> are used to slice out a starting index
+from the gather indices array as usual, except the starting index has only one
+element, `X`. Similarly, there is only one output offset index with the value
+`O`<sub>`0`</sub>. However, before being used as indices into the input array,
+these are expanded in accordance to "Gather Index Mapping" (`start_index_map` in
+the formal description) and "Offset Mapping" (`expand_offset_dims` in the formal
+description) into [`0`,`O`<sub>`0`</sub>] and [`X`,`0`] respectively, adding up
+to [`X`,`O`<sub>`0`</sub>]. In other words, the output index
+[`G`<sub>`0`</sub>,`G`<sub>`1`</sub>,`O`<sub>`0`</sub>] maps to the input index
+[`GatherIndices`[`G`<sub>`0`</sub>,`G`<sub>`1`</sub>,`0`],`X`] which gives us
+the semantics for `tf.gather_nd`.
+
+`slice_sizes` for this case is `[1,11]`. Intuitively this means that every
+index `X` in the gather indices array picks an entire row and the result is the
+concatenation of all these rows.
+
+## GetTupleElement
+
+See also
+[`XlaBuilder::GetTupleElement`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+Indexes into a tuple with a compile-time-constant value.
+
+The value must be a compile-time-constant so that shape inference can determine
+the type of the resulting value.
+
+This is analogous to `std::get<int N>(t)` in C++. Conceptually:
+
+```
+let v: f32[10] = f32[10]{0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
+let s: s32 = 5;
+let t: (f32[10], s32) = tuple(v, s);
+let element_1: s32 = gettupleelement(t, 1); // Inferred shape matches s32.
+```
+
+See also `tf.tuple`.
+
+## Infeed
+
+See also
+[`XlaBuilder::Infeed`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+<b> `Infeed(shape)` </b>
+
+| Argument | Type | Semantics |
+| -------- | ------- | ----------------------------------------------------- |
+| `shape` | `Shape` | Shape of the data read from the Infeed interface. The |
+: : : layout field of the shape must be set to match the :
+: : : layout of the data sent to the device; otherwise its :
+: : : behavior is undefined. :
+
+Reads a single data item from the implicit Infeed streaming interface of the
+device, interpreting the data as the given shape and its layout, and returns a
+`XlaOp` of the data. Multiple Infeed operations are allowed in a
+computation, but there must be a total order among the Infeed operations. For
+example, two Infeeds in the code below have a total order since there is a
+dependency between the while loops.
+
+```
+result1 = while (condition, init = init_value) {
+ Infeed(shape)
+}
+
+result2 = while (condition, init = result1) {
+ Infeed(shape)
+}
+```
+
+Nested tuple shapes are not supported. For an empty tuple shape, the Infeed
+operation is effectively a no-op and proceeds without reading any data from the
+Infeed of the device.
+
+> Note: We plan to allow multiple Infeed operations without a total order, in
+> which case the compiler will provide information about how the Infeed
+> operations are serialized in the compiled program.
+
+## Iota
+
+<b> `Iota()` </b>
+
+Builds a constant literal on device rather than a potentially large host
+transfer. Creates a rank 1 tensor of values starting at zero and incrementing
+by one.
+
+Arguments | Type | Semantics
+------------------ | --------------- | ---------------------------
+`type` | `PrimitiveType` | type U
+`size` | `int64` | The number of elements in the tensor.
+
+## Map
+
+See also
+[`XlaBuilder::Map`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+<b> `Map(operands..., computation)` </b>
+
+| Arguments | Type | Semantics |
+| ----------------- | ---------------------- | ------------------------------ |
+| `operands` | sequence of N `XlaOp`s | N arrays of types T_0..T_{N-1} |
+| `computation` | `XlaComputation` | computation of type `T_0, T_1, |
+: : : ..., T_{N + M -1} -> S` with N :
+: : : parameters of type T and M of :
+: : : arbitrary type :
+| `dimensions` | `int64` array | array of map dimensions |
+
+Applies a scalar function over the given `operands` arrays, producing an array
+of the same dimensions where each element is the result of the mapped function
+applied to the corresponding elements in the input arrays.
+
+The mapped function is an arbitrary computation with the restriction that it has
+N inputs of scalar type `T` and a single output with type `S`. The output has
+the same dimensions as the operands except that the element type T is replaced
+with S.
+
+For example: `Map(op1, op2, op3, computation, par1)` maps `elem_out <-
+computation(elem1, elem2, elem3, par1)` at each (multi-dimensional) index in the
+input arrays to produce the output array.
+
+## Pad
+
+See also
+[`XlaBuilder::Pad`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+<b> `Pad(operand, padding_value, padding_config)` </b>
+
+| Arguments | Type | Semantics |
+| ---------------- | --------------- | --------------------------------------- |
+| `operand` | `XlaOp` | array of type `T` |
+| `padding_value` | `XlaOp` | scalar of type `T` to fill in the added |
+: : : padding :
+| `padding_config` | `PaddingConfig` | padding amount on both edges (low, |
+: : : high) and between the elements of each :
+: : : dimension :
+
+Expands the given `operand` array by padding around the array as well as between
+the elements of the array with the given `padding_value`. `padding_config`
+specifies the amount of edge padding and the interior padding for each
+dimension.
+
+`PaddingConfig` is a repeated field of `PaddingConfigDimension`, which contains
+three fields for each dimension: `edge_padding_low`, `edge_padding_high`, and
+`interior_padding`. `edge_padding_low` and `edge_padding_high` specify the
+amount of padding added at the low-end (next to index 0) and the high-end (next
+to the highest index) of each dimension respectively. The amount of edge padding
+can be negative -- the absolute value of negative padding indicates the number
+of elements to remove from the specified dimension. `interior_padding` specifies
+the amount of padding added between any two elements in each dimension. Interior
+padding occurs logically before edge padding, so in the case of negative edge
+padding elements are removed from the interior-padded operand. This operation is
+a no-op if the edge padding pairs are all (0, 0) and the interior padding values
+are all 0. The figure below shows examples of different `edge_padding` and
+`interior_padding` values for a two-dimensional array.
+
+<div style="width:95%; margin:auto; margin-bottom:10px; margin-top:20px;">
+ <img style="width:100%" src="https://www.tensorflow.org/images/ops_pad.png">
+</div>
+
+## Recv
+
+See also
+[`XlaBuilder::Recv`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+<b> `Recv(shape, channel_handle)` </b>
+
+| Arguments | Type | Semantics |
+| ---------------- | --------------- | ------------------------------------ |
+| `shape` | `Shape` | shape of the data to receive |
+| `channel_handle` | `ChannelHandle` | unique identifier for each send/recv pair |
+
+Receives data of the given shape from a `Send` instruction in another
+computation that shares the same channel handle. Returns a
+XlaOp for the received data.
+
+The client API of `Recv` operation represents synchronous communication.
+However, the instruction is internally decomposed into 2 HLO instructions
+(`Recv` and `RecvDone`) to enable asynchronous data transfers. See also
+[`HloInstruction::CreateRecv` and `HloInstruction::CreateRecvDone`](https://www.tensorflow.org/code/tensorflow/compiler/xla/service/hlo_instruction.h).
+
+<b>`Recv(const Shape& shape, int64 channel_id)`</b>
+
+Allocates resources required to receive data from a `Send` instruction with the
+same channel_id. Returns a context for the allocated resources, which is used
+by a following `RecvDone` instruction to wait for the completion of the data
+transfer. The context is a tuple of {receive buffer (shape), request identifier
+(U32)} and it can only be used by a `RecvDone` instruction.
+
+<b> `RecvDone(HloInstruction context)` </b>
+
+Given a context created by a `Recv` instruction, waits for the data transfer to
+complete and returns the received data.
+
+## Reduce
+
+See also
+[`XlaBuilder::Reduce`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+Applies a reduction function to one or more arrays in parallel.
+
+<b> `Reduce(operands..., init_values..., computation, dimensions)` </b>
+
+Arguments | Type | Semantics
+------------- | --------------------- | ---------------------------------------
+`operands` | Sequence of N `XlaOp` | N arrays of types `T_0, ..., T_N`.
+`init_values` | Sequence of N `XlaOp` | N scalars of types `T_0, ..., T_N`.
+`computation` | `XlaComputation` | computation of type
+ : : `T_0, ..., T_N, T_0, ..., T_N -> Collate(T_0, ..., T_N)`
+`dimensions` | `int64` array | unordered array of dimensions to reduce
+
+Where:
+* N is required to be greater or equal to 1.
+* All input arrays must have the same dimensions.
+* If `N = 1`, `Collate(T)` is `T`.
+* If `N > 1`, `Collate(T_0, ..., T_N)` is a tuple of `N` elements of type `T`.
+
+The output of the op is `Collate(Q_0, ..., Q_N)` where `Q_i` is an array of type
+`T_i`, the dimensions of which are described below.
+
+This operation reduces one or more dimensions of each input array into scalars.
+The rank of each returned array is `rank(operand) - len(dimensions)`.
+`init_value` is the initial value used for every reduction and may be inserted
+anywhere during computation by the back-end. In most cases, `init_value` is an
+identity of the reduction function (for example, 0 for addition). The applied
+`computation` is always passed the `init_value` on the left-hand side.
+
+The evaluation order of the reduction function is arbitrary and may be
+non-deterministic. Therefore, the reduction function should not be overly
+sensitive to reassociation.
+
+Some reduction functions like addition are not strictly associative for floats.
+However, if the range of the data is limited, floating-point addition is close
+enough to being associative for most practical uses. It is possible to conceive
+of some completely non-associative reductions, however, and these will produce
+incorrect or unpredictable results in XLA reductions.
+
+As an example, when reducing across one dimension in a single 1D array with
+values [10, 11, 12, 13], with reduction function `f` (this is `computation`)
+then that could be computed as
+
+`f(10, f(11, f(12, f(init_value, 13)))`
+
+but there are also many other possibilities, e.g.
+
+`f(init_value, f(f(10, f(init_value, 11)), f(f(init_value, 12), f(init_value, 13))))`
+
+The following is a rough pseudo-code example of how reduction could be
+implemented, using summation as the reduction computation with an initial value
+of 0.
+
+```python
+result_shape <- remove all dims in dimensions from operand_shape
+
+# Iterate over all elements in result_shape. The number of r's here is equal
+# to the rank of the result
+for r0 in range(result_shape[0]), r1 in range(result_shape[1]), ...:
+ # Initialize this result element
+ result[r0, r1...] <- 0
+
+ # Iterate over all the reduction dimensions
+ for d0 in range(dimensions[0]), d1 in range(dimensions[1]), ...:
+ # Increment the result element with the value of the operand's element.
+ # The index of the operand's element is constructed from all ri's and di's
+ # in the right order (by construction ri's and di's together index over the
+ # whole operand shape).
+ result[r0, r1...] += operand[ri... di]
+```
+
+Here's an example of reducing a 2D array (matrix). The shape has rank 2,
+dimension 0 of size 2 and dimension 1 of size 3:
+
+<div style="width:95%; margin:auto; margin-bottom:10px; margin-top:20px;">
+ <img style="width:35%" src="https://www.tensorflow.org/images/ops_2d_matrix.png">
+</div>
+
+Results of reducing dimensions 0 or 1 with an "add" function:
+
+<div style="width:95%; margin:auto; margin-bottom:10px; margin-top:20px;">
+ <img style="width:35%" src="https://www.tensorflow.org/images/ops_reduce_from_2d_matrix.png">
+</div>
+
+Note that both reduction results are 1D arrays. The diagram shows one as column
+and another as row just for visual convenience.
+
+For a more complex example, here is a 3D array. Its rank is 3, dimension 0 of
+size 4, dimension 1 of size 2 and dimension 2 of size 3. For simplicity, the
+values 1 to 6 are replicated across dimension 0.
+
+<div style="width:95%; margin:auto; margin-bottom:10px; margin-top:20px;">
+ <img style="width:35%" src="https://www.tensorflow.org/images/ops_reduce_from_3d_matrix.png">
+</div>
+
+Similarly to the 2D example, we can reduce just one dimension. If we reduce
+dimension 0, for example, we get a rank-2 array where all values across
+dimension 0 were folded into a scalar:
+
+```text
+| 4 8 12 |
+| 16 20 24 |
+```
+
+If we reduce dimension 2, we also get a rank-2 array where all values across
+dimension 2 were folded into a scalar:
+
+```text
+| 6 15 |
+| 6 15 |
+| 6 15 |
+| 6 15 |
+```
+
+Note that the relative order between the remaining dimensions in the input is
+preserved in the output, but some dimensions may get assigned new numbers (since
+the rank changes).
+
+We can also reduce multiple dimensions. Add-reducing dimensions 0 and 1 produces
+the 1D array `| 20 28 36 |`.
+
+Reducing the 3D array over all its dimensions produces the scalar `84`.
+
+When `N > 1`, reduce function application is slightly more complex, as it is
+applied simultaneously to all inputs. For example, consider the following
+reduction function, which can be used to compute the max and the argmax of a
+a 1-D tensor in parallel:
+
+```
+f: (Float, Int, Float, Int) -> Float, Int
+f(max, argmax, value, index):
+ if value >= argmax:
+ return (value, index)
+ else:
+ return (max, argmax)
+```
+
+For 1-D Input arrays `V = Float[N], K = Int[N]`, and init values
+`I_V = Float, I_K = Int`, the result `f_(N-1)` of reducing across the only
+input dimension is equivalent to the following recursive application:
+```
+f_0 = f(I_V, I_K, V_0, K_0)
+f_1 = f(f_0.first, f_0.second, V_1, K_1)
+...
+f_(N-1) = f(f_(N-2).first, f_(N-2).second, V_(N-1), K_(N-1))
+```
+
+Applying this reduction to an array of values, and an array of sequential
+indices (i.e. iota), will co-iterate over the arrays, and return a tuple
+containing the maximal value and the matching index.
+
+## ReducePrecision
+
+See also
+[`XlaBuilder::ReducePrecision`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+Models the effect of converting floating-point values to a lower-precision
+format (such as IEEE-FP16) and back to the original format. The number of
+exponent and mantissa bits in the lower-precision format can be specified
+arbitrarily, although all bit sizes may not be supported on all hardware
+implementations.
+
+<b> `ReducePrecision(operand, mantissa_bits, exponent_bits)` </b>
+
+Arguments | Type | Semantics
+--------------- | ------- | -------------------------------------------------
+`operand` | `XlaOp` | array of floating-point type `T`.
+`exponent_bits` | `int32` | number of exponent bits in lower-precision format
+`mantissa_bits` | `int32` | number of mantissa bits in lower-precision format
+
+The result is an array of type `T`. The input values are rounded to the nearest
+value representable with the given number of mantissa bits (using "ties to even"
+semantics), and any values that exceed the range specified by the number of
+exponent bits are clamped to positive or negative infinity. `NaN` values are
+retained, although they may be converted to canonical `NaN` values.
+
+The lower-precision format must have at least one exponent bit (in order to
+distinguish a zero value from an infinity, since both have a zero mantissa), and
+must have a non-negative number of mantissa bits. The number of exponent or
+mantissa bits may exceed the corresponding value for type `T`; the corresponding
+portion of the conversion is then simply a no-op.
+
+## ReduceWindow
+
+See also
+[`XlaBuilder::ReduceWindow`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+Applies a reduction function to all elements in each window of the input
+multi-dimensional array, producing an output multi-dimensional array with the
+same number of elements as the number of valid positions of the window. A
+pooling layer can be expressed as a `ReduceWindow`. Similar to
+[`Reduce`](#reduce), the applied `computation` is always passed the `init_value`
+on the left-hand side.
+
+<b> `ReduceWindow(operand, init_value, computation, window_dimensions,
+window_strides, padding)` </b>
+
+| Arguments | Type | Semantics |
+| ------------------- | ------------------- | -------------------------------- |
+| `operand` | `XlaOp` | N dimensional array containing |
+: : : elements of type T. This is the :
+: : : base area on which the window is :
+: : : placed. :
+| `init_value` | `XlaOp` | Starting value for the |
+: : : reduction. See [Reduce](#reduce) :
+: : : for details. :
+| `computation` | `XlaComputation` | Reduction function of type `T, T |
+: : : -> T`, to apply to all elements :
+: : : in each window :
+| `window_dimensions` | `ArraySlice<int64>` | array of integers for window |
+: : : dimension values :
+| `window_strides` | `ArraySlice<int64>` | array of integers for window |
+: : : stride values :
+| `padding` | `Padding` | padding type for window |
+: : : (Padding\:\:kSame or :
+: : : Padding\:\:kValid) :
+
+Below code and figure shows an example of using `ReduceWindow`. Input is a
+matrix of size [4x6] and both window_dimensions and window_stride_dimensions are
+[2x3].
+
+```
+// Create a computation for the reduction (maximum).
+XlaComputation max;
+{
+ XlaBuilder builder(client_, "max");
+ auto y = builder.Parameter(0, ShapeUtil::MakeShape(F32, {}), "y");
+ auto x = builder.Parameter(1, ShapeUtil::MakeShape(F32, {}), "x");
+ builder.Max(y, x);
+ max = builder.Build().ConsumeValueOrDie();
+}
+
+// Create a ReduceWindow computation with the max reduction computation.
+XlaBuilder builder(client_, "reduce_window_2x3");
+auto shape = ShapeUtil::MakeShape(F32, {4, 6});
+auto input = builder.Parameter(0, shape, "input");
+builder.ReduceWindow(
+ input, *max,
+ /*init_val=*/builder.ConstantLiteral(LiteralUtil::MinValue(F32)),
+ /*window_dimensions=*/{2, 3},
+ /*window_stride_dimensions=*/{2, 3},
+ Padding::kValid);
+```
+
+<div style="width:95%; margin:auto; margin-bottom:10px; margin-top:20px;">
+ <img style="width:35%" src="https://www.tensorflow.org/images/ops_reduce_window.png">
+</div>
+
+Stride of 1 in a dimension specifies that the position of a window in the
+dimension is 1 element away from its adjacent window. In order to specify that
+no windows overlap with each other, window_stride_dimensions should be equal to
+window_dimensions. The figure below illustrates the use of two different stride
+values. Padding is applied to each dimension of the input and the calculations
+are the same as though the input came in with the dimensions it has after
+padding.
+
+<div style="width:95%; margin:auto; margin-bottom:10px; margin-top:20px;">
+ <img style="width:75%" src="https://www.tensorflow.org/images/ops_reduce_window_stride.png">
+</div>
+
+The evaluation order of the reduction function is arbitrary and may be
+non-deterministic. Therefore, the reduction function should not be overly
+sensitive to reassociation. See the discussion about associativity in the
+context of [`Reduce`](#reduce) for more details.
+
+## Reshape
+
+See also
+[`XlaBuilder::Reshape`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h)
+and the [`Collapse`](#collapse) operation.
+
+Reshapes the dimensions of an array into a new configuration.
+
+<b> `Reshape(operand, new_sizes)` </b>
+<b> `Reshape(operand, dimensions, new_sizes)` </b>
+
+Arguments | Type | Semantics
+------------ | -------------- | ---------------------------------------
+`operand` | `XlaOp` | array of type T
+`dimensions` | `int64` vector | order in which dimensions are collapsed
+`new_sizes` | `int64` vector | vector of sizes of new dimensions
+
+Conceptually, reshape first flattens an array into a one-dimensional vector of
+data values, and then refines this vector into a new shape. The input arguments
+are an arbitrary array of type T, a compile-time-constant vector of dimension
+indices, and a compile-time-constant vector of dimension sizes for the result.
+The values in the `dimension` vector, if given, must be a permutation of all of
+T's dimensions; the default if not given is `{0, ..., rank - 1}`. The order of
+the dimensions in `dimensions` is from slowest-varying dimension (most major) to
+fastest-varying dimension (most minor) in the loop nest which collapses the
+input array into a single dimension. The `new_sizes` vector determines the size
+of the output array. The value at index 0 in `new_sizes` is the size of
+dimension 0, the value at index 1 is the size of dimension 1, and so on. The
+product of the `new_size` dimensions must equal the product of the operand's
+dimension sizes. When refining the collapsed array into the multidimensional
+array defined by `new_sizes`, the dimensions in `new_sizes` are ordered from
+slowest varying (most major) and to fastest varying (most minor).
+
+For example, let v be an array of 24 elements:
+
+```
+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}}};
+
+In-order collapse:
+let v012_24 = Reshape(v, {0,1,2}, {24});
+then v012_24 == 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};
+
+let v012_83 = Reshape(v, {0,1,2}, {8,3});
+then v012_83 == 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}};
+
+Out-of-order collapse:
+let v021_24 = Reshape(v, {1,2,0}, {24});
+then v012_24 == f32[24] {10, 20, 30, 40, 11, 21, 31, 41, 12, 22, 32, 42,
+ 15, 25, 35, 45, 16, 26, 36, 46, 17, 27, 37, 47};
+
+let v021_83 = Reshape(v, {1,2,0}, {8,3});
+then v021_83 == f32[8x3] {{10, 20, 30}, {40, 11, 21},
+ {31, 41, 12}, {22, 32, 42},
+ {15, 25, 35}, {45, 16, 26},
+ {36, 46, 17}, {27, 37, 47}};
+
+
+let v021_262 = Reshape(v, {1,2,0}, {2,6,2});
+then v021_262 == f32[2x6x2] {{{10, 20}, {30, 40},
+ {11, 21}, {31, 41},
+ {12, 22}, {32, 42}},
+ {{15, 25}, {35, 45},
+ {16, 26}, {36, 46},
+ {17, 27}, {37, 47}}};
+```
+
+As a special case, reshape can transform a single-element array to a scalar and
+vice versa. For example,
+
+```
+Reshape(f32[1x1] {{5}}, {0,1}, {}) == 5;
+Reshape(5, {}, {1,1}) == f32[1x1] {{5}};
+```
+
+## Rev (reverse)
+
+See also
+[`XlaBuilder::Rev`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+<b>`Rev(operand, dimensions)`</b>
+
+Arguments | Type | Semantics
+------------ | ------------------- | ---------------------
+`operand` | `XlaOp` | array of type T
+`dimensions` | `ArraySlice<int64>` | dimensions to reverse
+
+Reverses the order of elements in the `operand` array along the specified
+`dimensions`, generating an output array of the same shape. Each element of the
+operand array at a multidimensional index is stored into the output array at a
+transformed index. The multidimensional index is transformed by reversing the
+index in each dimension to be reversed (i.e., if a dimension of size N is one of
+the reversing dimensions, its index i is transformed into N - 1 - i).
+
+One use for the `Rev` operation is to reverse the convolution weight array along
+the two window dimensions during the gradient computation in neural networks.
+
+## RngNormal
+
+See also
+[`XlaBuilder::RngNormal`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+Constructs an output of a given shape with random numbers generated following
+the $$N(\mu, \sigma)$$ normal distribution. The parameters $$\mu$$ and
+$$\sigma$$, and output shape have to have a floating point elemental type. The
+parameters furthermore have to be scalar valued.
+
+<b>`RngNormal(mu, sigma, shape)`</b>
+
+| Arguments | Type | Semantics |
+| --------- | ------- | --------------------------------------------------- |
+| `mu` | `XlaOp` | Scalar of type T specifying mean of generated |
+: : : numbers :
+| `sigma` | `XlaOp` | Scalar of type T specifying standard deviation of |
+: : : generated numbers :
+| `shape` | `Shape` | Output shape of type T |
+
+## RngUniform
+
+See also
+[`XlaBuilder::RngUniform`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+Constructs an output of a given shape with random numbers generated following
+the uniform distribution over the interval $$[a,b)$$. The parameters and output
+element type have to be a boolean type, an integral type or a floating point
+types, and the types have to be consistent. The CPU and GPU backends currently
+only support F64, F32, F16, BF16, S64, U64, S32 and U32. Furthermore, the
+parameters need to be scalar valued. If $$b <= a$$ the result is
+implementation-defined.
+
+<b>`RngUniform(a, b, shape)`</b>
+
+| Arguments | Type | Semantics |
+| --------- | ----------------------- | --------------------------------- |
+| `a` | `XlaOp` | Scalar of type T specifying lower |
+: : : limit of interval :
+| `b` | `XlaOp` | Scalar of type T specifying upper |
+: : : limit of interval :
+| `shape` | `Shape` | Output shape of type T |
+
+## Scatter
+
+The XLA scatter operation generates a result which is the value of the input
+tensor `operand`, with several slices (at indices specified by
+`scatter_indices`) updated with the values in `updates` using
+`update_computation`.
+
+See also
+[`XlaBuilder::Scatter`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+<b> `scatter(operand, scatter_indices, updates, update_computation, index_vector_dim, update_window_dims, inserted_window_dims, scatter_dims_to_operand_dims)` </b>
+
+|Arguments | Type | Semantics |
+|------------------|------------------------|----------------------------------|
+|`operand` | `XlaOp` | Tensor to be scattered into. |
+|`scatter_indices` | `XlaOp` | Tensor containing the starting |
+: : : indices of the slices that must :
+: : : be scattered to. :
+|`updates` | `XlaOp` | Tensor containing the values that|
+: : : must be used for scattering. :
+|`update_computation`| `XlaComputation` | Computation to be used for |
+: : : combining the existing values in :
+: : : the input tensor and the updates :
+: : : during scatter. This computation :
+: : : should be of type `T, T -> T`. :
+|`index_vector_dim`| `int64` | The dimension in |
+: : : `scatter_indices` that contains :
+: : : the starting indices. :
+|`update_window_dims`| `ArraySlice<int64>` | The set of dimensions in |
+: : : `updates` shape that are _window :
+: : : dimensions_. :
+|`inserted_window_dims`| `ArraySlice<int64>`| The set of _window dimensions_ |
+: : : that must be inserted into :
+: : : `updates` shape. :
+|`scatter_dims_to_operand_dims`| `ArraySlice<int64>` | A dimensions map from |
+: : : the scatter indices to the :
+: : : operand index space. This array :
+: : : is interpreted as mapping `i` to :
+: : : `scatter_dims_to_operand_dims[i]`:
+: : : . It has to be one-to-one and :
+: : : total. :
+
+If `index_vector_dim` is equal to `scatter_indices.rank` we implicitly consider
+`scatter_indices` to have a trailing `1` dimension.
+
+We define `update_scatter_dims` of type `ArraySlice<int64>` as the set of
+dimensions in `updates` shape that are not in `update_window_dims`, in ascending
+order.
+
+The arguments of scatter should follow these constraints:
+
+ - `updates` tensor must be of rank `update_window_dims.size +
+ scatter_indices.rank - 1`.
+
+ - Bounds of dimension `i` in `updates` must conform to the following:
+ - If `i` is present in `update_window_dims` (i.e. equal to
+ `update_window_dims`[`k`] for some `k`), then the bound of dimension
+ `i` in `updates` must not exceed the corresponding bound of `operand`
+ after accounting for the `inserted_window_dims` (i.e.
+ `adjusted_window_bounds`[`k`], where `adjusted_window_bounds` contains
+ the bounds of `operand` with the bounds at indices
+ `inserted_window_dims` removed).
+ - If `i` is present in `update_scatter_dims` (i.e. equal to
+ `update_scatter_dims`[`k`] for some `k`), then the bound of dimension
+ `i` in `updates` must be equal to the corresponding bound of
+ `scatter_indices`, skipping `index_vector_dim` (i.e.
+ `scatter_indices.shape.dims`[`k`], if `k` < `index_vector_dim` and
+ `scatter_indices.shape.dims`[`k+1`] otherwise).
+
+ - `update_window_dims` must be in ascending order, not have any repeating
+ dimension numbers, and be in the range `[0, updates.rank)`.
+
+ - `inserted_window_dims` must be in ascending order, not have any
+ repeating dimension numbers, and be in the range `[0, operand.rank)`.
+
+ - `scatter_dims_to_operand_dims.size` must be equal to
+ `scatter_indices`[`index_vector_dim`], and its values must be in the range
+ `[0, operand.rank)`.
+
+For a given index `U` in the `updates` tensor, the corresponding index `I` in
+the `operand` tensor into which this update has to be applied is computed as
+follows:
+
+ 1. Let `G` = { `U`[`k`] for `k` in `update_scatter_dims` }. Use `G` to look up
+ an index vector `S` in the `scatter_indices` tensor such that `S`[`i`] =
+ `scatter_indices`[Combine(`G`, `i`)] where Combine(A, b) inserts b at
+ positions `index_vector_dim` into A.
+ 2. Create an index `S`<sub>`in`</sub> into `operand` using `S` by scattering
+ `S` using the `scatter_dims_to_operand_dims` map. More formally:
+ 1. `S`<sub>`in`</sub>[`scatter_dims_to_operand_dims`[`k`]] = `S`[`k`] if
+ `k` < `scatter_dims_to_operand_dims.size`.
+ 2. `S`<sub>`in`</sub>[`_`] = `0` otherwise.
+ 3. Create an index `W`<sub>`in`</sub> into `operand` by scattering the indices
+ at `update_window_dims` in `U` according to `inserted_window_dims`.
+ More formally:
+ 1. `W`<sub>`in`</sub>[`window_dims_to_operand_dims`(`k`)] = `U`[`k`] if
+ `k` < `update_window_dims.size`, where `window_dims_to_operand_dims`
+ is the monotonic function with domain [`0`, `update_window_dims.size`)
+ and range [`0`, `operand.rank`) \\ `inserted_window_dims`. (For
+ example, if `update_window_dims.size` is `4`, `operand.rank` is `6`,
+ and `inserted_window_dims` is {`0`, `2`} then
+ `window_dims_to_operand_dims` is {`0`→`1`, `1`→`3`, `2`→`4`,
+ `3`→`5`}).
+ 2. `W`<sub>`in`</sub>[`_`] = `0` otherwise.
+ 4. `I` is `W`<sub>`in`</sub> + `S`<sub>`in`</sub> where + is element-wise
+ addition.
+
+In summary, the scatter operation can be defined as follows.
+
+ - Initialize `output` with `operand`, i.e. for all indices `O` in the
+ `operand` tensor:\
+ `output`[`O`] = `operand`[`O`]
+ - For every index `U` in the `updates` tensor and the corresponding index `O`
+ in the `operand` tensor:\
+ `output`[`O`] = `update_computation`(`output`[`O`], `updates`[`U`])
+
+The order in which updates are applied is non-deterministic. So, when multiple
+indices in `updates` refer to the same index in `operand`, the corresponding
+value in `output` will be non-deterministic.
+
+Note that the first parameter that is passed into the `update_computation` will
+always be the current value from the `output` tensor and the second parameter
+will always be the value from the `updates` tensor. This is important
+specifically for cases when the `update_computation` is _not commutative_.
+
+Informally, the scatter op can be viewed as an _inverse_ of the gather op, i.e.
+the scatter op updates the elements in the input that are extracted by the
+corresponding gather op.
+
+For a detailed informal description and examples, refer to the
+"Informal Description" section under `Gather`.
+
+## Select
+
+See also
+[`XlaBuilder::Select`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+Constructs an output array from elements of two input arrays, based on the
+values of a predicate array.
+
+<b> `Select(pred, on_true, on_false)` </b>
+
+Arguments | Type | Semantics
+---------- | ------- | ------------------
+`pred` | `XlaOp` | array of type PRED
+`on_true` | `XlaOp` | array of type T
+`on_false` | `XlaOp` | array of type T
+
+The arrays `on_true` and `on_false` must have the same shape. This is also the
+shape of the output array. The array `pred` must have the same dimensionality as
+`on_true` and `on_false`, with the `PRED` element type.
+
+For each element `P` of `pred`, the corresponding element of the output array is
+taken from `on_true` if the value of `P` is `true`, and from `on_false` if the
+value of `P` is `false`. As a restricted form of [broadcasting]
+(broadcasting.md), `pred` can be a scalar of type `PRED`. In this case, the
+output array is taken wholly from `on_true` if `pred` is `true`, and from
+`on_false` if `pred` is `false`.
+
+Example with non-scalar `pred`:
+
+```
+let pred: PRED[4] = {true, false, false, true};
+let v1: s32[4] = {1, 2, 3, 4};
+let v2: s32[4] = {100, 200, 300, 400};
+==>
+Select(pred, v1, v2) = s32[4]{1, 200, 300, 4};
+```
+
+Example with scalar `pred`:
+
+```
+let pred: PRED = true;
+let v1: s32[4] = {1, 2, 3, 4};
+let v2: s32[4] = {100, 200, 300, 400};
+==>
+Select(pred, v1, v2) = s32[4]{1, 2, 3, 4};
+```
+
+Selections between tuples are supported. Tuples are considered to be scalar
+types for this purpose. If `on_true` and `on_false` are tuples (which must have
+the same shape!) then `pred` has to be a scalar of type `PRED`.
+
+## SelectAndScatter
+
+See also
+[`XlaBuilder::SelectAndScatter`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+This operation can be considered as a composite operation that first computes
+`ReduceWindow` on the `operand` array to select an element from each window, and
+then scatters the `source` array to the indices of the selected elements to
+construct an output array with the same shape as the operand array. The binary
+`select` function is used to select an element from each window by applying it
+across each window, and it is called with the property that the first
+parameter's index vector is lexicographically less than the second parameter's
+index vector. The `select` function returns `true` if the first parameter is
+selected and returns `false` if the second parameter is selected, and the
+function must hold transitivity (i.e., if `select(a, b)` and `select(b, c)` are
+`true`, then `select(a, c)` is also `true`) so that the selected element does
+not depend on the order of the elements traversed for a given window.
+
+The function `scatter` is applied at each selected index in the output array. It
+takes two scalar parameters:
+
+1. Current value at the selected index in the output array
+2. The scatter value from `source` that applies to the selected index
+
+It combines the two parameters and returns a scalar value that's used to update
+the value at the selected index in the output array. Initially, all indices of
+the output array are set to `init_value`.
+
+The output array has the same shape as the `operand` array and the `source`
+array must have the same shape as the result of applying a `ReduceWindow`
+operation on the `operand` array. `SelectAndScatter` can be used to
+backpropagate the gradient values for a pooling layer in a neural network.
+
+<b>`SelectAndScatter(operand, select, window_dimensions, window_strides,
+padding, source, init_value, scatter)`</b>
+
+| Arguments | Type | Semantics |
+| ------------------- | ------------------- | -------------------------------- |
+| `operand` | `XlaOp` | array of type T over which the |
+: : : windows slide :
+| `select` | `XlaComputation` | binary computation of type `T, T |
+: : : -> PRED`, to apply to all :
+: : : elements in each window; returns :
+: : : `true` if the first parameter is :
+: : : selected and returns `false` if :
+: : : the second parameter is selected :
+| `window_dimensions` | `ArraySlice<int64>` | array of integers for window |
+: : : dimension values :
+| `window_strides` | `ArraySlice<int64>` | array of integers for window |
+: : : stride values :
+| `padding` | `Padding` | padding type for window |
+: : : (Padding\:\:kSame or :
+: : : Padding\:\:kValid) :
+| `source` | `XlaOp` | array of type T with the values |
+: : : to scatter :
+| `init_value` | `XlaOp` | scalar value of type T for the |
+: : : initial value of the output :
+: : : array :
+| `scatter` | `XlaComputation` | binary computation of type `T, T |
+: : : -> T`, to apply each scatter :
+: : : source element with its :
+: : : destination element :
+
+The figure below shows examples of using `SelectAndScatter`, with the `select`
+function computing the maximal value among its parameters. Note that when the
+windows overlap, as in the figure (2) below, an index of the `operand` array may
+be selected multiple times by different windows. In the figure, the element of
+value 9 is selected by both of the top windows (blue and red) and the binary
+addition `scatter` function produces the output element of value 8 (2 + 6).
+
+<div style="width:95%; margin:auto; margin-bottom:10px; margin-top:20px;">
+ <img style="width:100%"
+ src="https://www.tensorflow.org/images/ops_scatter_to_selected_window_element.png">
+</div>
+
+The evaluation order of the `scatter` function is arbitrary and may be
+non-deterministic. Therefore, the `scatter` function should not be overly
+sensitive to reassociation. See the discussion about associativity in the
+context of [`Reduce`](#reduce) for more details.
+
+## Send
+
+See also
+[`XlaBuilder::Send`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+<b> `Send(operand, channel_handle)` </b>
+
+Arguments | Type | Semantics
+---------------- | --------------- | -----------------------------------------
+`operand` | `XlaOp` | data to send (array of type T)
+`channel_handle` | `ChannelHandle` | unique identifier for each send/recv pair
+
+Sends the given operand data to a `Recv` instruction in another computation
+that shares the same channel handle. Does not return any data.
+
+Similar to the `Recv` operation, the client API of `Send` operation represents
+synchronous communication, and is internally decomposed into 2 HLO instructions
+(`Send` and `SendDone`) to enable asynchronous data transfers. See also
+[`HloInstruction::CreateSend` and `HloInstruction::CreateSendDone`](https://www.tensorflow.org/code/tensorflow/compiler/xla/service/hlo_instruction.h).
+
+<b>`Send(HloInstruction operand, int64 channel_id)`</b>
+
+Initiates an asynchronous transfer of the operand to the resources allocated by
+the `Recv` instruction with the same channel id. Returns a context, which is
+used by a following `SendDone` instruction to wait for the completion of the
+data transfer. The context is a tuple of {operand (shape), request identifier
+(U32)} and it can only be used by a `SendDone` instruction.
+
+<b> `SendDone(HloInstruction context)` </b>
+
+Given a context created by a `Send` instruction, waits for the data transfer to
+complete. The instruction does not return any data.
+
+<b> Scheduling of channel instructions </b>
+
+The execution order of the 4 instructions for each channel (`Recv`, `RecvDone`,
+`Send`, `SendDone`) is as below.
+
+<div style="width:95%; margin:auto; margin-bottom:10px; margin-top:20px;">
+ <img style="width:70%" src="../../images/send_recv_order.png">
+</div>
+
+* `Recv` happens before `Send`
+* `Send` happens before `RecvDone`
+* `Recv` happens before `RecvDone`
+* `Send` happens before `SendDone`
+
+When the backend compilers generate a linear schedule for each computation that
+communicates via channel instructions, there must not be cycles across the
+computations. For example, below schedules lead to deadlocks.
+
+<div style="width:95%; margin:auto; margin-bottom:10px; margin-top:20px;">
+ <img style="width:100%" src="../../images/send_recv_schedule.png">
+</div>
+
+## Slice
+
+See also
+[`XlaBuilder::Slice`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+Slicing extracts a sub-array from the input array. The sub-array is of the same
+rank as the input and contains the values inside a bounding box within the input
+array where the dimensions and indices of the bounding box are given as
+arguments to the slice operation.
+
+<b> `Slice(operand, start_indices, limit_indices)` </b>
+
+| Arguments | Type | Semantics |
+| --------------- | ------------------- | ------------------------------------ |
+| `operand` | `XlaOp` | N dimensional array of type T |
+| `start_indices` | `ArraySlice<int64>` | List of N integers containing the |
+: : : starting indices of the slice for :
+: : : each dimension. Values must be :
+: : : greater than or equal to zero. :
+| `limit_indices` | `ArraySlice<int64>` | List of N integers containing the |
+: : : ending indices (exclusive) for the :
+: : : slice for each dimension. Each value :
+: : : must be greater than or equal to the :
+: : : respective `start_indices` value for :
+: : : the dimension and less than or equal :
+: : : to the size of the dimension. :
+
+1-dimensional example:
+
+```
+let a = {0.0, 1.0, 2.0, 3.0, 4.0}
+Slice(a, {2}, {4}) produces:
+ {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} }
+
+Slice(b, {2, 1}, {4, 3}) produces:
+ { { 7.0, 8.0},
+ {10.0, 11.0} }
+```
+
+## Sort
+
+See also
+[`XlaBuilder::Sort`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+There are two versions of the Sort instruction: a single-operand and a
+two-operand version.
+
+<b>`Sort(operand)`</b>
+
+Arguments | Type | Semantics
+----------- | ------- | --------------------
+`operand` | `XlaOp` | The operand to sort.
+`dimension` | `int64` | The dimension along which to sort.
+
+Sorts the elements in the operand in ascending order along the provided
+dimension. For example, for a rank-2 (matrix) operand, a `dimension` value of 0
+will sort each column independently, and a `dimension` value of 1 will sort each
+row independently. If the operand's elements have floating point type, and the
+operand contains NaN elements, the order of elements in the output is
+implementation-defined.
+
+<b>`Sort(key, value)`</b>
+
+Sorts both the key and the value operands. The keys are sorted as in the
+single-operand version. The values are sorted according to the order of their
+corresponding keys. For example, if the inputs are `keys = [3, 1]` and
+`values = [42, 50]`, then the output of the sort is the tuple
+`{[1, 3], [50, 42]}`.
+
+The sort is not guaranteed to be stable, that is, if the keys array contains
+duplicates, the order of their corresponding values may not be preserved.
+
+Arguments | Type | Semantics
+----------- | ------- | -------------------
+`keys` | `XlaOp` | The sort keys.
+`values` | `XlaOp` | The values to sort.
+`dimension` | `int64` | The dimension along which to sort.
+
+The `keys` and `values` must have the same dimensions, but may have different
+element types.
+
+## Transpose
+
+See also the `tf.reshape` operation.
+
+<b>`Transpose(operand)`</b>
+
+Arguments | Type | Semantics
+------------- | ------------------- | ------------------------------
+`operand` | `XlaOp` | The operand to transpose.
+`permutation` | `ArraySlice<int64>` | How to permute the dimensions.
+
+
+Permutes the operand dimensions with the given permutation, so
+`∀ i . 0 ≤ i < rank ⇒ input_dimensions[permutation[i]] = output_dimensions[i]`.
+
+This is the same as Reshape(operand, permutation,
+ Permute(permutation, operand.shape.dimensions)).
+
+## Tuple
+
+See also
+[`XlaBuilder::Tuple`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+A tuple containing a variable number of data handles, each of which has its own
+shape.
+
+This is analogous to `std::tuple` in C++. Conceptually:
+
+```
+let v: f32[10] = f32[10]{0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
+let s: s32 = 5;
+let t: (f32[10], s32) = tuple(v, s);
+```
+
+Tuples can be deconstructed (accessed) via the [`GetTupleElement`]
+(#gettupleelement) operation.
+
+## While
+
+See also
+[`XlaBuilder::While`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h).
+
+<b> `While(condition, body, init)` </b>
+
+| Arguments | Type | Semantics |
+| ----------- | ---------------- | ---------------------------------------- |
+| `condition` | `XlaComputation` | XlaComputation of type `T -> PRED` which |
+: : : defines the termination condition of the :
+: : : loop. :
+| `body` | `XlaComputation` | XlaComputation of type `T -> T` which |
+: : : defines the body of the loop. :
+| `init` | `T` | Initial value for the parameter of |
+: : : `condition` and `body`. :
+
+Sequentially executes the `body` until the `condition` fails. This is similar to
+a typical while loop in many other languages except for the differences and
+restrictions listed below.
+
+* A `While` node returns a value of type `T`, which is the result from the
+ last execution of the `body`.
+* The shape of the type `T` is statically determined and must be the same
+ across all iterations.
+
+The T parameters of the computations are initialized with the `init` value in
+the first iteration and are automatically updated to the new result from `body`
+in each subsequent iteration.
+
+One main use case of the `While` node is to implement the repeated execution of
+training in neural networks. Simplified pseudocode is shown below with a graph
+that represents the computation. The code can be found in
+[`while_test.cc`](https://www.tensorflow.org/code/tensorflow/compiler/xla/tests/while_test.cc).
+The type `T` in this example is a `Tuple` consisting of an `int32` for the
+iteration count and a `vector[10]` for the accumulator. For 1000 iterations, the
+loop keeps adding a constant vector to the accumulator.
+
+```
+// Pseudocode for the computation.
+init = {0, zero_vector[10]} // Tuple of int32 and float[10].
+result = init;
+while (result(0) < 1000) {
+ iteration = result(0) + 1;
+ new_vector = result(1) + constant_vector[10];
+ result = {iteration, new_vector};
+}
+```
+
+<div style="width:95%; margin:auto; margin-bottom:10px; margin-top:20px;">
+ <img style="width:100%" src="https://www.tensorflow.org/images/ops_while.png">
+</div>
diff --git a/tensorflow/examples/android/BUILD b/tensorflow/examples/android/BUILD
index f327b64..f5f0d7c 100644
--- a/tensorflow/examples/android/BUILD
+++ b/tensorflow/examples/android/BUILD
@@ -68,6 +68,7 @@
srcs = glob([
"src/**/*.java",
]),
+ aapt_version = "aapt",
# Package assets from assets dir as well as all model targets. Remove undesired models
# (and corresponding Activities in source) to reduce APK size.
assets = [
diff --git a/tensorflow/examples/get_started/regression/test.py b/tensorflow/examples/get_started/regression/test.py
index 0b1477a..bb4db67 100644
--- a/tensorflow/examples/get_started/regression/test.py
+++ b/tensorflow/examples/get_started/regression/test.py
@@ -29,7 +29,7 @@
sys.modules["imports85"] = imports85
# pylint: disable=g-bad-import-order,g-import-not-at-top
-import tensorflow.contrib.data as data
+import tensorflow.data as data
import tensorflow.examples.get_started.regression.dnn_regression as dnn_regression
import tensorflow.examples.get_started.regression.linear_regression as linear_regression
diff --git a/tensorflow/go/op/wrappers.go b/tensorflow/go/op/wrappers.go
index 9dd487e..b4d4db3 100644
--- a/tensorflow/go/op/wrappers.go
+++ b/tensorflow/go/op/wrappers.go
@@ -2562,92 +2562,6 @@
return op.Output(0)
}
-// EditDistanceAttr is an optional argument to EditDistance.
-type EditDistanceAttr func(optionalAttr)
-
-// EditDistanceNormalize sets the optional normalize attribute to value.
-//
-// value: boolean (if true, edit distances are normalized by length of truth).
-//
-// The output is:
-// If not specified, defaults to true
-func EditDistanceNormalize(value bool) EditDistanceAttr {
- return func(m optionalAttr) {
- m["normalize"] = value
- }
-}
-
-// Computes the (possibly normalized) Levenshtein Edit Distance.
-//
-// The inputs are variable-length sequences provided by SparseTensors
-// (hypothesis_indices, hypothesis_values, hypothesis_shape)
-// and
-// (truth_indices, truth_values, truth_shape).
-//
-// The inputs are:
-//
-// Arguments:
-// hypothesis_indices: The indices of the hypothesis list SparseTensor.
-// This is an N x R int64 matrix.
-// hypothesis_values: The values of the hypothesis list SparseTensor.
-// This is an N-length vector.
-// hypothesis_shape: The shape of the hypothesis list SparseTensor.
-// This is an R-length vector.
-// truth_indices: The indices of the truth list SparseTensor.
-// This is an M x R int64 matrix.
-// truth_values: The values of the truth list SparseTensor.
-// This is an M-length vector.
-// truth_shape: truth indices, vector.
-//
-// Returns A dense float tensor with rank R - 1.
-//
-// For the example input:
-//
-// // hypothesis represents a 2x1 matrix with variable-length values:
-// // (0,0) = ["a"]
-// // (1,0) = ["b"]
-// hypothesis_indices = [[0, 0, 0],
-// [1, 0, 0]]
-// hypothesis_values = ["a", "b"]
-// hypothesis_shape = [2, 1, 1]
-//
-// // truth represents a 2x2 matrix with variable-length values:
-// // (0,0) = []
-// // (0,1) = ["a"]
-// // (1,0) = ["b", "c"]
-// // (1,1) = ["a"]
-// truth_indices = [[0, 1, 0],
-// [1, 0, 0],
-// [1, 0, 1],
-// [1, 1, 0]]
-// truth_values = ["a", "b", "c", "a"]
-// truth_shape = [2, 2, 2]
-// normalize = true
-//
-// The output will be:
-//
-// // output is a 2x2 matrix with edit distances normalized by truth lengths.
-// output = [[inf, 1.0], // (0,0): no truth, (0,1): no hypothesis
-// [0.5, 1.0]] // (1,0): addition, (1,1): no hypothesis
-func EditDistance(scope *Scope, hypothesis_indices tf.Output, hypothesis_values tf.Output, hypothesis_shape tf.Output, truth_indices tf.Output, truth_values tf.Output, truth_shape tf.Output, optional ...EditDistanceAttr) (output tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "EditDistance",
- Input: []tf.Input{
- hypothesis_indices, hypothesis_values, hypothesis_shape, truth_indices, truth_values, truth_shape,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
// Reverses specific dimensions of a tensor.
//
// Given a `tensor`, and a `bool` tensor `dims` representing the dimensions
@@ -3828,27 +3742,6 @@
return op.Output(0)
}
-// Retrieves the tree ensemble resource stamp token, number of trees and growing statistics.
-//
-// Arguments:
-// tree_ensemble_handle: Handle to the tree ensemble.
-//
-// Returns Stamp token of the tree ensemble resource.The number of trees in the tree ensemble resource.The number of trees that were finished successfully.The number of layers we attempted to build (but not necessarily succeeded).Rank size 2 tensor that contains start and end ids of the nodes in the latest
-// layer.
-func BoostedTreesGetEnsembleStates(scope *Scope, tree_ensemble_handle tf.Output) (stamp_token tf.Output, num_trees tf.Output, num_finalized_trees tf.Output, num_attempted_layers tf.Output, last_layer_nodes_range tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "BoostedTreesGetEnsembleStates",
- Input: []tf.Input{
- tree_ensemble_handle,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0), op.Output(1), op.Output(2), op.Output(3), op.Output(4)
-}
-
// Creates a tree ensemble model and returns a handle to it.
//
// Arguments:
@@ -3890,66 +3783,596 @@
return op.Output(0)
}
-// Computes the sum along sparse segments of a tensor.
+// BoostedTreesEnsembleResourceHandleOpAttr is an optional argument to BoostedTreesEnsembleResourceHandleOp.
+type BoostedTreesEnsembleResourceHandleOpAttr func(optionalAttr)
+
+// BoostedTreesEnsembleResourceHandleOpContainer sets the optional container attribute to value.
+// If not specified, defaults to ""
+func BoostedTreesEnsembleResourceHandleOpContainer(value string) BoostedTreesEnsembleResourceHandleOpAttr {
+ return func(m optionalAttr) {
+ m["container"] = value
+ }
+}
+
+// BoostedTreesEnsembleResourceHandleOpSharedName sets the optional shared_name attribute to value.
+// If not specified, defaults to ""
+func BoostedTreesEnsembleResourceHandleOpSharedName(value string) BoostedTreesEnsembleResourceHandleOpAttr {
+ return func(m optionalAttr) {
+ m["shared_name"] = value
+ }
+}
+
+// Creates a handle to a BoostedTreesEnsembleResource
+func BoostedTreesEnsembleResourceHandleOp(scope *Scope, optional ...BoostedTreesEnsembleResourceHandleOpAttr) (resource tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "BoostedTreesEnsembleResourceHandleOp",
+
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// ComputeAccidentalHitsAttr is an optional argument to ComputeAccidentalHits.
+type ComputeAccidentalHitsAttr func(optionalAttr)
+
+// ComputeAccidentalHitsSeed sets the optional seed attribute to value.
//
-// Read
-// [the section on segmentation](https://tensorflow.org/api_guides/python/math_ops#Segmentation)
-// for an explanation of segments.
+// 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 ComputeAccidentalHitsSeed(value int64) ComputeAccidentalHitsAttr {
+ return func(m optionalAttr) {
+ m["seed"] = value
+ }
+}
+
+// ComputeAccidentalHitsSeed2 sets the optional seed2 attribute to value.
//
-// Like `SegmentSum`, but `segment_ids` can have rank less than `data`'s first
-// dimension, selecting a subset of dimension 0, specified by `indices`.
+// value: An second seed to avoid seed collision.
+// If not specified, defaults to 0
+func ComputeAccidentalHitsSeed2(value int64) ComputeAccidentalHitsAttr {
+ return func(m optionalAttr) {
+ m["seed2"] = value
+ }
+}
+
+// Computes the ids of the positions in sampled_candidates that match true_labels.
//
-// For example:
-//
-// ```python
-// c = tf.constant([[1,2,3,4], [-1,-2,-3,-4], [5,6,7,8]])
-//
-// # Select two rows, one segment.
-// tf.sparse_segment_sum(c, tf.constant([0, 1]), tf.constant([0, 0]))
-// # => [[0 0 0 0]]
-//
-// # Select two rows, two segment.
-// tf.sparse_segment_sum(c, tf.constant([0, 1]), tf.constant([0, 1]))
-// # => [[ 1 2 3 4]
-// # [-1 -2 -3 -4]]
-//
-// # Select all rows, two segments.
-// tf.sparse_segment_sum(c, tf.constant([0, 1, 2]), tf.constant([0, 0, 1]))
-// # => [[0 0 0 0]
-// # [5 6 7 8]]
-//
-// # Which is equivalent to:
-// tf.segment_sum(c, tf.constant([0, 0, 1]))
-// ```
+// When doing log-odds NCE, the result of this op should be passed through a
+// SparseToDense op, then added to the logits of the sampled candidates. This has
+// the effect of 'removing' the sampled labels that match the true labels by
+// making the classifier sure that they are sampled labels.
//
// Arguments:
+// true_classes: The true_classes output of UnpackSparseLabels.
+// sampled_candidates: The sampled_candidates output of CandidateSampler.
+// num_true: Number of true labels per context.
//
-// indices: A 1-D tensor. Has same rank as `segment_ids`.
-// segment_ids: A 1-D tensor. Values should be sorted and can be repeated.
+// Returns A vector of indices corresponding to rows of true_candidates.A vector of IDs of positions in sampled_candidates that match a true_label
+// for the row with the corresponding index in indices.A vector of the same length as indices and ids, in which each element
+// is -FLOAT_MAX.
+func ComputeAccidentalHits(scope *Scope, true_classes tf.Output, sampled_candidates tf.Output, num_true int64, optional ...ComputeAccidentalHitsAttr) (indices tf.Output, ids tf.Output, weights tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"num_true": num_true}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "ComputeAccidentalHits",
+ Input: []tf.Input{
+ true_classes, sampled_candidates,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0), op.Output(1), op.Output(2)
+}
+
+// FixedUnigramCandidateSamplerAttr is an optional argument to FixedUnigramCandidateSampler.
+type FixedUnigramCandidateSamplerAttr func(optionalAttr)
+
+// FixedUnigramCandidateSamplerVocabFile sets the optional vocab_file attribute to value.
//
-// Returns Has same shape as data, except for dimension 0 which
-// has size `k`, the number of segments.
-func SparseSegmentSum(scope *Scope, data tf.Output, indices tf.Output, segment_ids tf.Output) (output tf.Output) {
+// value: Each valid line in this file (which should have a CSV-like format)
+// corresponds to a valid word ID. IDs are in sequential order, starting from
+// num_reserved_ids. The last entry in each line is expected to be a value
+// corresponding to the count or relative probability. Exactly one of vocab_file
+// and unigrams needs to be passed to this op.
+// If not specified, defaults to ""
+func FixedUnigramCandidateSamplerVocabFile(value string) FixedUnigramCandidateSamplerAttr {
+ return func(m optionalAttr) {
+ m["vocab_file"] = value
+ }
+}
+
+// FixedUnigramCandidateSamplerDistortion sets the optional distortion attribute to value.
+//
+// value: The distortion is used to skew the unigram probability distribution.
+// Each weight is first raised to the distortion's power before adding to the
+// internal unigram distribution. As a result, distortion = 1.0 gives regular
+// unigram sampling (as defined by the vocab file), and distortion = 0.0 gives
+// a uniform distribution.
+// If not specified, defaults to 1
+func FixedUnigramCandidateSamplerDistortion(value float32) FixedUnigramCandidateSamplerAttr {
+ return func(m optionalAttr) {
+ m["distortion"] = value
+ }
+}
+
+// FixedUnigramCandidateSamplerNumReservedIds sets the optional num_reserved_ids attribute to value.
+//
+// value: Optionally some reserved IDs can be added in the range [0,
+// ..., num_reserved_ids) by the users. One use case is that a special unknown
+// word token is used as ID 0. These IDs will have a sampling probability of 0.
+// If not specified, defaults to 0
+func FixedUnigramCandidateSamplerNumReservedIds(value int64) FixedUnigramCandidateSamplerAttr {
+ return func(m optionalAttr) {
+ m["num_reserved_ids"] = value
+ }
+}
+
+// FixedUnigramCandidateSamplerNumShards sets the optional num_shards attribute to value.
+//
+// value: A sampler can be used to sample from a subset of the original range
+// in order to speed up the whole computation through parallelism. This parameter
+// (together with 'shard') indicates the number of partitions that are being
+// used in the overall computation.
+// If not specified, defaults to 1
+//
+// REQUIRES: value >= 1
+func FixedUnigramCandidateSamplerNumShards(value int64) FixedUnigramCandidateSamplerAttr {
+ return func(m optionalAttr) {
+ m["num_shards"] = value
+ }
+}
+
+// FixedUnigramCandidateSamplerShard sets the optional shard attribute to value.
+//
+// value: A sampler can be used to sample from a subset of the original range
+// in order to speed up the whole computation through parallelism. This parameter
+// (together with 'num_shards') indicates the particular partition number of a
+// sampler op, when partitioning is being used.
+// If not specified, defaults to 0
+//
+// REQUIRES: value >= 0
+func FixedUnigramCandidateSamplerShard(value int64) FixedUnigramCandidateSamplerAttr {
+ return func(m optionalAttr) {
+ m["shard"] = value
+ }
+}
+
+// FixedUnigramCandidateSamplerUnigrams sets the optional unigrams attribute to value.
+//
+// value: A list of unigram counts or probabilities, one per ID in sequential
+// order. Exactly one of vocab_file and unigrams should be passed to this op.
+// If not specified, defaults to <>
+func FixedUnigramCandidateSamplerUnigrams(value []float32) FixedUnigramCandidateSamplerAttr {
+ return func(m optionalAttr) {
+ m["unigrams"] = value
+ }
+}
+
+// FixedUnigramCandidateSamplerSeed 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 FixedUnigramCandidateSamplerSeed(value int64) FixedUnigramCandidateSamplerAttr {
+ return func(m optionalAttr) {
+ m["seed"] = value
+ }
+}
+
+// FixedUnigramCandidateSamplerSeed2 sets the optional seed2 attribute to value.
+//
+// value: An second seed to avoid seed collision.
+// If not specified, defaults to 0
+func FixedUnigramCandidateSamplerSeed2(value int64) FixedUnigramCandidateSamplerAttr {
+ return func(m optionalAttr) {
+ m["seed2"] = value
+ }
+}
+
+// Generates labels for candidate sampling with a learned unigram distribution.
+//
+// A unigram sampler could use a fixed unigram distribution read from a
+// file or passed in as an in-memory array instead of building up the distribution
+// from data on the fly. There is also an option to skew the distribution by
+// applying a distortion power to the weights.
+//
+// The vocabulary file should be in CSV-like format, with the last field
+// being the weight associated with the word.
+//
+// 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 FixedUnigramCandidateSampler(scope *Scope, true_classes tf.Output, num_true int64, num_sampled int64, unique bool, range_max int64, optional ...FixedUnigramCandidateSamplerAttr) (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: "FixedUnigramCandidateSampler",
+ Input: []tf.Input{
+ true_classes,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0), op.Output(1), op.Output(2)
+}
+
+// LogUniformCandidateSamplerAttr is an optional argument to LogUniformCandidateSampler.
+type LogUniformCandidateSamplerAttr func(optionalAttr)
+
+// LogUniformCandidateSamplerSeed 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 LogUniformCandidateSamplerSeed(value int64) LogUniformCandidateSamplerAttr {
+ return func(m optionalAttr) {
+ m["seed"] = value
+ }
+}
+
+// LogUniformCandidateSamplerSeed2 sets the optional seed2 attribute to value.
+//
+// value: An second seed to avoid seed collision.
+// If not specified, defaults to 0
+func LogUniformCandidateSamplerSeed2(value int64) LogUniformCandidateSamplerAttr {
+ return func(m optionalAttr) {
+ m["seed2"] = value
+ }
+}
+
+// Generates labels for candidate sampling with a log-uniform 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 LogUniformCandidateSampler(scope *Scope, true_classes tf.Output, num_true int64, num_sampled int64, unique bool, range_max int64, optional ...LogUniformCandidateSamplerAttr) (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: "LogUniformCandidateSampler",
+ Input: []tf.Input{
+ true_classes,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0), op.Output(1), op.Output(2)
+}
+
+// UniformCandidateSamplerAttr is an optional argument to UniformCandidateSampler.
+type UniformCandidateSamplerAttr func(optionalAttr)
+
+// UniformCandidateSamplerSeed 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 UniformCandidateSamplerSeed(value int64) UniformCandidateSamplerAttr {
+ return func(m optionalAttr) {
+ m["seed"] = value
+ }
+}
+
+// UniformCandidateSamplerSeed2 sets the optional seed2 attribute to value.
+//
+// value: An second seed to avoid seed collision.
+// If not specified, defaults to 0
+func UniformCandidateSamplerSeed2(value int64) UniformCandidateSamplerAttr {
+ return func(m optionalAttr) {
+ m["seed2"] = value
+ }
+}
+
+// Generates labels for candidate sampling with a uniform 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 UniformCandidateSampler(scope *Scope, true_classes tf.Output, num_true int64, num_sampled int64, unique bool, range_max int64, optional ...UniformCandidateSamplerAttr) (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: "UniformCandidateSampler",
+ Input: []tf.Input{
+ true_classes,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0), op.Output(1), op.Output(2)
+}
+
+// GenerateVocabRemappingAttr is an optional argument to GenerateVocabRemapping.
+type GenerateVocabRemappingAttr func(optionalAttr)
+
+// GenerateVocabRemappingOldVocabSize sets the optional old_vocab_size attribute to value.
+//
+// value: Number of entries in the old vocab file to consider. If -1,
+// use the entire old vocabulary.
+// If not specified, defaults to -1
+//
+// REQUIRES: value >= -1
+func GenerateVocabRemappingOldVocabSize(value int64) GenerateVocabRemappingAttr {
+ return func(m optionalAttr) {
+ m["old_vocab_size"] = value
+ }
+}
+
+// Given a path to new and old vocabulary files, returns a remapping Tensor of
+//
+// length `num_new_vocab`, where `remapping[i]` contains the row number in the old
+// vocabulary that corresponds to row `i` in the new vocabulary (starting at line
+// `new_vocab_offset` and up to `num_new_vocab` entities), or `-1` if entry `i`
+// in the new vocabulary is not in the old vocabulary. The old vocabulary is
+// constrained to the first `old_vocab_size` entries if `old_vocab_size` is not the
+// default value of -1.
+//
+// `num_vocab_offset` enables
+// use in the partitioned variable case, and should generally be set through
+// examining partitioning info. The format of the files should be a text file,
+// with each line containing a single entity within the vocabulary.
+//
+// For example, with `new_vocab_file` a text file containing each of the following
+// elements on a single line: `[f0, f1, f2, f3]`, old_vocab_file = [f1, f0, f3],
+// `num_new_vocab = 3, new_vocab_offset = 1`, the returned remapping would be
+// `[0, -1, 2]`.
+//
+// The op also returns a count of how many entries in the new vocabulary
+// were present in the old vocabulary, which is used to calculate the number of
+// values to initialize in a weight matrix remapping
+//
+// This functionality can be used to remap both row vocabularies (typically,
+// features) and column vocabularies (typically, classes) from TensorFlow
+// checkpoints. Note that the partitioning logic relies on contiguous vocabularies
+// corresponding to div-partitioned variables. Moreover, the underlying remapping
+// uses an IndexTable (as opposed to an inexact CuckooTable), so client code should
+// use the corresponding index_table_from_file() as the FeatureColumn framework
+// does (as opposed to tf.feature_to_id(), which uses a CuckooTable).
+//
+// Arguments:
+// new_vocab_file: Path to the new vocab file.
+// old_vocab_file: Path to the old vocab file.
+// new_vocab_offset: How many entries into the new vocab file to start reading.
+// num_new_vocab: Number of entries in the new vocab file to remap.
+//
+// Returns A Tensor of length num_new_vocab where the element at index i
+// is equal to the old ID that maps to the new ID i. This element is -1 for any
+// new ID that is not found in the old vocabulary.Number of new vocab entries found in old vocab.
+func GenerateVocabRemapping(scope *Scope, new_vocab_file tf.Output, old_vocab_file tf.Output, new_vocab_offset int64, num_new_vocab int64, optional ...GenerateVocabRemappingAttr) (remapping tf.Output, num_present tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"new_vocab_offset": new_vocab_offset, "num_new_vocab": num_new_vocab}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "GenerateVocabRemapping",
+ Input: []tf.Input{
+ new_vocab_file, old_vocab_file,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0), op.Output(1)
+}
+
+// Broadcasts a tensor value to one or more other devices.
+func CollectiveBcastSend(scope *Scope, input tf.Output, group_size int64, group_key int64, instance_key int64, shape tf.Shape) (data tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"group_size": group_size, "group_key": group_key, "instance_key": instance_key, "shape": shape}
+ opspec := tf.OpSpec{
+ Type: "CollectiveBcastSend",
+ Input: []tf.Input{
+ input,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// Mutually reduces multiple tensors of identical type and shape.
+func CollectiveReduce(scope *Scope, input tf.Output, group_size int64, group_key int64, instance_key int64, merge_op string, final_op string, subdiv_offsets []int64) (data tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"group_size": group_size, "group_key": group_key, "instance_key": instance_key, "merge_op": merge_op, "final_op": final_op, "subdiv_offsets": subdiv_offsets}
+ opspec := tf.OpSpec{
+ Type: "CollectiveReduce",
+ Input: []tf.Input{
+ input,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// AbortAttr is an optional argument to Abort.
+type AbortAttr func(optionalAttr)
+
+// AbortErrorMsg sets the optional error_msg attribute to value.
+//
+// value: A string which is the message associated with the exception.
+// If not specified, defaults to ""
+func AbortErrorMsg(value string) AbortAttr {
+ return func(m optionalAttr) {
+ m["error_msg"] = value
+ }
+}
+
+// AbortExitWithoutError sets the optional exit_without_error attribute to value.
+// If not specified, defaults to false
+func AbortExitWithoutError(value bool) AbortAttr {
+ return func(m optionalAttr) {
+ m["exit_without_error"] = value
+ }
+}
+
+// Raise a exception to abort the process when called.
+//
+// If exit_without_error is true, the process will exit normally,
+// otherwise it will exit with a SIGABORT signal.
+//
+// Returns nothing but an exception.
+//
+// Returns the created operation.
+func Abort(scope *Scope, optional ...AbortAttr) (o *tf.Operation) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "Abort",
+
+ Attrs: attrs,
+ }
+ return scope.AddOperation(opspec)
+}
+
+// Forwards the input to the output.
+//
+// This operator represents the loop termination condition used by the
+// "pivot" switches of a loop.
+//
+// Arguments:
+// input: A boolean scalar, representing the branch predicate of the Switch op.
+//
+// Returns The same tensor as `input`.
+func LoopCond(scope *Scope, input tf.Output) (output tf.Output) {
if scope.Err() != nil {
return
}
opspec := tf.OpSpec{
- Type: "SparseSegmentSum",
+ Type: "LoopCond",
Input: []tf.Input{
- data, indices, segment_ids,
+ input,
},
}
op := scope.AddOperation(opspec)
return op.Output(0)
}
-// Computes hyperbolic sine of x element-wise.
-func Sinh(scope *Scope, x tf.Output) (y tf.Output) {
+// Returns a tensor of zeros with the same shape and type as x.
+//
+// Arguments:
+// x: a tensor of type T.
+//
+// Returns a tensor of the same shape and type as x but filled with zeros.
+func ZerosLike(scope *Scope, x tf.Output) (y tf.Output) {
if scope.Err() != nil {
return
}
opspec := tf.OpSpec{
- Type: "Sinh",
+ Type: "ZerosLike",
Input: []tf.Input{
x,
},
@@ -3958,101 +4381,15 @@
return op.Output(0)
}
-// Computes the minimum along segments of a tensor.
-//
-// Read
-// [the section on segmentation](https://tensorflow.org/api_guides/python/math_ops#segmentation)
-// for an explanation of segments.
-//
-// This operator is similar to the unsorted segment sum operator found
-// [(here)](../../../api_docs/python/math_ops.md#UnsortedSegmentSum).
-// Instead of computing the sum over segments, it computes the minimum such that:
-//
-// \\(output_i = \min_{j...} data_[j...]\\) where min is over tuples `j...` such
-// that `segment_ids[j...] == i`.
-//
-// If the minimum is empty for a given segment ID `i`, it outputs the largest
-// possible value for the specific numeric type,
-// `output[i] = numeric_limits<T>::max()`.
-//
-// If the given segment ID `i` is negative, then the corresponding value is
-// dropped, and will not be included in the result.
-//
-// Arguments:
-//
-// segment_ids: A tensor whose shape is a prefix of `data.shape`.
-//
-//
-// Returns Has same shape as data, except for the first `segment_ids.rank`
-// dimensions, which are replaced with a single dimension which has size
-// `num_segments`.
-func UnsortedSegmentMin(scope *Scope, data tf.Output, segment_ids tf.Output, num_segments tf.Output) (output tf.Output) {
+// Returns a copy of the input tensor.
+func Snapshot(scope *Scope, input tf.Output) (output tf.Output) {
if scope.Err() != nil {
return
}
opspec := tf.OpSpec{
- Type: "UnsortedSegmentMin",
+ Type: "Snapshot",
Input: []tf.Input{
- data, segment_ids, num_segments,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// Computes rectified linear 6: `min(max(features, 0), 6)`.
-func Relu6(scope *Scope, features tf.Output) (activations tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "Relu6",
- Input: []tf.Input{
- features,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// Computes the sum along segments of a tensor.
-//
-// Read
-// [the section on segmentation](https://tensorflow.org/api_guides/python/math_ops#Segmentation)
-// for an explanation of segments.
-//
-// Computes a tensor such that
-// \\(output[i] = \sum_{j...} data[j...]\\) where the sum is over tuples `j...` such
-// that `segment_ids[j...] == i`. Unlike `SegmentSum`, `segment_ids`
-// need not be sorted and need not cover all values in the full
-// range of valid values.
-//
-// If the sum is empty for a given segment ID `i`, `output[i] = 0`.
-// If the given segment ID `i` is negative, the value is dropped and will not be
-// added to the sum of the segment.
-//
-// `num_segments` should equal the number of distinct segment IDs.
-//
-// <div style="width:70%; margin:auto; margin-bottom:10px; margin-top:20px;">
-// <img style="width:100%" src="https://www.tensorflow.org/images/UnsortedSegmentSum.png" alt>
-// </div>
-//
-// Arguments:
-//
-// segment_ids: A tensor whose shape is a prefix of `data.shape`.
-//
-//
-// Returns Has same shape as data, except for the first `segment_ids.rank`
-// dimensions, which are replaced with a single dimension which has size
-// `num_segments`.
-func UnsortedSegmentSum(scope *Scope, data tf.Output, segment_ids tf.Output, num_segments tf.Output) (output tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "UnsortedSegmentSum",
- Input: []tf.Input{
- data, segment_ids, num_segments,
+ input,
},
}
op := scope.AddOperation(opspec)
@@ -4413,6 +4750,162 @@
return op.Output(0)
}
+// EditDistanceAttr is an optional argument to EditDistance.
+type EditDistanceAttr func(optionalAttr)
+
+// EditDistanceNormalize sets the optional normalize attribute to value.
+//
+// value: boolean (if true, edit distances are normalized by length of truth).
+//
+// The output is:
+// If not specified, defaults to true
+func EditDistanceNormalize(value bool) EditDistanceAttr {
+ return func(m optionalAttr) {
+ m["normalize"] = value
+ }
+}
+
+// Computes the (possibly normalized) Levenshtein Edit Distance.
+//
+// The inputs are variable-length sequences provided by SparseTensors
+// (hypothesis_indices, hypothesis_values, hypothesis_shape)
+// and
+// (truth_indices, truth_values, truth_shape).
+//
+// The inputs are:
+//
+// Arguments:
+// hypothesis_indices: The indices of the hypothesis list SparseTensor.
+// This is an N x R int64 matrix.
+// hypothesis_values: The values of the hypothesis list SparseTensor.
+// This is an N-length vector.
+// hypothesis_shape: The shape of the hypothesis list SparseTensor.
+// This is an R-length vector.
+// truth_indices: The indices of the truth list SparseTensor.
+// This is an M x R int64 matrix.
+// truth_values: The values of the truth list SparseTensor.
+// This is an M-length vector.
+// truth_shape: truth indices, vector.
+//
+// Returns A dense float tensor with rank R - 1.
+//
+// For the example input:
+//
+// // hypothesis represents a 2x1 matrix with variable-length values:
+// // (0,0) = ["a"]
+// // (1,0) = ["b"]
+// hypothesis_indices = [[0, 0, 0],
+// [1, 0, 0]]
+// hypothesis_values = ["a", "b"]
+// hypothesis_shape = [2, 1, 1]
+//
+// // truth represents a 2x2 matrix with variable-length values:
+// // (0,0) = []
+// // (0,1) = ["a"]
+// // (1,0) = ["b", "c"]
+// // (1,1) = ["a"]
+// truth_indices = [[0, 1, 0],
+// [1, 0, 0],
+// [1, 0, 1],
+// [1, 1, 0]]
+// truth_values = ["a", "b", "c", "a"]
+// truth_shape = [2, 2, 2]
+// normalize = true
+//
+// The output will be:
+//
+// // output is a 2x2 matrix with edit distances normalized by truth lengths.
+// output = [[inf, 1.0], // (0,0): no truth, (0,1): no hypothesis
+// [0.5, 1.0]] // (1,0): addition, (1,1): no hypothesis
+func EditDistance(scope *Scope, hypothesis_indices tf.Output, hypothesis_values tf.Output, hypothesis_shape tf.Output, truth_indices tf.Output, truth_values tf.Output, truth_shape tf.Output, optional ...EditDistanceAttr) (output tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "EditDistance",
+ Input: []tf.Input{
+ hypothesis_indices, hypothesis_values, hypothesis_shape, truth_indices, truth_values, truth_shape,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// DepthwiseConv2dNativeBackpropInputAttr is an optional argument to DepthwiseConv2dNativeBackpropInput.
+type DepthwiseConv2dNativeBackpropInputAttr func(optionalAttr)
+
+// DepthwiseConv2dNativeBackpropInputDataFormat 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, height, width, channels].
+// Alternatively, the format could be "NCHW", the data storage order of:
+// [batch, channels, height, width].
+// If not specified, defaults to "NHWC"
+func DepthwiseConv2dNativeBackpropInputDataFormat(value string) DepthwiseConv2dNativeBackpropInputAttr {
+ return func(m optionalAttr) {
+ m["data_format"] = value
+ }
+}
+
+// DepthwiseConv2dNativeBackpropInputDilations sets the optional dilations attribute to value.
+//
+// value: 1-D tensor of length 4. The dilation factor for each dimension of
+// `input`. If set to k > 1, there will be k-1 skipped cells between each filter
+// element on that dimension. The dimension order is determined by the value of
+// `data_format`, see above for details. Dilations in the batch and depth
+// dimensions must be 1.
+// If not specified, defaults to <i:1 i:1 i:1 i:1 >
+func DepthwiseConv2dNativeBackpropInputDilations(value []int64) DepthwiseConv2dNativeBackpropInputAttr {
+ return func(m optionalAttr) {
+ m["dilations"] = value
+ }
+}
+
+// Computes the gradients of depthwise convolution with respect to the input.
+//
+// Arguments:
+// input_sizes: An integer vector representing the shape of `input`, based
+// on `data_format`. For example, if `data_format` is 'NHWC' then
+// `input` is a 4-D `[batch, height, width, channels]` tensor.
+// filter: 4-D with shape
+// `[filter_height, filter_width, in_channels, depthwise_multiplier]`.
+// out_backprop: 4-D with shape based on `data_format`.
+// For example, if `data_format` is 'NHWC' then
+// out_backprop shape is `[batch, out_height, out_width, out_channels]`.
+// Gradients w.r.t. the output of the convolution.
+// strides: The stride of the sliding window for each dimension of the input
+// of the convolution.
+// padding: The type of padding algorithm to use.
+//
+// Returns 4-D with shape according to `data_format`. For example, if
+// `data_format` is 'NHWC', output shape is `[batch, in_height,
+// in_width, in_channels]`. Gradient w.r.t. the input of the
+// convolution.
+func DepthwiseConv2dNativeBackpropInput(scope *Scope, input_sizes tf.Output, filter tf.Output, out_backprop tf.Output, strides []int64, padding string, optional ...DepthwiseConv2dNativeBackpropInputAttr) (output tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"strides": strides, "padding": padding}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "DepthwiseConv2dNativeBackpropInput",
+ Input: []tf.Input{
+ input_sizes, filter, out_backprop,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
// ApproximateEqualAttr is an optional argument to ApproximateEqual.
type ApproximateEqualAttr func(optionalAttr)
@@ -4581,6 +5074,78 @@
return op.Output(0), op.Output(1), op.Output(2)
}
+// AllCandidateSamplerAttr is an optional argument to AllCandidateSampler.
+type AllCandidateSamplerAttr func(optionalAttr)
+
+// AllCandidateSamplerSeed 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 AllCandidateSamplerSeed(value int64) AllCandidateSamplerAttr {
+ return func(m optionalAttr) {
+ m["seed"] = value
+ }
+}
+
+// AllCandidateSamplerSeed2 sets the optional seed2 attribute to value.
+//
+// value: An second seed to avoid seed collision.
+// If not specified, defaults to 0
+func AllCandidateSamplerSeed2(value int64) AllCandidateSamplerAttr {
+ 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 produce.
+// 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.
+//
+// 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 AllCandidateSampler(scope *Scope, true_classes tf.Output, num_true int64, num_sampled int64, unique bool, optional ...AllCandidateSamplerAttr) (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}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "AllCandidateSampler",
+ Input: []tf.Input{
+ true_classes,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0), op.Output(1), op.Output(2)
+}
+
// Returns x + y element-wise.
//
// *NOTE*: `Add` supports broadcasting. `AddN` does not. More about broadcasting
@@ -4599,21 +5164,6 @@
return op.Output(0)
}
-// Computes exponential of x element-wise. \\(y = e^x\\).
-func Exp(scope *Scope, x tf.Output) (y tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "Exp",
- Input: []tf.Input{
- x,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
// Returns an element-wise indication of the sign of a number.
//
// `y = sign(x) = -1` if `x < 0`; 0 if `x == 0`; 1 if `x > 0`.
@@ -4768,104 +5318,6 @@
return op.Output(0)
}
-// Computes the maximum along segments of a tensor.
-//
-// Read
-// [the section on segmentation](https://tensorflow.org/api_guides/python/math_ops#Segmentation)
-// for an explanation of segments.
-//
-// This operator is similar to the unsorted segment sum operator found
-// [(here)](../../../api_docs/python/math_ops.md#UnsortedSegmentSum).
-// Instead of computing the sum over segments, it computes the maximum such that:
-//
-// \\(output_i = \max_{j...} data[j...]\\) where max is over tuples `j...` such
-// that `segment_ids[j...] == i`.
-//
-// If the maximum is empty for a given segment ID `i`, it outputs the smallest
-// possible value for the specific numeric type,
-// `output[i] = numeric_limits<T>::lowest()`.
-//
-// If the given segment ID `i` is negative, then the corresponding value is
-// dropped, and will not be included in the result.
-//
-// <div style="width:70%; margin:auto; margin-bottom:10px; margin-top:20px;">
-// <img style="width:100%" src="https://www.tensorflow.org/images/UnsortedSegmentMax.png" alt>
-// </div>
-//
-// Arguments:
-//
-// segment_ids: A tensor whose shape is a prefix of `data.shape`.END
-// }
-// out_arg {
-// name: "output"
-// description: <<END
-// Has same shape as data, except for the first `segment_ids.rank`
-// dimensions, which are replaced with a single dimension which has size
-// `num_segments`.
-//
-func UnsortedSegmentMax(scope *Scope, data tf.Output, segment_ids tf.Output, num_segments tf.Output) (output tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "UnsortedSegmentMax",
- Input: []tf.Input{
- data, segment_ids, num_segments,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// NthElementAttr is an optional argument to NthElement.
-type NthElementAttr func(optionalAttr)
-
-// NthElementReverse sets the optional reverse attribute to value.
-//
-// value: When set to True, find the nth-largest value in the vector and vice
-// versa.
-// If not specified, defaults to false
-func NthElementReverse(value bool) NthElementAttr {
- return func(m optionalAttr) {
- m["reverse"] = value
- }
-}
-
-// Finds values of the `n`-th order statistic for the last dimension.
-//
-// If the input is a vector (rank-1), finds the entries which is the nth-smallest
-// value in the vector and outputs their values as scalar tensor.
-//
-// For matrices (resp. higher rank input), computes the entries which is the
-// nth-smallest value in each row (resp. vector along the last dimension). Thus,
-//
-// values.shape = input.shape[:-1]
-//
-// Arguments:
-// input: 1-D or higher with last dimension at least `n+1`.
-// n: 0-D. Position of sorted vector to select along the last dimension (along
-// each row for matrices). Valid range of n is `[0, input.shape[:-1])`
-//
-// Returns The `n`-th order statistic along each last dimensional slice.
-func NthElement(scope *Scope, input tf.Output, n tf.Output, optional ...NthElementAttr) (values tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "NthElement",
- Input: []tf.Input{
- input, n,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
// Computes the sum along sparse segments of a tensor.
//
// Like `SparseSegmentSum`, but allows missing ids in `segment_ids`. If an id is
@@ -5142,6 +5594,74 @@
return op.Output(0)
}
+// Computes hyperbolic sine of x element-wise.
+func Sinh(scope *Scope, x tf.Output) (y tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "Sinh",
+ Input: []tf.Input{
+ x,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// Computes the sum along sparse segments of a tensor.
+//
+// Read
+// [the section on segmentation](https://tensorflow.org/api_guides/python/math_ops#Segmentation)
+// for an explanation of segments.
+//
+// Like `SegmentSum`, but `segment_ids` can have rank less than `data`'s first
+// dimension, selecting a subset of dimension 0, specified by `indices`.
+//
+// For example:
+//
+// ```python
+// c = tf.constant([[1,2,3,4], [-1,-2,-3,-4], [5,6,7,8]])
+//
+// # Select two rows, one segment.
+// tf.sparse_segment_sum(c, tf.constant([0, 1]), tf.constant([0, 0]))
+// # => [[0 0 0 0]]
+//
+// # Select two rows, two segment.
+// tf.sparse_segment_sum(c, tf.constant([0, 1]), tf.constant([0, 1]))
+// # => [[ 1 2 3 4]
+// # [-1 -2 -3 -4]]
+//
+// # Select all rows, two segments.
+// tf.sparse_segment_sum(c, tf.constant([0, 1, 2]), tf.constant([0, 0, 1]))
+// # => [[0 0 0 0]
+// # [5 6 7 8]]
+//
+// # Which is equivalent to:
+// tf.segment_sum(c, tf.constant([0, 0, 1]))
+// ```
+//
+// Arguments:
+//
+// indices: A 1-D tensor. Has same rank as `segment_ids`.
+// segment_ids: A 1-D tensor. Values should be sorted and can be repeated.
+//
+// Returns Has same shape as data, except for dimension 0 which
+// has size `k`, the number of segments.
+func SparseSegmentSum(scope *Scope, data tf.Output, indices tf.Output, segment_ids tf.Output) (output tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "SparseSegmentSum",
+ Input: []tf.Input{
+ data, indices, segment_ids,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
// Computes natural logarithm of (1 + x) element-wise.
//
// I.e., \\(y = \log_e (1 + x)\\).
@@ -6701,6 +7221,63 @@
return components
}
+// Computes rectified linear 6: `min(max(features, 0), 6)`.
+func Relu6(scope *Scope, features tf.Output) (activations tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "Relu6",
+ Input: []tf.Input{
+ features,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// Computes the minimum along segments of a tensor.
+//
+// Read
+// [the section on segmentation](https://tensorflow.org/api_guides/python/math_ops#segmentation)
+// for an explanation of segments.
+//
+// This operator is similar to the unsorted segment sum operator found
+// [(here)](../../../api_docs/python/math_ops.md#UnsortedSegmentSum).
+// Instead of computing the sum over segments, it computes the minimum such that:
+//
+// \\(output_i = \min_{j...} data_[j...]\\) where min is over tuples `j...` such
+// that `segment_ids[j...] == i`.
+//
+// If the minimum is empty for a given segment ID `i`, it outputs the largest
+// possible value for the specific numeric type,
+// `output[i] = numeric_limits<T>::max()`.
+//
+// If the given segment ID `i` is negative, then the corresponding value is
+// dropped, and will not be included in the result.
+//
+// Arguments:
+//
+// segment_ids: A tensor whose shape is a prefix of `data.shape`.
+//
+//
+// Returns Has same shape as data, except for the first `segment_ids.rank`
+// dimensions, which are replaced with a single dimension which has size
+// `num_segments`.
+func UnsortedSegmentMin(scope *Scope, data tf.Output, segment_ids tf.Output, num_segments tf.Output) (output tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "UnsortedSegmentMin",
+ Input: []tf.Input{
+ data, segment_ids, num_segments,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
// Computes rectified linear gradients for a Relu operation.
//
// Arguments:
@@ -7674,6 +8251,44 @@
return op.Output(0)
}
+// Bucketizes 'input' based on 'boundaries'.
+//
+// For example, if the inputs are
+// boundaries = [0, 10, 100]
+// input = [[-5, 10000]
+// [150, 10]
+// [5, 100]]
+//
+// then the output will be
+// output = [[0, 3]
+// [3, 2]
+// [1, 3]]
+//
+// Arguments:
+// input: Any shape of Tensor contains with int or float type.
+// boundaries: A sorted list of floats gives the boundary of the buckets.
+//
+// Returns Same shape with 'input', each value of input replaced with bucket index.
+//
+// @compatibility(numpy)
+// Equivalent to np.digitize.
+// @end_compatibility
+func Bucketize(scope *Scope, input tf.Output, boundaries []float32) (output tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"boundaries": boundaries}
+ opspec := tf.OpSpec{
+ Type: "Bucketize",
+ Input: []tf.Input{
+ input,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
// FusedBatchNormV2Attr is an optional argument to FusedBatchNormV2.
type FusedBatchNormV2Attr func(optionalAttr)
@@ -8790,6 +9405,119 @@
return op.Output(0)
}
+// Computes exponential of x element-wise. \\(y = e^x\\).
+func Exp(scope *Scope, x tf.Output) (y tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "Exp",
+ Input: []tf.Input{
+ x,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// NthElementAttr is an optional argument to NthElement.
+type NthElementAttr func(optionalAttr)
+
+// NthElementReverse sets the optional reverse attribute to value.
+//
+// value: When set to True, find the nth-largest value in the vector and vice
+// versa.
+// If not specified, defaults to false
+func NthElementReverse(value bool) NthElementAttr {
+ return func(m optionalAttr) {
+ m["reverse"] = value
+ }
+}
+
+// Finds values of the `n`-th order statistic for the last dimension.
+//
+// If the input is a vector (rank-1), finds the entries which is the nth-smallest
+// value in the vector and outputs their values as scalar tensor.
+//
+// For matrices (resp. higher rank input), computes the entries which is the
+// nth-smallest value in each row (resp. vector along the last dimension). Thus,
+//
+// values.shape = input.shape[:-1]
+//
+// Arguments:
+// input: 1-D or higher with last dimension at least `n+1`.
+// n: 0-D. Position of sorted vector to select along the last dimension (along
+// each row for matrices). Valid range of n is `[0, input.shape[:-1])`
+//
+// Returns The `n`-th order statistic along each last dimensional slice.
+func NthElement(scope *Scope, input tf.Output, n tf.Output, optional ...NthElementAttr) (values tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "NthElement",
+ Input: []tf.Input{
+ input, n,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// Computes the maximum along segments of a tensor.
+//
+// Read
+// [the section on segmentation](https://tensorflow.org/api_guides/python/math_ops#Segmentation)
+// for an explanation of segments.
+//
+// This operator is similar to the unsorted segment sum operator found
+// [(here)](../../../api_docs/python/math_ops.md#UnsortedSegmentSum).
+// Instead of computing the sum over segments, it computes the maximum such that:
+//
+// \\(output_i = \max_{j...} data[j...]\\) where max is over tuples `j...` such
+// that `segment_ids[j...] == i`.
+//
+// If the maximum is empty for a given segment ID `i`, it outputs the smallest
+// possible value for the specific numeric type,
+// `output[i] = numeric_limits<T>::lowest()`.
+//
+// If the given segment ID `i` is negative, then the corresponding value is
+// dropped, and will not be included in the result.
+//
+// <div style="width:70%; margin:auto; margin-bottom:10px; margin-top:20px;">
+// <img style="width:100%" src="https://www.tensorflow.org/images/UnsortedSegmentMax.png" alt>
+// </div>
+//
+// Arguments:
+//
+// segment_ids: A tensor whose shape is a prefix of `data.shape`.END
+// }
+// out_arg {
+// name: "output"
+// description: <<END
+// Has same shape as data, except for the first `segment_ids.rank`
+// dimensions, which are replaced with a single dimension which has size
+// `num_segments`.
+//
+func UnsortedSegmentMax(scope *Scope, data tf.Output, segment_ids tf.Output, num_segments tf.Output) (output tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "UnsortedSegmentMax",
+ Input: []tf.Input{
+ data, segment_ids, num_segments,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
// Transforms a vector of brain.Example protos (as strings) into typed tensors.
//
// Arguments:
@@ -9491,56 +10219,62 @@
return op.Output(0)
}
-// ResourceApplyFtrlAttr is an optional argument to ResourceApplyFtrl.
-type ResourceApplyFtrlAttr func(optionalAttr)
+// FusedResizeAndPadConv2DAttr is an optional argument to FusedResizeAndPadConv2D.
+type FusedResizeAndPadConv2DAttr func(optionalAttr)
-// ResourceApplyFtrlUseLocking sets the optional use_locking attribute to value.
+// FusedResizeAndPadConv2DResizeAlignCorners sets the optional resize_align_corners 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.
+// value: If true, the centers of the 4 corner pixels of the input and output tensors are
+// aligned, preserving the values at the corner pixels. Defaults to false.
// If not specified, defaults to false
-func ResourceApplyFtrlUseLocking(value bool) ResourceApplyFtrlAttr {
+func FusedResizeAndPadConv2DResizeAlignCorners(value bool) FusedResizeAndPadConv2DAttr {
return func(m optionalAttr) {
- m["use_locking"] = value
+ m["resize_align_corners"] = value
}
}
-// Update '*var' according to the Ftrl-proximal scheme.
+// Performs a resize and padding as a preprocess during a convolution.
//
-// 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
+// It's often possible to do spatial transformations more efficiently as part of
+// the packing stage of a convolution, so this op allows for an optimized
+// implementation where these stages are fused together. This prevents the need to
+// write out the intermediate results as whole tensors, reducing memory pressure,
+// and we can get some latency gains by merging the transformation calculations.
+// The data_format attribute for Conv2D isn't supported by this op, and defaults to
+// 'NHWC' order.
+// Internally this op uses a single per-graph scratch buffer, which means that it
+// will block if multiple versions are being run in parallel. This is because this
+// operator is primarily an optimization to minimize memory usage.
//
// 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 regulariation. Must be a scalar.
-// lr_power: Scaling factor. Must be a scalar.
+// input: 4-D with shape `[batch, in_height, in_width, in_channels]`.
+// size: A 1-D int32 Tensor of 2 elements: `new_height, new_width`. The
+// new size for the images.
+// paddings: A two-column matrix specifying the padding sizes. The number of
+// rows must be the same as the rank of `input`.
+// filter: 4-D with shape
+// `[filter_height, filter_width, in_channels, out_channels]`.
//
-// Returns the created operation.
-func ResourceApplyFtrl(scope *Scope, var_ tf.Output, accum tf.Output, linear tf.Output, grad tf.Output, lr tf.Output, l1 tf.Output, l2 tf.Output, lr_power tf.Output, optional ...ResourceApplyFtrlAttr) (o *tf.Operation) {
+// strides: 1-D of length 4. The stride of the sliding window for each dimension
+// of `input`. Must be in the same order as the dimension specified with format.
+// padding: The type of padding algorithm to use.
+func FusedResizeAndPadConv2D(scope *Scope, input tf.Output, size tf.Output, paddings tf.Output, filter tf.Output, mode string, strides []int64, padding string, optional ...FusedResizeAndPadConv2DAttr) (output tf.Output) {
if scope.Err() != nil {
return
}
- attrs := map[string]interface{}{}
+ attrs := map[string]interface{}{"mode": mode, "strides": strides, "padding": padding}
for _, a := range optional {
a(attrs)
}
opspec := tf.OpSpec{
- Type: "ResourceApplyFtrl",
+ Type: "FusedResizeAndPadConv2D",
Input: []tf.Input{
- var_, accum, linear, grad, lr, l1, l2, lr_power,
+ input, size, paddings, filter,
},
Attrs: attrs,
}
- return scope.AddOperation(opspec)
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
}
// RandomUniformAttr is an optional argument to RandomUniform.
@@ -9597,6 +10331,58 @@
return op.Output(0)
}
+// ResourceApplyFtrlAttr is an optional argument to ResourceApplyFtrl.
+type ResourceApplyFtrlAttr func(optionalAttr)
+
+// ResourceApplyFtrlUseLocking 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 ResourceApplyFtrlUseLocking(value bool) ResourceApplyFtrlAttr {
+ return func(m optionalAttr) {
+ m["use_locking"] = value
+ }
+}
+
+// Update '*var' according to the Ftrl-proximal scheme.
+//
+// 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.
+// lr: Scaling factor. Must be a scalar.
+// l1: L1 regulariation. Must be a scalar.
+// l2: L2 regulariation. Must be a scalar.
+// lr_power: Scaling factor. Must be a scalar.
+//
+// Returns the created operation.
+func ResourceApplyFtrl(scope *Scope, var_ tf.Output, accum tf.Output, linear tf.Output, grad tf.Output, lr tf.Output, l1 tf.Output, l2 tf.Output, lr_power tf.Output, optional ...ResourceApplyFtrlAttr) (o *tf.Operation) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "ResourceApplyFtrl",
+ Input: []tf.Input{
+ var_, accum, linear, grad, lr, l1, l2, lr_power,
+ },
+ Attrs: attrs,
+ }
+ return scope.AddOperation(opspec)
+}
+
// Encode audio data using the WAV file format.
//
// This operation will generate a string suitable to be saved out to create a .wav
@@ -9733,23 +10519,6 @@
return scope.AddOperation(opspec)
}
-// Broadcasts a tensor value to one or more other devices.
-func CollectiveBcastSend(scope *Scope, input tf.Output, group_size int64, group_key int64, instance_key int64, shape tf.Shape) (data tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"group_size": group_size, "group_key": group_key, "instance_key": instance_key, "shape": shape}
- opspec := tf.OpSpec{
- Type: "CollectiveBcastSend",
- Input: []tf.Input{
- input,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
// Split a `SparseTensor` into `num_split` tensors along one dimension.
//
// If the `shape[split_dim]` is not an integer multiple of `num_split`. Slices
@@ -9873,6 +10642,118 @@
return scope.AddOperation(opspec)
}
+// Calculates gains for each feature and returns the best possible split information for the feature.
+//
+// The split information is the best threshold (bucket id), gains and left/right node contributions per node for each feature.
+//
+// It is possible that not all nodes can be split on each feature. Hence, the list of possible nodes can differ between the features. Therefore, we return `node_ids_list` for each feature, containing the list of nodes that this feature can be used to split.
+//
+// In this manner, the output is the best split per features and per node, so that it needs to be combined later to produce the best split for each node (among all possible features).
+//
+// The length of output lists are all of the same length, `num_features`.
+// The output shapes are compatible in a way that the first dimension of all tensors of all lists are the same and equal to the number of possible split nodes for each feature.
+//
+// Arguments:
+// node_id_range: A Rank 1 tensor (shape=[2]) to specify the range [first, last) of node ids to process within `stats_summary_list`. The nodes are iterated between the two nodes specified by the tensor, as like `for node_id in range(node_id_range[0], node_id_range[1])` (Note that the last index node_id_range[1] is exclusive).
+// stats_summary_list: A list of Rank 3 tensor (#shape=[max_splits, bucket, 2]) for accumulated stats summary (gradient/hessian) per node per buckets for each feature. The first dimension of the tensor is the maximum number of splits, and thus not all elements of it will be used, but only the indexes specified by node_ids will be used.
+// l1: l1 regularization factor on leaf weights, per instance based.
+// l2: l2 regularization factor on leaf weights, per instance based.
+// tree_complexity: adjustment to the gain, per leaf based.
+// min_node_weight: mininum avg of hessians in a node before required for the node to be considered for splitting.
+// max_splits: the number of nodes that can be split in the whole tree. Used as a dimension of output tensors.
+//
+// Returns An output list of Rank 1 tensors indicating possible split node ids for each feature. The length of the list is num_features, but each tensor has different size as each feature provides different possible nodes. See above for details like shapes and sizes.An output list of Rank 1 tensors indicating the best gains for each feature to split for certain nodes. See above for details like shapes and sizes.An output list of Rank 1 tensors indicating the bucket id to compare with (as a threshold) for split in each node. See above for details like shapes and sizes.A list of Rank 2 tensors indicating the contribution of the left nodes when branching from parent nodes (given by the tensor element in the output node_ids_list) to the left direction by the given threshold for each feature. This value will be used to make the left node value by adding to the parent node value. Second dimension size is 1 for 1-dimensional logits, but would be larger for multi-class problems. See above for details like shapes and sizes.A list of Rank 2 tensors, with the same shape/conditions as left_node_contribs_list, but just that the value is for the right node.
+func BoostedTreesCalculateBestGainsPerFeature(scope *Scope, node_id_range tf.Output, stats_summary_list []tf.Output, l1 tf.Output, l2 tf.Output, tree_complexity tf.Output, min_node_weight tf.Output, max_splits int64) (node_ids_list []tf.Output, gains_list []tf.Output, thresholds_list []tf.Output, left_node_contribs_list []tf.Output, right_node_contribs_list []tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"max_splits": max_splits}
+ opspec := tf.OpSpec{
+ Type: "BoostedTreesCalculateBestGainsPerFeature",
+ Input: []tf.Input{
+ node_id_range, tf.OutputList(stats_summary_list), l1, l2, tree_complexity, min_node_weight,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ if scope.Err() != nil {
+ return
+ }
+ var idx int
+ var err error
+ if node_ids_list, idx, err = makeOutputList(op, idx, "node_ids_list"); err != nil {
+ scope.UpdateErr("BoostedTreesCalculateBestGainsPerFeature", err)
+ return
+ }
+ if gains_list, idx, err = makeOutputList(op, idx, "gains_list"); err != nil {
+ scope.UpdateErr("BoostedTreesCalculateBestGainsPerFeature", err)
+ return
+ }
+ if thresholds_list, idx, err = makeOutputList(op, idx, "thresholds_list"); err != nil {
+ scope.UpdateErr("BoostedTreesCalculateBestGainsPerFeature", err)
+ return
+ }
+ if left_node_contribs_list, idx, err = makeOutputList(op, idx, "left_node_contribs_list"); err != nil {
+ scope.UpdateErr("BoostedTreesCalculateBestGainsPerFeature", err)
+ return
+ }
+ if right_node_contribs_list, idx, err = makeOutputList(op, idx, "right_node_contribs_list"); err != nil {
+ scope.UpdateErr("BoostedTreesCalculateBestGainsPerFeature", err)
+ return
+ }
+ return node_ids_list, gains_list, thresholds_list, left_node_contribs_list, right_node_contribs_list
+}
+
+// EncodePngAttr is an optional argument to EncodePng.
+type EncodePngAttr func(optionalAttr)
+
+// EncodePngCompression sets the optional compression attribute to value.
+//
+// value: Compression level.
+// If not specified, defaults to -1
+func EncodePngCompression(value int64) EncodePngAttr {
+ return func(m optionalAttr) {
+ m["compression"] = value
+ }
+}
+
+// PNG-encode an image.
+//
+// `image` is a 3-D uint8 or uint16 Tensor of shape `[height, width, channels]`
+// where `channels` is:
+//
+// * 1: for grayscale.
+// * 2: for grayscale + alpha.
+// * 3: for RGB.
+// * 4: for RGBA.
+//
+// The ZLIB compression level, `compression`, can be -1 for the PNG-encoder
+// default or a value from 0 to 9. 9 is the highest compression level, generating
+// the smallest output, but is slower.
+//
+// Arguments:
+// image: 3-D with shape `[height, width, channels]`.
+//
+// Returns 0-D. PNG-encoded image.
+func EncodePng(scope *Scope, image tf.Output, optional ...EncodePngAttr) (contents tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "EncodePng",
+ Input: []tf.Input{
+ image,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
// DataFormatVecPermuteAttr is an optional argument to DataFormatVecPermute.
type DataFormatVecPermuteAttr func(optionalAttr)
@@ -10109,6 +10990,112 @@
return op.Output(0)
}
+// This op consumes a lock created by `MutexLock`.
+//
+// This op exists to consume a tensor created by `MutexLock` (other than
+// direct control dependencies). It should be the only that consumes the tensor,
+// and will raise an error if it is not. Its only purpose is to keep the
+// mutex lock tensor alive until it is consumed by this op.
+//
+// **NOTE**: This operation must run on the same device as its input. This may
+// be enforced via the `colocate_with` mechanism.
+//
+// Arguments:
+// mutex_lock: A tensor returned by `MutexLock`.
+//
+// Returns the created operation.
+func ConsumeMutexLock(scope *Scope, mutex_lock tf.Output) (o *tf.Operation) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "ConsumeMutexLock",
+ Input: []tf.Input{
+ mutex_lock,
+ },
+ }
+ return scope.AddOperation(opspec)
+}
+
+// ResourceScatterNdAddAttr is an optional argument to ResourceScatterNdAdd.
+type ResourceScatterNdAddAttr func(optionalAttr)
+
+// ResourceScatterNdAddUseLocking sets the optional use_locking attribute to value.
+//
+// value: An optional bool. Defaults to True. If True, the assignment will
+// be protected by a lock; otherwise the behavior is undefined,
+// but may exhibit less contention.
+// If not specified, defaults to true
+func ResourceScatterNdAddUseLocking(value bool) ResourceScatterNdAddAttr {
+ return func(m optionalAttr) {
+ m["use_locking"] = value
+ }
+}
+
+// Adds sparse `updates` to individual values or slices within a given
+//
+// variable according to `indices`.
+//
+// `ref` is a `Tensor` with rank `P` and `indices` is a `Tensor` of rank `Q`.
+//
+// `indices` must be integer tensor, containing indices into `ref`.
+// 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`.
+//
+// `updates` is `Tensor` of rank `Q-1+P-K` with shape:
+//
+// ```
+// [d_0, ..., d_{Q-2}, ref.shape[K], ..., ref.shape[P-1]].
+// ```
+//
+// For example, say we want to update 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], use_resource=True)
+// indices = tf.constant([[4], [3], [1] ,[7]])
+// updates = tf.constant([9, 10, 11, 12])
+// update = tf.scatter_nd_add(ref, indices, updates)
+// with tf.Session() as sess:
+// print sess.run(update)
+// ```
+//
+// The resulting update to ref would look like this:
+//
+// [1, 12, 3, 14, 14, 6, 7, 20]
+//
+// See `tf.scatter_nd` for more details about how to make updates to
+// slices.
+//
+// Arguments:
+// ref: A resource handle. Must be from a VarHandleOp.
+// indices: A Tensor. Must be one of the following types: int32, int64.
+// A tensor of indices into ref.
+// updates: A Tensor. Must have the same type as ref. A tensor of
+// values to add to ref.
+//
+// Returns the created operation.
+func ResourceScatterNdAdd(scope *Scope, ref tf.Output, indices tf.Output, updates tf.Output, optional ...ResourceScatterNdAddAttr) (o *tf.Operation) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "ResourceScatterNdAdd",
+ Input: []tf.Input{
+ ref, indices, updates,
+ },
+ Attrs: attrs,
+ }
+ return scope.AddOperation(opspec)
+}
+
// Updates the tree ensemble by either adding a layer to the last tree being grown
//
// or by starting a new tree.
@@ -10850,102 +11837,6 @@
return scope.AddOperation(opspec)
}
-// ResourceScatterNdAddAttr is an optional argument to ResourceScatterNdAdd.
-type ResourceScatterNdAddAttr func(optionalAttr)
-
-// ResourceScatterNdAddUseLocking sets the optional use_locking attribute to value.
-//
-// value: An optional bool. Defaults to True. If True, the assignment will
-// be protected by a lock; otherwise the behavior is undefined,
-// but may exhibit less contention.
-// If not specified, defaults to true
-func ResourceScatterNdAddUseLocking(value bool) ResourceScatterNdAddAttr {
- return func(m optionalAttr) {
- m["use_locking"] = value
- }
-}
-
-// Adds sparse `updates` to individual values or slices within a given
-//
-// variable according to `indices`.
-//
-// `ref` is a `Tensor` with rank `P` and `indices` is a `Tensor` of rank `Q`.
-//
-// `indices` must be integer tensor, containing indices into `ref`.
-// 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`.
-//
-// `updates` is `Tensor` of rank `Q-1+P-K` with shape:
-//
-// ```
-// [d_0, ..., d_{Q-2}, ref.shape[K], ..., ref.shape[P-1]].
-// ```
-//
-// For example, say we want to update 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], use_resource=True)
-// indices = tf.constant([[4], [3], [1] ,[7]])
-// updates = tf.constant([9, 10, 11, 12])
-// update = tf.scatter_nd_add(ref, indices, updates)
-// with tf.Session() as sess:
-// print sess.run(update)
-// ```
-//
-// The resulting update to ref would look like this:
-//
-// [1, 12, 3, 14, 14, 6, 7, 20]
-//
-// See `tf.scatter_nd` for more details about how to make updates to
-// slices.
-//
-// Arguments:
-// ref: A resource handle. Must be from a VarHandleOp.
-// indices: A Tensor. Must be one of the following types: int32, int64.
-// A tensor of indices into ref.
-// updates: A Tensor. Must have the same type as ref. A tensor of
-// values to add to ref.
-//
-// Returns the created operation.
-func ResourceScatterNdAdd(scope *Scope, ref tf.Output, indices tf.Output, updates tf.Output, optional ...ResourceScatterNdAddAttr) (o *tf.Operation) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "ResourceScatterNdAdd",
- Input: []tf.Input{
- ref, indices, updates,
- },
- Attrs: attrs,
- }
- return scope.AddOperation(opspec)
-}
-
-// Mutually reduces multiple tensors of identical type and shape.
-func CollectiveReduce(scope *Scope, input tf.Output, group_size int64, group_key int64, instance_key int64, merge_op string, final_op string, subdiv_offsets []int64) (data tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"group_size": group_size, "group_key": group_key, "instance_key": instance_key, "merge_op": merge_op, "final_op": final_op, "subdiv_offsets": subdiv_offsets}
- opspec := tf.OpSpec{
- Type: "CollectiveReduce",
- Input: []tf.Input{
- input,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
// StatelessRandomNormalAttr is an optional argument to StatelessRandomNormal.
type StatelessRandomNormalAttr func(optionalAttr)
@@ -10989,6 +11880,143 @@
return op.Output(0)
}
+// Creates a sequence of numbers.
+//
+// This operation creates a sequence of numbers that begins at `start` and
+// extends by increments of `delta` up to but not including `limit`.
+//
+// For example:
+//
+// ```
+// # 'start' is 3
+// # 'limit' is 18
+// # 'delta' is 3
+// tf.range(start, limit, delta) ==> [3, 6, 9, 12, 15]
+// ```
+//
+// Arguments:
+// start: 0-D (scalar). First entry in the sequence.
+// limit: 0-D (scalar). Upper limit of sequence, exclusive.
+// delta: 0-D (scalar). Optional. Default is 1. Number that increments `start`.
+//
+// Returns 1-D.
+func Range(scope *Scope, start tf.Output, limit tf.Output, delta tf.Output) (output tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "Range",
+ Input: []tf.Input{
+ start, limit, delta,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// ResourceApplyMomentumAttr is an optional argument to ResourceApplyMomentum.
+type ResourceApplyMomentumAttr func(optionalAttr)
+
+// ResourceApplyMomentumUseLocking 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 ResourceApplyMomentumUseLocking(value bool) ResourceApplyMomentumAttr {
+ return func(m optionalAttr) {
+ m["use_locking"] = value
+ }
+}
+
+// ResourceApplyMomentumUseNesterov sets the optional use_nesterov attribute to value.
+//
+// value: If `True`, the tensor passed to compute grad will be
+// var - lr * momentum * accum, so in the end, the var you get is actually
+// var - lr * momentum * accum.
+// If not specified, defaults to false
+func ResourceApplyMomentumUseNesterov(value bool) ResourceApplyMomentumAttr {
+ return func(m optionalAttr) {
+ m["use_nesterov"] = value
+ }
+}
+
+// Update '*var' according to the momentum scheme. Set use_nesterov = True if you
+//
+// want to use Nesterov momentum.
+//
+// accum = accum * momentum + grad
+// var -= lr * accum
+//
+// Arguments:
+// var_: Should be from a Variable().
+// accum: Should be from a Variable().
+// lr: Scaling factor. Must be a scalar.
+// grad: The gradient.
+// momentum: Momentum. Must be a scalar.
+//
+// Returns the created operation.
+func ResourceApplyMomentum(scope *Scope, var_ tf.Output, accum tf.Output, lr tf.Output, grad tf.Output, momentum tf.Output, optional ...ResourceApplyMomentumAttr) (o *tf.Operation) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "ResourceApplyMomentum",
+ Input: []tf.Input{
+ var_, accum, lr, grad, momentum,
+ },
+ Attrs: attrs,
+ }
+ return scope.AddOperation(opspec)
+}
+
+// Exits the current frame to its parent frame.
+//
+// Exit makes its input `data` available to the parent frame.
+//
+// Arguments:
+// data: The tensor to be made available to the parent frame.
+//
+// Returns The same tensor as `data`.
+func Exit(scope *Scope, data tf.Output) (output tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "Exit",
+ Input: []tf.Input{
+ data,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// Produce a string tensor that encodes the state of a Reader.
+//
+// Not all Readers support being serialized, so this can produce an
+// Unimplemented error.
+//
+// Arguments:
+// reader_handle: Handle to a Reader.
+func ReaderSerializeStateV2(scope *Scope, reader_handle tf.Output) (state tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "ReaderSerializeStateV2",
+ Input: []tf.Input{
+ reader_handle,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
// DepthwiseConv2dNativeBackpropFilterAttr is an optional argument to DepthwiseConv2dNativeBackpropFilter.
type DepthwiseConv2dNativeBackpropFilterAttr func(optionalAttr)
@@ -11122,68 +12150,6 @@
return op.Output(0)
}
-// StringSplitV2Attr is an optional argument to StringSplitV2.
-type StringSplitV2Attr func(optionalAttr)
-
-// StringSplitV2Maxsplit sets the optional maxsplit attribute to value.
-//
-// value: An `int`. If `maxsplit > 0`, limit of the split of the result.
-// If not specified, defaults to -1
-func StringSplitV2Maxsplit(value int64) StringSplitV2Attr {
- return func(m optionalAttr) {
- m["maxsplit"] = value
- }
-}
-
-// Split elements of `source` based on `sep` into a `SparseTensor`.
-//
-// Let N be the size of source (typically N will be the batch size). Split each
-// element of `source` based on `sep` and return a `SparseTensor`
-// containing the split tokens. Empty tokens are ignored.
-//
-// For example, N = 2, source[0] is 'hello world' and source[1] is 'a b c',
-// then the output will be
-// ```
-// st.indices = [0, 0;
-// 0, 1;
-// 1, 0;
-// 1, 1;
-// 1, 2]
-// st.shape = [2, 3]
-// st.values = ['hello', 'world', 'a', 'b', 'c']
-// ```
-//
-// If `sep` is given, consecutive delimiters are not grouped together and are
-// deemed to delimit empty strings. For example, source of `"1<>2<><>3"` and
-// sep of `"<>"` returns `["1", "2", "", "3"]`. If `sep` is None or an empty
-// string, consecutive whitespace are regarded as a single separator, and the
-// result will contain no empty strings at the startor end if the string has
-// leading or trailing whitespace.
-//
-// Note that the above mentioned behavior matches python's str.split.
-//
-// Arguments:
-// input: `1-D` string `Tensor`, the strings to split.
-// sep: `0-D` string `Tensor`, the delimiter character.
-func StringSplitV2(scope *Scope, input tf.Output, sep tf.Output, optional ...StringSplitV2Attr) (indices tf.Output, values tf.Output, shape tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "StringSplitV2",
- Input: []tf.Input{
- input, sep,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0), op.Output(1), op.Output(2)
-}
-
// MaxPoolAttr is an optional argument to MaxPool.
type MaxPoolAttr func(optionalAttr)
@@ -11664,6 +12630,51 @@
return op.Output(0)
}
+// ProdAttr is an optional argument to Prod.
+type ProdAttr func(optionalAttr)
+
+// ProdKeepDims sets the optional keep_dims attribute to value.
+//
+// value: If true, retain reduced dimensions with length 1.
+// If not specified, defaults to false
+func ProdKeepDims(value bool) ProdAttr {
+ return func(m optionalAttr) {
+ m["keep_dims"] = value
+ }
+}
+
+// Computes the product of elements across dimensions of a tensor.
+//
+// Reduces `input` along the dimensions given in `axis`. Unless
+// `keep_dims` is true, the rank of the tensor is reduced by 1 for each entry in
+// `axis`. If `keep_dims` is true, the reduced dimensions are
+// retained with length 1.
+//
+// Arguments:
+// input: The tensor to reduce.
+// axis: The dimensions to reduce. Must be in the range
+// `[-rank(input), rank(input))`.
+//
+// Returns The reduced tensor.
+func Prod(scope *Scope, input tf.Output, axis tf.Output, optional ...ProdAttr) (output tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "Prod",
+ Input: []tf.Input{
+ input, axis,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
// ResizeBilinearAttr is an optional argument to ResizeBilinear.
type ResizeBilinearAttr func(optionalAttr)
@@ -11708,21 +12719,6 @@
return op.Output(0)
}
-// Computes softsign: `features / (abs(features) + 1)`.
-func Softsign(scope *Scope, features tf.Output) (activations tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "Softsign",
- Input: []tf.Input{
- features,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
// Creates a TensorList which, when stacked, has the value of `tensor`.
//
// Each tensor in the result list corresponds to one row of the input tensor.
@@ -11743,81 +12739,6 @@
return op.Output(0)
}
-// GenerateVocabRemappingAttr is an optional argument to GenerateVocabRemapping.
-type GenerateVocabRemappingAttr func(optionalAttr)
-
-// GenerateVocabRemappingOldVocabSize sets the optional old_vocab_size attribute to value.
-//
-// value: Number of entries in the old vocab file to consider. If -1,
-// use the entire old vocabulary.
-// If not specified, defaults to -1
-//
-// REQUIRES: value >= -1
-func GenerateVocabRemappingOldVocabSize(value int64) GenerateVocabRemappingAttr {
- return func(m optionalAttr) {
- m["old_vocab_size"] = value
- }
-}
-
-// Given a path to new and old vocabulary files, returns a remapping Tensor of
-//
-// length `num_new_vocab`, where `remapping[i]` contains the row number in the old
-// vocabulary that corresponds to row `i` in the new vocabulary (starting at line
-// `new_vocab_offset` and up to `num_new_vocab` entities), or `-1` if entry `i`
-// in the new vocabulary is not in the old vocabulary. The old vocabulary is
-// constrained to the first `old_vocab_size` entries if `old_vocab_size` is not the
-// default value of -1.
-//
-// `num_vocab_offset` enables
-// use in the partitioned variable case, and should generally be set through
-// examining partitioning info. The format of the files should be a text file,
-// with each line containing a single entity within the vocabulary.
-//
-// For example, with `new_vocab_file` a text file containing each of the following
-// elements on a single line: `[f0, f1, f2, f3]`, old_vocab_file = [f1, f0, f3],
-// `num_new_vocab = 3, new_vocab_offset = 1`, the returned remapping would be
-// `[0, -1, 2]`.
-//
-// The op also returns a count of how many entries in the new vocabulary
-// were present in the old vocabulary, which is used to calculate the number of
-// values to initialize in a weight matrix remapping
-//
-// This functionality can be used to remap both row vocabularies (typically,
-// features) and column vocabularies (typically, classes) from TensorFlow
-// checkpoints. Note that the partitioning logic relies on contiguous vocabularies
-// corresponding to div-partitioned variables. Moreover, the underlying remapping
-// uses an IndexTable (as opposed to an inexact CuckooTable), so client code should
-// use the corresponding index_table_from_file() as the FeatureColumn framework
-// does (as opposed to tf.feature_to_id(), which uses a CuckooTable).
-//
-// Arguments:
-// new_vocab_file: Path to the new vocab file.
-// old_vocab_file: Path to the old vocab file.
-// new_vocab_offset: How many entries into the new vocab file to start reading.
-// num_new_vocab: Number of entries in the new vocab file to remap.
-//
-// Returns A Tensor of length num_new_vocab where the element at index i
-// is equal to the old ID that maps to the new ID i. This element is -1 for any
-// new ID that is not found in the old vocabulary.Number of new vocab entries found in old vocab.
-func GenerateVocabRemapping(scope *Scope, new_vocab_file tf.Output, old_vocab_file tf.Output, new_vocab_offset int64, num_new_vocab int64, optional ...GenerateVocabRemappingAttr) (remapping tf.Output, num_present tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"new_vocab_offset": new_vocab_offset, "num_new_vocab": num_new_vocab}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "GenerateVocabRemapping",
- Input: []tf.Input{
- new_vocab_file, old_vocab_file,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0), op.Output(1)
-}
-
// Assigns sparse updates to the variable referenced by `resource`.
//
// This operation computes
@@ -12024,65 +12945,6 @@
return scope.AddOperation(opspec)
}
-// ComputeAccidentalHitsAttr is an optional argument to ComputeAccidentalHits.
-type ComputeAccidentalHitsAttr func(optionalAttr)
-
-// ComputeAccidentalHitsSeed 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 ComputeAccidentalHitsSeed(value int64) ComputeAccidentalHitsAttr {
- return func(m optionalAttr) {
- m["seed"] = value
- }
-}
-
-// ComputeAccidentalHitsSeed2 sets the optional seed2 attribute to value.
-//
-// value: An second seed to avoid seed collision.
-// If not specified, defaults to 0
-func ComputeAccidentalHitsSeed2(value int64) ComputeAccidentalHitsAttr {
- return func(m optionalAttr) {
- m["seed2"] = value
- }
-}
-
-// Computes the ids of the positions in sampled_candidates that match true_labels.
-//
-// When doing log-odds NCE, the result of this op should be passed through a
-// SparseToDense op, then added to the logits of the sampled candidates. This has
-// the effect of 'removing' the sampled labels that match the true labels by
-// making the classifier sure that they are sampled labels.
-//
-// Arguments:
-// true_classes: The true_classes output of UnpackSparseLabels.
-// sampled_candidates: The sampled_candidates output of CandidateSampler.
-// num_true: Number of true labels per context.
-//
-// Returns A vector of indices corresponding to rows of true_candidates.A vector of IDs of positions in sampled_candidates that match a true_label
-// for the row with the corresponding index in indices.A vector of the same length as indices and ids, in which each element
-// is -FLOAT_MAX.
-func ComputeAccidentalHits(scope *Scope, true_classes tf.Output, sampled_candidates tf.Output, num_true int64, optional ...ComputeAccidentalHitsAttr) (indices tf.Output, ids tf.Output, weights tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"num_true": num_true}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "ComputeAccidentalHits",
- Input: []tf.Input{
- true_classes, sampled_candidates,
- },
- Attrs: attrs,
- }
- 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)
@@ -12879,6 +13741,27 @@
return op.Output(0)
}
+// Retrieves the tree ensemble resource stamp token, number of trees and growing statistics.
+//
+// Arguments:
+// tree_ensemble_handle: Handle to the tree ensemble.
+//
+// Returns Stamp token of the tree ensemble resource.The number of trees in the tree ensemble resource.The number of trees that were finished successfully.The number of layers we attempted to build (but not necessarily succeeded).Rank size 2 tensor that contains start and end ids of the nodes in the latest
+// layer.
+func BoostedTreesGetEnsembleStates(scope *Scope, tree_ensemble_handle tf.Output) (stamp_token tf.Output, num_trees tf.Output, num_finalized_trees tf.Output, num_attempted_layers tf.Output, last_layer_nodes_range tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "BoostedTreesGetEnsembleStates",
+ Input: []tf.Input{
+ tree_ensemble_handle,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0), op.Output(1), op.Output(2), op.Output(3), op.Output(4)
+}
+
// ResourceApplyPowerSignAttr is an optional argument to ResourceApplyPowerSign.
type ResourceApplyPowerSignAttr func(optionalAttr)
@@ -13661,78 +14544,6 @@
return op.Output(0)
}
-// AllCandidateSamplerAttr is an optional argument to AllCandidateSampler.
-type AllCandidateSamplerAttr func(optionalAttr)
-
-// AllCandidateSamplerSeed 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 AllCandidateSamplerSeed(value int64) AllCandidateSamplerAttr {
- return func(m optionalAttr) {
- m["seed"] = value
- }
-}
-
-// AllCandidateSamplerSeed2 sets the optional seed2 attribute to value.
-//
-// value: An second seed to avoid seed collision.
-// If not specified, defaults to 0
-func AllCandidateSamplerSeed2(value int64) AllCandidateSamplerAttr {
- 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 produce.
-// 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.
-//
-// 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 AllCandidateSampler(scope *Scope, true_classes tf.Output, num_true int64, num_sampled int64, unique bool, optional ...AllCandidateSamplerAttr) (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}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "AllCandidateSampler",
- Input: []tf.Input{
- true_classes,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0), op.Output(1), op.Output(2)
-}
-
// Adds two `SparseTensor` objects to produce another `SparseTensor`.
//
// The input `SparseTensor` objects' indices are assumed ordered in standard
@@ -14543,6 +15354,78 @@
return op.Output(0)
}
+// Returns the last element of the input list as well as a list with all but that element.
+//
+// Fails if the list is empty.
+//
+// input_handle: the input list
+// tensor: the withdrawn last element of the list
+// element_dtype: the type of elements in the list
+// element_shape: the shape of the output tensor
+func TensorListPopBack(scope *Scope, input_handle tf.Output, element_dtype tf.DataType) (output_handle tf.Output, tensor tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"element_dtype": element_dtype}
+ opspec := tf.OpSpec{
+ Type: "TensorListPopBack",
+ Input: []tf.Input{
+ input_handle,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0), op.Output(1)
+}
+
+// MaxPoolGradGradAttr is an optional argument to MaxPoolGradGrad.
+type MaxPoolGradGradAttr func(optionalAttr)
+
+// MaxPoolGradGradDataFormat 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 MaxPoolGradGradDataFormat(value string) MaxPoolGradGradAttr {
+ return func(m optionalAttr) {
+ m["data_format"] = value
+ }
+}
+
+// Computes second-order gradients of the maxpooling function.
+//
+// Arguments:
+// orig_input: The original input tensor.
+// orig_output: The original output tensor.
+// grad: 4-D. Gradients of gradients w.r.t. the input of `max_pool`.
+// 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 Gradients of gradients w.r.t. the input to `max_pool`.
+func MaxPoolGradGrad(scope *Scope, orig_input tf.Output, orig_output tf.Output, grad tf.Output, ksize []int64, strides []int64, padding string, optional ...MaxPoolGradGradAttr) (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: "MaxPoolGradGrad",
+ Input: []tf.Input{
+ orig_input, orig_output, grad,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
// TensorArrayGatherV3Attr is an optional argument to TensorArrayGatherV3.
type TensorArrayGatherV3Attr func(optionalAttr)
@@ -14589,33 +15472,6 @@
return op.Output(0)
}
-// This op consumes a lock created by `MutexLock`.
-//
-// This op exists to consume a tensor created by `MutexLock` (other than
-// direct control dependencies). It should be the only that consumes the tensor,
-// and will raise an error if it is not. Its only purpose is to keep the
-// mutex lock tensor alive until it is consumed by this op.
-//
-// **NOTE**: This operation must run on the same device as its input. This may
-// be enforced via the `colocate_with` mechanism.
-//
-// Arguments:
-// mutex_lock: A tensor returned by `MutexLock`.
-//
-// Returns the created operation.
-func ConsumeMutexLock(scope *Scope, mutex_lock tf.Output) (o *tf.Operation) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "ConsumeMutexLock",
- Input: []tf.Input{
- mutex_lock,
- },
- }
- return scope.AddOperation(opspec)
-}
-
// Returns x / y element-wise for integer types.
//
// Truncation designates that negative numbers will round fractional quantities
@@ -15686,79 +16542,6 @@
return op.Output(0)
}
-// LogUniformCandidateSamplerAttr is an optional argument to LogUniformCandidateSampler.
-type LogUniformCandidateSamplerAttr func(optionalAttr)
-
-// LogUniformCandidateSamplerSeed 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 LogUniformCandidateSamplerSeed(value int64) LogUniformCandidateSamplerAttr {
- return func(m optionalAttr) {
- m["seed"] = value
- }
-}
-
-// LogUniformCandidateSamplerSeed2 sets the optional seed2 attribute to value.
-//
-// value: An second seed to avoid seed collision.
-// If not specified, defaults to 0
-func LogUniformCandidateSamplerSeed2(value int64) LogUniformCandidateSamplerAttr {
- return func(m optionalAttr) {
- m["seed2"] = value
- }
-}
-
-// Generates labels for candidate sampling with a log-uniform 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 LogUniformCandidateSampler(scope *Scope, true_classes tf.Output, num_true int64, num_sampled int64, unique bool, range_max int64, optional ...LogUniformCandidateSamplerAttr) (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: "LogUniformCandidateSampler",
- Input: []tf.Input{
- true_classes,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0), op.Output(1), op.Output(2)
-}
-
// Returns the max of x and y (i.e. x > y ? x : y) element-wise.
//
// *NOTE*: `Maximum` supports broadcasting. More about broadcasting
@@ -16044,109 +16827,6 @@
return op.Output(0)
}
-// ProdAttr is an optional argument to Prod.
-type ProdAttr func(optionalAttr)
-
-// ProdKeepDims sets the optional keep_dims attribute to value.
-//
-// value: If true, retain reduced dimensions with length 1.
-// If not specified, defaults to false
-func ProdKeepDims(value bool) ProdAttr {
- return func(m optionalAttr) {
- m["keep_dims"] = value
- }
-}
-
-// Computes the product of elements across dimensions of a tensor.
-//
-// Reduces `input` along the dimensions given in `axis`. Unless
-// `keep_dims` is true, the rank of the tensor is reduced by 1 for each entry in
-// `axis`. If `keep_dims` is true, the reduced dimensions are
-// retained with length 1.
-//
-// Arguments:
-// input: The tensor to reduce.
-// axis: The dimensions to reduce. Must be in the range
-// `[-rank(input), rank(input))`.
-//
-// Returns The reduced tensor.
-func Prod(scope *Scope, input tf.Output, axis tf.Output, optional ...ProdAttr) (output tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "Prod",
- Input: []tf.Input{
- input, axis,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// FusedResizeAndPadConv2DAttr is an optional argument to FusedResizeAndPadConv2D.
-type FusedResizeAndPadConv2DAttr func(optionalAttr)
-
-// FusedResizeAndPadConv2DResizeAlignCorners sets the optional resize_align_corners attribute to value.
-//
-// value: If true, the centers of the 4 corner pixels of the input and output tensors are
-// aligned, preserving the values at the corner pixels. Defaults to false.
-// If not specified, defaults to false
-func FusedResizeAndPadConv2DResizeAlignCorners(value bool) FusedResizeAndPadConv2DAttr {
- return func(m optionalAttr) {
- m["resize_align_corners"] = value
- }
-}
-
-// Performs a resize and padding as a preprocess during a convolution.
-//
-// It's often possible to do spatial transformations more efficiently as part of
-// the packing stage of a convolution, so this op allows for an optimized
-// implementation where these stages are fused together. This prevents the need to
-// write out the intermediate results as whole tensors, reducing memory pressure,
-// and we can get some latency gains by merging the transformation calculations.
-// The data_format attribute for Conv2D isn't supported by this op, and defaults to
-// 'NHWC' order.
-// Internally this op uses a single per-graph scratch buffer, which means that it
-// will block if multiple versions are being run in parallel. This is because this
-// operator is primarily an optimization to minimize memory usage.
-//
-// Arguments:
-// input: 4-D with shape `[batch, in_height, in_width, in_channels]`.
-// size: A 1-D int32 Tensor of 2 elements: `new_height, new_width`. The
-// new size for the images.
-// paddings: A two-column matrix specifying the padding sizes. The number of
-// rows must be the same as the rank of `input`.
-// filter: 4-D with shape
-// `[filter_height, filter_width, in_channels, out_channels]`.
-//
-// strides: 1-D of length 4. The stride of the sliding window for each dimension
-// of `input`. Must be in the same order as the dimension specified with format.
-// padding: The type of padding algorithm to use.
-func FusedResizeAndPadConv2D(scope *Scope, input tf.Output, size tf.Output, paddings tf.Output, filter tf.Output, mode string, strides []int64, padding string, optional ...FusedResizeAndPadConv2DAttr) (output tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"mode": mode, "strides": strides, "padding": padding}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "FusedResizeAndPadConv2D",
- Input: []tf.Input{
- input, size, paddings, filter,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
// Returns a list of tensors with the same shapes and contents as the input
//
// tensors.
@@ -17616,175 +18296,6 @@
return op.Output(0)
}
-// BoostedTreesEnsembleResourceHandleOpAttr is an optional argument to BoostedTreesEnsembleResourceHandleOp.
-type BoostedTreesEnsembleResourceHandleOpAttr func(optionalAttr)
-
-// BoostedTreesEnsembleResourceHandleOpContainer sets the optional container attribute to value.
-// If not specified, defaults to ""
-func BoostedTreesEnsembleResourceHandleOpContainer(value string) BoostedTreesEnsembleResourceHandleOpAttr {
- return func(m optionalAttr) {
- m["container"] = value
- }
-}
-
-// BoostedTreesEnsembleResourceHandleOpSharedName sets the optional shared_name attribute to value.
-// If not specified, defaults to ""
-func BoostedTreesEnsembleResourceHandleOpSharedName(value string) BoostedTreesEnsembleResourceHandleOpAttr {
- return func(m optionalAttr) {
- m["shared_name"] = value
- }
-}
-
-// Creates a handle to a BoostedTreesEnsembleResource
-func BoostedTreesEnsembleResourceHandleOp(scope *Scope, optional ...BoostedTreesEnsembleResourceHandleOpAttr) (resource tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "BoostedTreesEnsembleResourceHandleOp",
-
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// ResourceApplyMomentumAttr is an optional argument to ResourceApplyMomentum.
-type ResourceApplyMomentumAttr func(optionalAttr)
-
-// ResourceApplyMomentumUseLocking 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 ResourceApplyMomentumUseLocking(value bool) ResourceApplyMomentumAttr {
- return func(m optionalAttr) {
- m["use_locking"] = value
- }
-}
-
-// ResourceApplyMomentumUseNesterov sets the optional use_nesterov attribute to value.
-//
-// value: If `True`, the tensor passed to compute grad will be
-// var - lr * momentum * accum, so in the end, the var you get is actually
-// var - lr * momentum * accum.
-// If not specified, defaults to false
-func ResourceApplyMomentumUseNesterov(value bool) ResourceApplyMomentumAttr {
- return func(m optionalAttr) {
- m["use_nesterov"] = value
- }
-}
-
-// Update '*var' according to the momentum scheme. Set use_nesterov = True if you
-//
-// want to use Nesterov momentum.
-//
-// accum = accum * momentum + grad
-// var -= lr * accum
-//
-// Arguments:
-// var_: Should be from a Variable().
-// accum: Should be from a Variable().
-// lr: Scaling factor. Must be a scalar.
-// grad: The gradient.
-// momentum: Momentum. Must be a scalar.
-//
-// Returns the created operation.
-func ResourceApplyMomentum(scope *Scope, var_ tf.Output, accum tf.Output, lr tf.Output, grad tf.Output, momentum tf.Output, optional ...ResourceApplyMomentumAttr) (o *tf.Operation) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "ResourceApplyMomentum",
- Input: []tf.Input{
- var_, accum, lr, grad, momentum,
- },
- Attrs: attrs,
- }
- return scope.AddOperation(opspec)
-}
-
-// MaxPoolGradGradAttr is an optional argument to MaxPoolGradGrad.
-type MaxPoolGradGradAttr func(optionalAttr)
-
-// MaxPoolGradGradDataFormat 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 MaxPoolGradGradDataFormat(value string) MaxPoolGradGradAttr {
- return func(m optionalAttr) {
- m["data_format"] = value
- }
-}
-
-// Computes second-order gradients of the maxpooling function.
-//
-// Arguments:
-// orig_input: The original input tensor.
-// orig_output: The original output tensor.
-// grad: 4-D. Gradients of gradients w.r.t. the input of `max_pool`.
-// 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 Gradients of gradients w.r.t. the input to `max_pool`.
-func MaxPoolGradGrad(scope *Scope, orig_input tf.Output, orig_output tf.Output, grad tf.Output, ksize []int64, strides []int64, padding string, optional ...MaxPoolGradGradAttr) (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: "MaxPoolGradGrad",
- Input: []tf.Input{
- orig_input, orig_output, grad,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// Returns the last element of the input list as well as a list with all but that element.
-//
-// Fails if the list is empty.
-//
-// input_handle: the input list
-// tensor: the withdrawn last element of the list
-// element_dtype: the type of elements in the list
-// element_shape: the shape of the output tensor
-func TensorListPopBack(scope *Scope, input_handle tf.Output, element_dtype tf.DataType) (output_handle tf.Output, tensor tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"element_dtype": element_dtype}
- opspec := tf.OpSpec{
- Type: "TensorListPopBack",
- Input: []tf.Input{
- input_handle,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0), op.Output(1)
-}
-
// Returns element-wise integer closest to x.
//
// If the result is midway between two representable values,
@@ -19075,31 +19586,6 @@
return op.Output(0)
}
-// Read an element from the TensorArray into output `value`.
-//
-// Arguments:
-// handle: The handle to a TensorArray.
-//
-// flow_in: A float scalar that enforces proper chaining of operations.
-// dtype: The type of the elem that is returned.
-//
-// Returns The tensor that is read from the TensorArray.
-func TensorArrayReadV3(scope *Scope, handle tf.Output, index tf.Output, flow_in tf.Output, dtype tf.DataType) (value tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"dtype": dtype}
- opspec := tf.OpSpec{
- Type: "TensorArrayReadV3",
- Input: []tf.Input{
- handle, index, flow_in,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
// QuantizeV2Attr is an optional argument to QuantizeV2.
type QuantizeV2Attr func(optionalAttr)
@@ -20497,6 +20983,201 @@
return op.Output(0)
}
+// EnterAttr is an optional argument to Enter.
+type EnterAttr func(optionalAttr)
+
+// EnterIsConstant sets the optional is_constant attribute to value.
+//
+// value: If true, the output is constant within the child frame.
+// If not specified, defaults to false
+func EnterIsConstant(value bool) EnterAttr {
+ return func(m optionalAttr) {
+ m["is_constant"] = value
+ }
+}
+
+// EnterParallelIterations sets the optional parallel_iterations attribute to value.
+//
+// value: The number of iterations allowed to run in parallel.
+// If not specified, defaults to 10
+func EnterParallelIterations(value int64) EnterAttr {
+ return func(m optionalAttr) {
+ m["parallel_iterations"] = value
+ }
+}
+
+// Creates or finds a child frame, and makes `data` available to the child frame.
+//
+// This op is used together with `Exit` to create loops in the graph.
+// The unique `frame_name` is used by the `Executor` to identify frames. If
+// `is_constant` is true, `output` is a constant in the child frame; otherwise
+// it may be changed in the child frame. At most `parallel_iterations` iterations
+// are run in parallel in the child frame.
+//
+// Arguments:
+// data: The tensor to be made available to the child frame.
+// frame_name: The name of the child frame.
+//
+// Returns The same tensor as `data`.
+func Enter(scope *Scope, data tf.Output, frame_name string, optional ...EnterAttr) (output tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"frame_name": frame_name}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "Enter",
+ Input: []tf.Input{
+ data,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// Add all input tensors element wise.
+//
+// Arguments:
+// inputs: Must all be the same size and shape.
+func AddN(scope *Scope, inputs []tf.Output) (sum tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "AddN",
+ Input: []tf.Input{
+ tf.OutputList(inputs),
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// TryRpcAttr is an optional argument to TryRpc.
+type TryRpcAttr func(optionalAttr)
+
+// TryRpcProtocol sets the optional protocol attribute to value.
+//
+// value: RPC protocol to use. Empty string means use the default protocol.
+// Options include 'grpc'.
+// If not specified, defaults to ""
+func TryRpcProtocol(value string) TryRpcAttr {
+ return func(m optionalAttr) {
+ m["protocol"] = value
+ }
+}
+
+// TryRpcFailFast sets the optional fail_fast attribute to value.
+//
+// value: `boolean`. If `true` (default), then failures to connect
+// (i.e., the server does not immediately respond) cause an RPC failure.
+// If not specified, defaults to true
+func TryRpcFailFast(value bool) TryRpcAttr {
+ return func(m optionalAttr) {
+ m["fail_fast"] = value
+ }
+}
+
+// TryRpcTimeoutInMs sets the optional timeout_in_ms attribute to value.
+//
+// value: `int`. If `0` (default), then the kernel will run the RPC
+// request and only time out if the RPC deadline passes or the session times out.
+// If this value is greater than `0`, then the op will raise an exception if
+// the RPC takes longer than `timeout_in_ms`.
+// If not specified, defaults to 0
+func TryRpcTimeoutInMs(value int64) TryRpcAttr {
+ return func(m optionalAttr) {
+ m["timeout_in_ms"] = value
+ }
+}
+
+// Perform batches of RPC requests.
+//
+// This op asynchronously performs either a single RPC request, or a batch
+// of requests. RPC requests are defined by three main parameters:
+//
+// - `address` (the host+port or BNS address of the request)
+// - `method` (the method name for the request)
+// - `request` (the serialized proto string, or vector of strings,
+// of the RPC request argument).
+//
+// For example, if you have an RPC service running on port localhost:2345,
+// and its interface is configured with the following proto declaration:
+//
+// ```
+// service MyService {
+// rpc MyMethod(MyRequestProto) returns (MyResponseProto) {
+// }
+// };
+// ```
+//
+// then call this op with arguments:
+//
+// ```
+// address = "localhost:2345"
+// method = "MyService/MyMethod"
+// ```
+//
+// The `request` tensor is a string tensor representing serialized `MyRequestProto`
+// strings; and the output string tensor `response` will have the same shape
+// and contain (upon successful completion) corresponding serialized
+// `MyResponseProto` strings.
+//
+// For example, to send a single, empty, `MyRequestProto`, call
+// this op with `request = ""`. To send 5 **parallel** empty requests,
+// call this op with `request = ["", "", "", "", ""]`.
+//
+// More generally, one can create a batch of `MyRequestProto` serialized protos
+// from regular batched tensors using the `encode_proto` op, and convert
+// the response `MyResponseProto` serialized protos to batched tensors
+// using the `decode_proto` op.
+//
+// **NOTE** Working with serialized proto strings is faster than instantiating
+// actual proto objects in memory, so no performance degradation is expected
+// compared to writing custom kernels for this workflow.
+//
+// Unlike the standard `Rpc` op, if the connection fails or the remote worker
+// returns an error status, this op does **not** reraise the exception.
+// Instead, the `status_code` and `status_message` entry for the corresponding RPC
+// call is set with the error returned from the RPC call. The `response` tensor
+// will contain valid response values for those minibatch entries whose RPCs did
+// not fail; the rest of the entries will have empty strings.
+//
+// Arguments:
+// address: `0-D` or `1-D`. The address (i.e. host_name:port) of the RPC server.
+// If this tensor has more than 1 element, then multiple parallel rpc requests
+// are sent. This argument broadcasts with `method` and `request`.
+// method: `0-D` or `1-D`. The method address on the RPC server.
+// If this tensor has more than 1 element, then multiple parallel rpc requests
+// are sent. This argument broadcasts with `address` and `request`.
+// request: `0-D` or `1-D`. Serialized proto strings: the rpc request argument.
+// If this tensor has more than 1 element, then multiple parallel rpc requests
+// are sent. This argument broadcasts with `address` and `method`.
+//
+// Returns Same shape as `request`. Serialized proto strings: the rpc responses.Same shape as `request`. Values correspond to tensorflow Status enum codes.Same shape as `request`. Values correspond to Status messages
+// returned from the RPC calls.
+func TryRpc(scope *Scope, address tf.Output, method tf.Output, request tf.Output, optional ...TryRpcAttr) (response tf.Output, status_code tf.Output, status_message tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "TryRpc",
+ Input: []tf.Input{
+ address, method, request,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0), op.Output(1), op.Output(2)
+}
+
// Delete the tensor specified by its handle in the session.
//
// Arguments:
@@ -20974,76 +21655,6 @@
return op.Output(0)
}
-// DepthwiseConv2dNativeBackpropInputAttr is an optional argument to DepthwiseConv2dNativeBackpropInput.
-type DepthwiseConv2dNativeBackpropInputAttr func(optionalAttr)
-
-// DepthwiseConv2dNativeBackpropInputDataFormat 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, height, width, channels].
-// Alternatively, the format could be "NCHW", the data storage order of:
-// [batch, channels, height, width].
-// If not specified, defaults to "NHWC"
-func DepthwiseConv2dNativeBackpropInputDataFormat(value string) DepthwiseConv2dNativeBackpropInputAttr {
- return func(m optionalAttr) {
- m["data_format"] = value
- }
-}
-
-// DepthwiseConv2dNativeBackpropInputDilations sets the optional dilations attribute to value.
-//
-// value: 1-D tensor of length 4. The dilation factor for each dimension of
-// `input`. If set to k > 1, there will be k-1 skipped cells between each filter
-// element on that dimension. The dimension order is determined by the value of
-// `data_format`, see above for details. Dilations in the batch and depth
-// dimensions must be 1.
-// If not specified, defaults to <i:1 i:1 i:1 i:1 >
-func DepthwiseConv2dNativeBackpropInputDilations(value []int64) DepthwiseConv2dNativeBackpropInputAttr {
- return func(m optionalAttr) {
- m["dilations"] = value
- }
-}
-
-// Computes the gradients of depthwise convolution with respect to the input.
-//
-// Arguments:
-// input_sizes: An integer vector representing the shape of `input`, based
-// on `data_format`. For example, if `data_format` is 'NHWC' then
-// `input` is a 4-D `[batch, height, width, channels]` tensor.
-// filter: 4-D with shape
-// `[filter_height, filter_width, in_channels, depthwise_multiplier]`.
-// out_backprop: 4-D with shape based on `data_format`.
-// For example, if `data_format` is 'NHWC' then
-// out_backprop shape is `[batch, out_height, out_width, out_channels]`.
-// Gradients w.r.t. the output of the convolution.
-// strides: The stride of the sliding window for each dimension of the input
-// of the convolution.
-// padding: The type of padding algorithm to use.
-//
-// Returns 4-D with shape according to `data_format`. For example, if
-// `data_format` is 'NHWC', output shape is `[batch, in_height,
-// in_width, in_channels]`. Gradient w.r.t. the input of the
-// convolution.
-func DepthwiseConv2dNativeBackpropInput(scope *Scope, input_sizes tf.Output, filter tf.Output, out_backprop tf.Output, strides []int64, padding string, optional ...DepthwiseConv2dNativeBackpropInputAttr) (output tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"strides": strides, "padding": padding}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "DepthwiseConv2dNativeBackpropInput",
- Input: []tf.Input{
- input_sizes, filter, out_backprop,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
// Stops gradient computation.
//
// When executed in a graph, this op outputs its input tensor as-is.
@@ -21313,23 +21924,44 @@
return op.Output(0), op.Output(1), op.Output(2)
}
-// Forwards the input to the output.
+// Computes the sum along segments of a tensor.
//
-// This operator represents the loop termination condition used by the
-// "pivot" switches of a loop.
+// Read
+// [the section on segmentation](https://tensorflow.org/api_guides/python/math_ops#Segmentation)
+// for an explanation of segments.
+//
+// Computes a tensor such that
+// \\(output[i] = \sum_{j...} data[j...]\\) where the sum is over tuples `j...` such
+// that `segment_ids[j...] == i`. Unlike `SegmentSum`, `segment_ids`
+// need not be sorted and need not cover all values in the full
+// range of valid values.
+//
+// If the sum is empty for a given segment ID `i`, `output[i] = 0`.
+// If the given segment ID `i` is negative, the value is dropped and will not be
+// added to the sum of the segment.
+//
+// `num_segments` should equal the number of distinct segment IDs.
+//
+// <div style="width:70%; margin:auto; margin-bottom:10px; margin-top:20px;">
+// <img style="width:100%" src="https://www.tensorflow.org/images/UnsortedSegmentSum.png" alt>
+// </div>
//
// Arguments:
-// input: A boolean scalar, representing the branch predicate of the Switch op.
//
-// Returns The same tensor as `input`.
-func LoopCond(scope *Scope, input tf.Output) (output tf.Output) {
+// segment_ids: A tensor whose shape is a prefix of `data.shape`.
+//
+//
+// Returns Has same shape as data, except for the first `segment_ids.rank`
+// dimensions, which are replaced with a single dimension which has size
+// `num_segments`.
+func UnsortedSegmentSum(scope *Scope, data tf.Output, segment_ids tf.Output, num_segments tf.Output) (output tf.Output) {
if scope.Err() != nil {
return
}
opspec := tf.OpSpec{
- Type: "LoopCond",
+ Type: "UnsortedSegmentSum",
Input: []tf.Input{
- input,
+ data, segment_ids, num_segments,
},
}
op := scope.AddOperation(opspec)
@@ -21963,40 +22595,6 @@
return op.Output(0)
}
-// Creates a sequence of numbers.
-//
-// This operation creates a sequence of numbers that begins at `start` and
-// extends by increments of `delta` up to but not including `limit`.
-//
-// For example:
-//
-// ```
-// # 'start' is 3
-// # 'limit' is 18
-// # 'delta' is 3
-// tf.range(start, limit, delta) ==> [3, 6, 9, 12, 15]
-// ```
-//
-// Arguments:
-// start: 0-D (scalar). First entry in the sequence.
-// limit: 0-D (scalar). Upper limit of sequence, exclusive.
-// delta: 0-D (scalar). Optional. Default is 1. Number that increments `start`.
-//
-// Returns 1-D.
-func Range(scope *Scope, start tf.Output, limit tf.Output, delta tf.Output) (output tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "Range",
- Input: []tf.Input{
- start, limit, delta,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
// DestroyResourceOpAttr is an optional argument to DestroyResourceOp.
type DestroyResourceOpAttr func(optionalAttr)
@@ -23061,156 +23659,6 @@
return op.Output(0)
}
-// Bucketizes 'input' based on 'boundaries'.
-//
-// For example, if the inputs are
-// boundaries = [0, 10, 100]
-// input = [[-5, 10000]
-// [150, 10]
-// [5, 100]]
-//
-// then the output will be
-// output = [[0, 3]
-// [3, 2]
-// [1, 3]]
-//
-// Arguments:
-// input: Any shape of Tensor contains with int or float type.
-// boundaries: A sorted list of floats gives the boundary of the buckets.
-//
-// Returns Same shape with 'input', each value of input replaced with bucket index.
-//
-// @compatibility(numpy)
-// Equivalent to np.digitize.
-// @end_compatibility
-func Bucketize(scope *Scope, input tf.Output, boundaries []float32) (output tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"boundaries": boundaries}
- opspec := tf.OpSpec{
- Type: "Bucketize",
- Input: []tf.Input{
- input,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// Calculates gains for each feature and returns the best possible split information for the feature.
-//
-// The split information is the best threshold (bucket id), gains and left/right node contributions per node for each feature.
-//
-// It is possible that not all nodes can be split on each feature. Hence, the list of possible nodes can differ between the features. Therefore, we return `node_ids_list` for each feature, containing the list of nodes that this feature can be used to split.
-//
-// In this manner, the output is the best split per features and per node, so that it needs to be combined later to produce the best split for each node (among all possible features).
-//
-// The length of output lists are all of the same length, `num_features`.
-// The output shapes are compatible in a way that the first dimension of all tensors of all lists are the same and equal to the number of possible split nodes for each feature.
-//
-// Arguments:
-// node_id_range: A Rank 1 tensor (shape=[2]) to specify the range [first, last) of node ids to process within `stats_summary_list`. The nodes are iterated between the two nodes specified by the tensor, as like `for node_id in range(node_id_range[0], node_id_range[1])` (Note that the last index node_id_range[1] is exclusive).
-// stats_summary_list: A list of Rank 3 tensor (#shape=[max_splits, bucket, 2]) for accumulated stats summary (gradient/hessian) per node per buckets for each feature. The first dimension of the tensor is the maximum number of splits, and thus not all elements of it will be used, but only the indexes specified by node_ids will be used.
-// l1: l1 regularization factor on leaf weights, per instance based.
-// l2: l2 regularization factor on leaf weights, per instance based.
-// tree_complexity: adjustment to the gain, per leaf based.
-// min_node_weight: mininum avg of hessians in a node before required for the node to be considered for splitting.
-// max_splits: the number of nodes that can be split in the whole tree. Used as a dimension of output tensors.
-//
-// Returns An output list of Rank 1 tensors indicating possible split node ids for each feature. The length of the list is num_features, but each tensor has different size as each feature provides different possible nodes. See above for details like shapes and sizes.An output list of Rank 1 tensors indicating the best gains for each feature to split for certain nodes. See above for details like shapes and sizes.An output list of Rank 1 tensors indicating the bucket id to compare with (as a threshold) for split in each node. See above for details like shapes and sizes.A list of Rank 2 tensors indicating the contribution of the left nodes when branching from parent nodes (given by the tensor element in the output node_ids_list) to the left direction by the given threshold for each feature. This value will be used to make the left node value by adding to the parent node value. Second dimension size is 1 for 1-dimensional logits, but would be larger for multi-class problems. See above for details like shapes and sizes.A list of Rank 2 tensors, with the same shape/conditions as left_node_contribs_list, but just that the value is for the right node.
-func BoostedTreesCalculateBestGainsPerFeature(scope *Scope, node_id_range tf.Output, stats_summary_list []tf.Output, l1 tf.Output, l2 tf.Output, tree_complexity tf.Output, min_node_weight tf.Output, max_splits int64) (node_ids_list []tf.Output, gains_list []tf.Output, thresholds_list []tf.Output, left_node_contribs_list []tf.Output, right_node_contribs_list []tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"max_splits": max_splits}
- opspec := tf.OpSpec{
- Type: "BoostedTreesCalculateBestGainsPerFeature",
- Input: []tf.Input{
- node_id_range, tf.OutputList(stats_summary_list), l1, l2, tree_complexity, min_node_weight,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- if scope.Err() != nil {
- return
- }
- var idx int
- var err error
- if node_ids_list, idx, err = makeOutputList(op, idx, "node_ids_list"); err != nil {
- scope.UpdateErr("BoostedTreesCalculateBestGainsPerFeature", err)
- return
- }
- if gains_list, idx, err = makeOutputList(op, idx, "gains_list"); err != nil {
- scope.UpdateErr("BoostedTreesCalculateBestGainsPerFeature", err)
- return
- }
- if thresholds_list, idx, err = makeOutputList(op, idx, "thresholds_list"); err != nil {
- scope.UpdateErr("BoostedTreesCalculateBestGainsPerFeature", err)
- return
- }
- if left_node_contribs_list, idx, err = makeOutputList(op, idx, "left_node_contribs_list"); err != nil {
- scope.UpdateErr("BoostedTreesCalculateBestGainsPerFeature", err)
- return
- }
- if right_node_contribs_list, idx, err = makeOutputList(op, idx, "right_node_contribs_list"); err != nil {
- scope.UpdateErr("BoostedTreesCalculateBestGainsPerFeature", err)
- return
- }
- return node_ids_list, gains_list, thresholds_list, left_node_contribs_list, right_node_contribs_list
-}
-
-// EncodePngAttr is an optional argument to EncodePng.
-type EncodePngAttr func(optionalAttr)
-
-// EncodePngCompression sets the optional compression attribute to value.
-//
-// value: Compression level.
-// If not specified, defaults to -1
-func EncodePngCompression(value int64) EncodePngAttr {
- return func(m optionalAttr) {
- m["compression"] = value
- }
-}
-
-// PNG-encode an image.
-//
-// `image` is a 3-D uint8 or uint16 Tensor of shape `[height, width, channels]`
-// where `channels` is:
-//
-// * 1: for grayscale.
-// * 2: for grayscale + alpha.
-// * 3: for RGB.
-// * 4: for RGBA.
-//
-// The ZLIB compression level, `compression`, can be -1 for the PNG-encoder
-// default or a value from 0 to 9. 9 is the highest compression level, generating
-// the smallest output, but is slower.
-//
-// Arguments:
-// image: 3-D with shape `[height, width, channels]`.
-//
-// Returns 0-D. PNG-encoded image.
-func EncodePng(scope *Scope, image tf.Output, optional ...EncodePngAttr) (contents tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "EncodePng",
- Input: []tf.Input{
- image,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
// Updates the table to associates keys with values.
//
// The tensor `keys` must be of the same type as the keys of the table.
@@ -24004,6 +24452,31 @@
return op.Output(0)
}
+// Read an element from the TensorArray into output `value`.
+//
+// Arguments:
+// handle: The handle to a TensorArray.
+//
+// flow_in: A float scalar that enforces proper chaining of operations.
+// dtype: The type of the elem that is returned.
+//
+// Returns The tensor that is read from the TensorArray.
+func TensorArrayReadV3(scope *Scope, handle tf.Output, index tf.Output, flow_in tf.Output, dtype tf.DataType) (value tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"dtype": dtype}
+ opspec := tf.OpSpec{
+ Type: "TensorArrayReadV3",
+ Input: []tf.Input{
+ handle, index, flow_in,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
// Computes the gradient for the tanh of `x` wrt its input.
//
// Specifically, `grad = dy * (1 - y*y)`, where `y = tanh(x)`, and `dy`
@@ -26678,6 +27151,260 @@
return op.Output(0)
}
+// LearnedUnigramCandidateSamplerAttr is an optional argument to LearnedUnigramCandidateSampler.
+type LearnedUnigramCandidateSamplerAttr func(optionalAttr)
+
+// LearnedUnigramCandidateSamplerSeed 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 LearnedUnigramCandidateSamplerSeed(value int64) LearnedUnigramCandidateSamplerAttr {
+ return func(m optionalAttr) {
+ m["seed"] = value
+ }
+}
+
+// LearnedUnigramCandidateSamplerSeed2 sets the optional seed2 attribute to value.
+//
+// value: An second seed to avoid seed collision.
+// If not specified, defaults to 0
+func LearnedUnigramCandidateSamplerSeed2(value int64) LearnedUnigramCandidateSamplerAttr {
+ 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 LearnedUnigramCandidateSampler(scope *Scope, true_classes tf.Output, num_true int64, num_sampled int64, unique bool, range_max int64, optional ...LearnedUnigramCandidateSamplerAttr) (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: "LearnedUnigramCandidateSampler",
+ Input: []tf.Input{
+ true_classes,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0), op.Output(1), op.Output(2)
+}
+
+// SerializeSparseAttr is an optional argument to SerializeSparse.
+type SerializeSparseAttr func(optionalAttr)
+
+// SerializeSparseOutType 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 SerializeSparseOutType(value tf.DataType) SerializeSparseAttr {
+ return func(m optionalAttr) {
+ m["out_type"] = value
+ }
+}
+
+// Serialize a `SparseTensor` into a `[3]` `Tensor` object.
+//
+// Arguments:
+// sparse_indices: 2-D. The `indices` of the `SparseTensor`.
+// sparse_values: 1-D. The `values` of the `SparseTensor`.
+// sparse_shape: 1-D. The `shape` of the `SparseTensor`.
+func SerializeSparse(scope *Scope, sparse_indices tf.Output, sparse_values tf.Output, sparse_shape tf.Output, optional ...SerializeSparseAttr) (serialized_sparse tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "SerializeSparse",
+ Input: []tf.Input{
+ sparse_indices, sparse_values, sparse_shape,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// RandomShuffleQueueV2Attr is an optional argument to RandomShuffleQueueV2.
+type RandomShuffleQueueV2Attr func(optionalAttr)
+
+// RandomShuffleQueueV2Shapes sets the optional shapes attribute to value.
+//
+// value: The shape of each component in a value. The length of this attr must
+// be either 0 or the same as the length of component_types. If the length of
+// this attr is 0, the shapes of queue elements are not constrained, and
+// only one element may be dequeued at a time.
+// If not specified, defaults to <>
+//
+// REQUIRES: len(value) >= 0
+func RandomShuffleQueueV2Shapes(value []tf.Shape) RandomShuffleQueueV2Attr {
+ return func(m optionalAttr) {
+ m["shapes"] = value
+ }
+}
+
+// RandomShuffleQueueV2Capacity sets the optional capacity attribute to value.
+//
+// value: The upper bound on the number of elements in this queue.
+// Negative numbers mean no limit.
+// If not specified, defaults to -1
+func RandomShuffleQueueV2Capacity(value int64) RandomShuffleQueueV2Attr {
+ return func(m optionalAttr) {
+ m["capacity"] = value
+ }
+}
+
+// RandomShuffleQueueV2MinAfterDequeue sets the optional min_after_dequeue attribute to value.
+//
+// value: Dequeue will block unless there would be this
+// many elements after the dequeue or the queue is closed. This
+// ensures a minimum level of mixing of elements.
+// If not specified, defaults to 0
+func RandomShuffleQueueV2MinAfterDequeue(value int64) RandomShuffleQueueV2Attr {
+ return func(m optionalAttr) {
+ m["min_after_dequeue"] = value
+ }
+}
+
+// RandomShuffleQueueV2Seed sets the optional seed attribute to value.
+//
+// value: If either seed or seed2 is set to be non-zero, the random number
+// generator is seeded by the given seed. Otherwise, a random seed is used.
+// If not specified, defaults to 0
+func RandomShuffleQueueV2Seed(value int64) RandomShuffleQueueV2Attr {
+ return func(m optionalAttr) {
+ m["seed"] = value
+ }
+}
+
+// RandomShuffleQueueV2Seed2 sets the optional seed2 attribute to value.
+//
+// value: A second seed to avoid seed collision.
+// If not specified, defaults to 0
+func RandomShuffleQueueV2Seed2(value int64) RandomShuffleQueueV2Attr {
+ return func(m optionalAttr) {
+ m["seed2"] = value
+ }
+}
+
+// RandomShuffleQueueV2Container 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 RandomShuffleQueueV2Container(value string) RandomShuffleQueueV2Attr {
+ return func(m optionalAttr) {
+ m["container"] = value
+ }
+}
+
+// RandomShuffleQueueV2SharedName sets the optional shared_name attribute to value.
+//
+// value: If non-empty, this queue will be shared under the given name
+// across multiple sessions.
+// If not specified, defaults to ""
+func RandomShuffleQueueV2SharedName(value string) RandomShuffleQueueV2Attr {
+ return func(m optionalAttr) {
+ m["shared_name"] = value
+ }
+}
+
+// A queue that randomizes the order of elements.
+//
+// Arguments:
+// component_types: The type of each component in a value.
+//
+// Returns The handle to the queue.
+func RandomShuffleQueueV2(scope *Scope, component_types []tf.DataType, optional ...RandomShuffleQueueV2Attr) (handle tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"component_types": component_types}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "RandomShuffleQueueV2",
+
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// Draw bounding boxes on a batch of images.
+//
+// Outputs a copy of `images` but draws on top of the pixels zero or more bounding
+// boxes specified by the locations in `boxes`. The coordinates of the each
+// bounding box in `boxes` are encoded as `[y_min, x_min, y_max, x_max]`. The
+// bounding box coordinates are floats in `[0.0, 1.0]` relative to the width and
+// height of the underlying image.
+//
+// For example, if an image is 100 x 200 pixels (height x width) and the bounding
+// box is `[0.1, 0.2, 0.5, 0.9]`, the upper-left and bottom-right coordinates of
+// the bounding box will be `(40, 10)` to `(180, 50)` (in (x,y) coordinates).
+//
+// Parts of the bounding box may fall outside the image.
+//
+// Arguments:
+// images: 4-D with shape `[batch, height, width, depth]`. A batch of images.
+// boxes: 3-D with shape `[batch, num_bounding_boxes, 4]` containing bounding
+// boxes.
+//
+// Returns 4-D with the same shape as `images`. The batch of input images with
+// bounding boxes drawn on the images.
+func DrawBoundingBoxes(scope *Scope, images tf.Output, boxes tf.Output) (output tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "DrawBoundingBoxes",
+ Input: []tf.Input{
+ images, boxes,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
// Gets the next output from the given iterator.
//
// This operation is a synchronous version IteratorGetNext. It should only be used
@@ -27436,178 +28163,6 @@
return op.Output(0)
}
-// EncodeProtoAttr is an optional argument to EncodeProto.
-type EncodeProtoAttr func(optionalAttr)
-
-// EncodeProtoDescriptorSource sets the optional descriptor_source attribute to value.
-// If not specified, defaults to "local://"
-func EncodeProtoDescriptorSource(value string) EncodeProtoAttr {
- return func(m optionalAttr) {
- m["descriptor_source"] = value
- }
-}
-
-// The op serializes protobuf messages provided in the input tensors.
-//
-// The types of the tensors in `values` must match the schema for the
-// fields specified in `field_names`. All the tensors in `values` must
-// have a common shape prefix, *batch_shape*.
-//
-// The `sizes` tensor specifies repeat counts for each field. The repeat
-// count (last dimension) of a each tensor in `values` must be greater
-// than or equal to corresponding repeat count in `sizes`.
-//
-// A `message_type` name must be provided to give context for the field
-// names. The actual message descriptor can be looked up either in the
-// linked-in descriptor pool or a filename provided by the caller using
-// the `descriptor_source` attribute.
-//
-// The `descriptor_source` attribute selects a source of protocol
-// descriptors to consult when looking up `message_type`. This may be a
-// filename containing a serialized `FileDescriptorSet` message,
-// or the special value `local://`, in which case only descriptors linked
-// into the code will be searched; the filename can be on any filesystem
-// accessible to TensorFlow.
-//
-// You can build a `descriptor_source` file using the `--descriptor_set_out`
-// and `--include_imports` options to the protocol compiler `protoc`.
-//
-// The `local://` database only covers descriptors linked into the
-// code via C++ libraries, not Python imports. You can link in a proto descriptor
-// by creating a cc_library target with alwayslink=1.
-//
-// There are a few special cases in the value mapping:
-//
-// Submessage and group fields must be pre-serialized as TensorFlow strings.
-//
-// TensorFlow lacks support for unsigned int64s, so they must be
-// represented as `tf.int64` with the same twos-complement bit pattern
-// (the obvious way).
-//
-// Unsigned int32 values can be represented exactly with `tf.int64`, or
-// with sign wrapping if the input is of type `tf.int32`.
-//
-// Arguments:
-// sizes: Tensor of int32 with shape `[batch_shape, len(field_names)]`.
-// values: List of tensors containing values for the corresponding field.
-// field_names: List of strings containing proto field names.
-// message_type: Name of the proto message type to decode.
-//
-// Returns Tensor of serialized protos with shape `batch_shape`.
-func EncodeProto(scope *Scope, sizes tf.Output, values []tf.Output, field_names []string, message_type string, optional ...EncodeProtoAttr) (bytes tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"field_names": field_names, "message_type": message_type}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "EncodeProto",
- Input: []tf.Input{
- sizes, tf.OutputList(values),
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// Creates a TensorArray for storing the gradients of values in the given handle.
-//
-// If the given TensorArray gradient already exists, returns a reference to it.
-//
-// Locks the size of the original TensorArray by disabling its dynamic size flag.
-//
-// **A note about the input flow_in:**
-//
-// The handle flow_in forces the execution of the gradient lookup to occur
-// only after certain other operations have occurred. For example, when
-// the forward TensorArray is dynamically sized, writes to this TensorArray
-// may resize the object. The gradient TensorArray is statically sized based
-// on the size of the forward TensorArray when this operation executes.
-// Furthermore, the size of the forward TensorArray is frozen by this call.
-// As a result, the flow is used to ensure that the call to generate the gradient
-// TensorArray only happens after all writes are executed.
-//
-// In the case of dynamically sized TensorArrays, gradient computation should
-// only be performed on read operations that have themselves been chained via
-// flow to occur only after all writes have executed. That way the final size
-// of the forward TensorArray is known when this operation is called.
-//
-// **A note about the source attribute:**
-//
-// TensorArray gradient calls use an accumulator TensorArray object. If
-// multiple gradients are calculated and run in the same session, the multiple
-// gradient nodes may accidentally flow through the same accumulator TensorArray.
-// This double counts and generally breaks the TensorArray gradient flow.
-//
-// The solution is to identify which gradient call this particular
-// TensorArray gradient is being called in. This is performed by identifying
-// a unique string (e.g. "gradients", "gradients_1", ...) from the input
-// gradient Tensor's name. This string is used as a suffix when creating
-// the TensorArray gradient object here (the attribute `source`).
-//
-// The attribute `source` is added as a suffix to the forward TensorArray's
-// name when performing the creation / lookup, so that each separate gradient
-// calculation gets its own TensorArray accumulator.
-//
-// Arguments:
-// handle: The handle to the forward TensorArray.
-// flow_in: A float scalar that enforces proper chaining of operations.
-// source: The gradient source string, used to decide which gradient TensorArray
-// to return.
-func TensorArrayGradV3(scope *Scope, handle tf.Output, flow_in tf.Output, source string) (grad_handle tf.Output, flow_out tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"source": source}
- opspec := tf.OpSpec{
- Type: "TensorArrayGradV3",
- Input: []tf.Input{
- handle, flow_in,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0), op.Output(1)
-}
-
-// Creates a dataset that splits a SparseTensor into elements row-wise.
-func SparseTensorSliceDataset(scope *Scope, indices tf.Output, values tf.Output, dense_shape tf.Output) (handle tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "SparseTensorSliceDataset",
- Input: []tf.Input{
- indices, values, dense_shape,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// Returns x / y element-wise for real types.
-//
-// If `x` and `y` are reals, this will return the floating-point division.
-//
-// *NOTE*: `Div` supports broadcasting. More about broadcasting
-// [here](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html)
-func RealDiv(scope *Scope, x tf.Output, y tf.Output) (z tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "RealDiv",
- Input: []tf.Input{
- x, y,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
// Adds v into specified rows of x.
//
// Computes y = x; y[i, :] += v; return y.
@@ -27903,6 +28458,255 @@
return op.Output(0)
}
+// StringSplitV2Attr is an optional argument to StringSplitV2.
+type StringSplitV2Attr func(optionalAttr)
+
+// StringSplitV2Maxsplit sets the optional maxsplit attribute to value.
+//
+// value: An `int`. If `maxsplit > 0`, limit of the split of the result.
+// If not specified, defaults to -1
+func StringSplitV2Maxsplit(value int64) StringSplitV2Attr {
+ return func(m optionalAttr) {
+ m["maxsplit"] = value
+ }
+}
+
+// Split elements of `source` based on `sep` into a `SparseTensor`.
+//
+// Let N be the size of source (typically N will be the batch size). Split each
+// element of `source` based on `sep` and return a `SparseTensor`
+// containing the split tokens. Empty tokens are ignored.
+//
+// For example, N = 2, source[0] is 'hello world' and source[1] is 'a b c',
+// then the output will be
+// ```
+// st.indices = [0, 0;
+// 0, 1;
+// 1, 0;
+// 1, 1;
+// 1, 2]
+// st.shape = [2, 3]
+// st.values = ['hello', 'world', 'a', 'b', 'c']
+// ```
+//
+// If `sep` is given, consecutive delimiters are not grouped together and are
+// deemed to delimit empty strings. For example, source of `"1<>2<><>3"` and
+// sep of `"<>"` returns `["1", "2", "", "3"]`. If `sep` is None or an empty
+// string, consecutive whitespace are regarded as a single separator, and the
+// result will contain no empty strings at the startor end if the string has
+// leading or trailing whitespace.
+//
+// Note that the above mentioned behavior matches python's str.split.
+//
+// Arguments:
+// input: `1-D` string `Tensor`, the strings to split.
+// sep: `0-D` string `Tensor`, the delimiter character.
+func StringSplitV2(scope *Scope, input tf.Output, sep tf.Output, optional ...StringSplitV2Attr) (indices tf.Output, values tf.Output, shape tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "StringSplitV2",
+ Input: []tf.Input{
+ input, sep,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0), op.Output(1), op.Output(2)
+}
+
+// Computes softsign: `features / (abs(features) + 1)`.
+func Softsign(scope *Scope, features tf.Output) (activations tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "Softsign",
+ Input: []tf.Input{
+ features,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// EncodeProtoAttr is an optional argument to EncodeProto.
+type EncodeProtoAttr func(optionalAttr)
+
+// EncodeProtoDescriptorSource sets the optional descriptor_source attribute to value.
+// If not specified, defaults to "local://"
+func EncodeProtoDescriptorSource(value string) EncodeProtoAttr {
+ return func(m optionalAttr) {
+ m["descriptor_source"] = value
+ }
+}
+
+// The op serializes protobuf messages provided in the input tensors.
+//
+// The types of the tensors in `values` must match the schema for the
+// fields specified in `field_names`. All the tensors in `values` must
+// have a common shape prefix, *batch_shape*.
+//
+// The `sizes` tensor specifies repeat counts for each field. The repeat
+// count (last dimension) of a each tensor in `values` must be greater
+// than or equal to corresponding repeat count in `sizes`.
+//
+// A `message_type` name must be provided to give context for the field
+// names. The actual message descriptor can be looked up either in the
+// linked-in descriptor pool or a filename provided by the caller using
+// the `descriptor_source` attribute.
+//
+// The `descriptor_source` attribute selects a source of protocol
+// descriptors to consult when looking up `message_type`. This may be a
+// filename containing a serialized `FileDescriptorSet` message,
+// or the special value `local://`, in which case only descriptors linked
+// into the code will be searched; the filename can be on any filesystem
+// accessible to TensorFlow.
+//
+// You can build a `descriptor_source` file using the `--descriptor_set_out`
+// and `--include_imports` options to the protocol compiler `protoc`.
+//
+// The `local://` database only covers descriptors linked into the
+// code via C++ libraries, not Python imports. You can link in a proto descriptor
+// by creating a cc_library target with alwayslink=1.
+//
+// There are a few special cases in the value mapping:
+//
+// Submessage and group fields must be pre-serialized as TensorFlow strings.
+//
+// TensorFlow lacks support for unsigned int64s, so they must be
+// represented as `tf.int64` with the same twos-complement bit pattern
+// (the obvious way).
+//
+// Unsigned int32 values can be represented exactly with `tf.int64`, or
+// with sign wrapping if the input is of type `tf.int32`.
+//
+// Arguments:
+// sizes: Tensor of int32 with shape `[batch_shape, len(field_names)]`.
+// values: List of tensors containing values for the corresponding field.
+// field_names: List of strings containing proto field names.
+// message_type: Name of the proto message type to decode.
+//
+// Returns Tensor of serialized protos with shape `batch_shape`.
+func EncodeProto(scope *Scope, sizes tf.Output, values []tf.Output, field_names []string, message_type string, optional ...EncodeProtoAttr) (bytes tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"field_names": field_names, "message_type": message_type}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "EncodeProto",
+ Input: []tf.Input{
+ sizes, tf.OutputList(values),
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// Creates a TensorArray for storing the gradients of values in the given handle.
+//
+// If the given TensorArray gradient already exists, returns a reference to it.
+//
+// Locks the size of the original TensorArray by disabling its dynamic size flag.
+//
+// **A note about the input flow_in:**
+//
+// The handle flow_in forces the execution of the gradient lookup to occur
+// only after certain other operations have occurred. For example, when
+// the forward TensorArray is dynamically sized, writes to this TensorArray
+// may resize the object. The gradient TensorArray is statically sized based
+// on the size of the forward TensorArray when this operation executes.
+// Furthermore, the size of the forward TensorArray is frozen by this call.
+// As a result, the flow is used to ensure that the call to generate the gradient
+// TensorArray only happens after all writes are executed.
+//
+// In the case of dynamically sized TensorArrays, gradient computation should
+// only be performed on read operations that have themselves been chained via
+// flow to occur only after all writes have executed. That way the final size
+// of the forward TensorArray is known when this operation is called.
+//
+// **A note about the source attribute:**
+//
+// TensorArray gradient calls use an accumulator TensorArray object. If
+// multiple gradients are calculated and run in the same session, the multiple
+// gradient nodes may accidentally flow through the same accumulator TensorArray.
+// This double counts and generally breaks the TensorArray gradient flow.
+//
+// The solution is to identify which gradient call this particular
+// TensorArray gradient is being called in. This is performed by identifying
+// a unique string (e.g. "gradients", "gradients_1", ...) from the input
+// gradient Tensor's name. This string is used as a suffix when creating
+// the TensorArray gradient object here (the attribute `source`).
+//
+// The attribute `source` is added as a suffix to the forward TensorArray's
+// name when performing the creation / lookup, so that each separate gradient
+// calculation gets its own TensorArray accumulator.
+//
+// Arguments:
+// handle: The handle to the forward TensorArray.
+// flow_in: A float scalar that enforces proper chaining of operations.
+// source: The gradient source string, used to decide which gradient TensorArray
+// to return.
+func TensorArrayGradV3(scope *Scope, handle tf.Output, flow_in tf.Output, source string) (grad_handle tf.Output, flow_out tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"source": source}
+ opspec := tf.OpSpec{
+ Type: "TensorArrayGradV3",
+ Input: []tf.Input{
+ handle, flow_in,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0), op.Output(1)
+}
+
+// Creates a dataset that splits a SparseTensor into elements row-wise.
+func SparseTensorSliceDataset(scope *Scope, indices tf.Output, values tf.Output, dense_shape tf.Output) (handle tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "SparseTensorSliceDataset",
+ Input: []tf.Input{
+ indices, values, dense_shape,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// Returns x / y element-wise for real types.
+//
+// If `x` and `y` are reals, this will return the floating-point division.
+//
+// *NOTE*: `Div` supports broadcasting. More about broadcasting
+// [here](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html)
+func RealDiv(scope *Scope, x tf.Output, y tf.Output) (z tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "RealDiv",
+ Input: []tf.Input{
+ x, y,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
// Creates a dataset that concatenates `input_dataset` with `another_dataset`.
func ConcatenateDataset(scope *Scope, input_dataset tf.Output, another_dataset tf.Output, output_types []tf.DataType, output_shapes []tf.Shape) (handle tf.Output) {
if scope.Err() != nil {
@@ -30829,260 +31633,6 @@
return op.Output(0)
}
-// SerializeSparseAttr is an optional argument to SerializeSparse.
-type SerializeSparseAttr func(optionalAttr)
-
-// SerializeSparseOutType 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 SerializeSparseOutType(value tf.DataType) SerializeSparseAttr {
- return func(m optionalAttr) {
- m["out_type"] = value
- }
-}
-
-// Serialize a `SparseTensor` into a `[3]` `Tensor` object.
-//
-// Arguments:
-// sparse_indices: 2-D. The `indices` of the `SparseTensor`.
-// sparse_values: 1-D. The `values` of the `SparseTensor`.
-// sparse_shape: 1-D. The `shape` of the `SparseTensor`.
-func SerializeSparse(scope *Scope, sparse_indices tf.Output, sparse_values tf.Output, sparse_shape tf.Output, optional ...SerializeSparseAttr) (serialized_sparse tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "SerializeSparse",
- Input: []tf.Input{
- sparse_indices, sparse_values, sparse_shape,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// RandomShuffleQueueV2Attr is an optional argument to RandomShuffleQueueV2.
-type RandomShuffleQueueV2Attr func(optionalAttr)
-
-// RandomShuffleQueueV2Shapes sets the optional shapes attribute to value.
-//
-// value: The shape of each component in a value. The length of this attr must
-// be either 0 or the same as the length of component_types. If the length of
-// this attr is 0, the shapes of queue elements are not constrained, and
-// only one element may be dequeued at a time.
-// If not specified, defaults to <>
-//
-// REQUIRES: len(value) >= 0
-func RandomShuffleQueueV2Shapes(value []tf.Shape) RandomShuffleQueueV2Attr {
- return func(m optionalAttr) {
- m["shapes"] = value
- }
-}
-
-// RandomShuffleQueueV2Capacity sets the optional capacity attribute to value.
-//
-// value: The upper bound on the number of elements in this queue.
-// Negative numbers mean no limit.
-// If not specified, defaults to -1
-func RandomShuffleQueueV2Capacity(value int64) RandomShuffleQueueV2Attr {
- return func(m optionalAttr) {
- m["capacity"] = value
- }
-}
-
-// RandomShuffleQueueV2MinAfterDequeue sets the optional min_after_dequeue attribute to value.
-//
-// value: Dequeue will block unless there would be this
-// many elements after the dequeue or the queue is closed. This
-// ensures a minimum level of mixing of elements.
-// If not specified, defaults to 0
-func RandomShuffleQueueV2MinAfterDequeue(value int64) RandomShuffleQueueV2Attr {
- return func(m optionalAttr) {
- m["min_after_dequeue"] = value
- }
-}
-
-// RandomShuffleQueueV2Seed sets the optional seed attribute to value.
-//
-// value: If either seed or seed2 is set to be non-zero, the random number
-// generator is seeded by the given seed. Otherwise, a random seed is used.
-// If not specified, defaults to 0
-func RandomShuffleQueueV2Seed(value int64) RandomShuffleQueueV2Attr {
- return func(m optionalAttr) {
- m["seed"] = value
- }
-}
-
-// RandomShuffleQueueV2Seed2 sets the optional seed2 attribute to value.
-//
-// value: A second seed to avoid seed collision.
-// If not specified, defaults to 0
-func RandomShuffleQueueV2Seed2(value int64) RandomShuffleQueueV2Attr {
- return func(m optionalAttr) {
- m["seed2"] = value
- }
-}
-
-// RandomShuffleQueueV2Container 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 RandomShuffleQueueV2Container(value string) RandomShuffleQueueV2Attr {
- return func(m optionalAttr) {
- m["container"] = value
- }
-}
-
-// RandomShuffleQueueV2SharedName sets the optional shared_name attribute to value.
-//
-// value: If non-empty, this queue will be shared under the given name
-// across multiple sessions.
-// If not specified, defaults to ""
-func RandomShuffleQueueV2SharedName(value string) RandomShuffleQueueV2Attr {
- return func(m optionalAttr) {
- m["shared_name"] = value
- }
-}
-
-// A queue that randomizes the order of elements.
-//
-// Arguments:
-// component_types: The type of each component in a value.
-//
-// Returns The handle to the queue.
-func RandomShuffleQueueV2(scope *Scope, component_types []tf.DataType, optional ...RandomShuffleQueueV2Attr) (handle tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"component_types": component_types}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "RandomShuffleQueueV2",
-
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// Draw bounding boxes on a batch of images.
-//
-// Outputs a copy of `images` but draws on top of the pixels zero or more bounding
-// boxes specified by the locations in `boxes`. The coordinates of the each
-// bounding box in `boxes` are encoded as `[y_min, x_min, y_max, x_max]`. The
-// bounding box coordinates are floats in `[0.0, 1.0]` relative to the width and
-// height of the underlying image.
-//
-// For example, if an image is 100 x 200 pixels (height x width) and the bounding
-// box is `[0.1, 0.2, 0.5, 0.9]`, the upper-left and bottom-right coordinates of
-// the bounding box will be `(40, 10)` to `(180, 50)` (in (x,y) coordinates).
-//
-// Parts of the bounding box may fall outside the image.
-//
-// Arguments:
-// images: 4-D with shape `[batch, height, width, depth]`. A batch of images.
-// boxes: 3-D with shape `[batch, num_bounding_boxes, 4]` containing bounding
-// boxes.
-//
-// Returns 4-D with the same shape as `images`. The batch of input images with
-// bounding boxes drawn on the images.
-func DrawBoundingBoxes(scope *Scope, images tf.Output, boxes tf.Output) (output tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "DrawBoundingBoxes",
- Input: []tf.Input{
- images, boxes,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// LearnedUnigramCandidateSamplerAttr is an optional argument to LearnedUnigramCandidateSampler.
-type LearnedUnigramCandidateSamplerAttr func(optionalAttr)
-
-// LearnedUnigramCandidateSamplerSeed 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 LearnedUnigramCandidateSamplerSeed(value int64) LearnedUnigramCandidateSamplerAttr {
- return func(m optionalAttr) {
- m["seed"] = value
- }
-}
-
-// LearnedUnigramCandidateSamplerSeed2 sets the optional seed2 attribute to value.
-//
-// value: An second seed to avoid seed collision.
-// If not specified, defaults to 0
-func LearnedUnigramCandidateSamplerSeed2(value int64) LearnedUnigramCandidateSamplerAttr {
- 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 LearnedUnigramCandidateSampler(scope *Scope, true_classes tf.Output, num_true int64, num_sampled int64, unique bool, range_max int64, optional ...LearnedUnigramCandidateSamplerAttr) (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: "LearnedUnigramCandidateSampler",
- Input: []tf.Input{
- true_classes,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0), op.Output(1), op.Output(2)
-}
-
// Computes gradients for the scaled exponential linear (Selu) operation.
//
// Arguments:
@@ -32441,79 +32991,6 @@
return weights, biases
}
-// UniformCandidateSamplerAttr is an optional argument to UniformCandidateSampler.
-type UniformCandidateSamplerAttr func(optionalAttr)
-
-// UniformCandidateSamplerSeed 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 UniformCandidateSamplerSeed(value int64) UniformCandidateSamplerAttr {
- return func(m optionalAttr) {
- m["seed"] = value
- }
-}
-
-// UniformCandidateSamplerSeed2 sets the optional seed2 attribute to value.
-//
-// value: An second seed to avoid seed collision.
-// If not specified, defaults to 0
-func UniformCandidateSamplerSeed2(value int64) UniformCandidateSamplerAttr {
- return func(m optionalAttr) {
- m["seed2"] = value
- }
-}
-
-// Generates labels for candidate sampling with a uniform 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 UniformCandidateSampler(scope *Scope, true_classes tf.Output, num_true int64, num_sampled int64, unique bool, range_max int64, optional ...UniformCandidateSamplerAttr) (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: "UniformCandidateSampler",
- Input: []tf.Input{
- true_classes,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0), op.Output(1), op.Output(2)
-}
-
// CTCLossAttr is an optional argument to CTCLoss.
type CTCLossAttr func(optionalAttr)
@@ -32664,480 +33141,3 @@
op := scope.AddOperation(opspec)
return op.Output(0), op.Output(1)
}
-
-// Add all input tensors element wise.
-//
-// Arguments:
-// inputs: Must all be the same size and shape.
-func AddN(scope *Scope, inputs []tf.Output) (sum tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "AddN",
- Input: []tf.Input{
- tf.OutputList(inputs),
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// TryRpcAttr is an optional argument to TryRpc.
-type TryRpcAttr func(optionalAttr)
-
-// TryRpcProtocol sets the optional protocol attribute to value.
-//
-// value: RPC protocol to use. Empty string means use the default protocol.
-// Options include 'grpc'.
-// If not specified, defaults to ""
-func TryRpcProtocol(value string) TryRpcAttr {
- return func(m optionalAttr) {
- m["protocol"] = value
- }
-}
-
-// TryRpcFailFast sets the optional fail_fast attribute to value.
-//
-// value: `boolean`. If `true` (default), then failures to connect
-// (i.e., the server does not immediately respond) cause an RPC failure.
-// If not specified, defaults to true
-func TryRpcFailFast(value bool) TryRpcAttr {
- return func(m optionalAttr) {
- m["fail_fast"] = value
- }
-}
-
-// TryRpcTimeoutInMs sets the optional timeout_in_ms attribute to value.
-//
-// value: `int`. If `0` (default), then the kernel will run the RPC
-// request and only time out if the RPC deadline passes or the session times out.
-// If this value is greater than `0`, then the op will raise an exception if
-// the RPC takes longer than `timeout_in_ms`.
-// If not specified, defaults to 0
-func TryRpcTimeoutInMs(value int64) TryRpcAttr {
- return func(m optionalAttr) {
- m["timeout_in_ms"] = value
- }
-}
-
-// Perform batches of RPC requests.
-//
-// This op asynchronously performs either a single RPC request, or a batch
-// of requests. RPC requests are defined by three main parameters:
-//
-// - `address` (the host+port or BNS address of the request)
-// - `method` (the method name for the request)
-// - `request` (the serialized proto string, or vector of strings,
-// of the RPC request argument).
-//
-// For example, if you have an RPC service running on port localhost:2345,
-// and its interface is configured with the following proto declaration:
-//
-// ```
-// service MyService {
-// rpc MyMethod(MyRequestProto) returns (MyResponseProto) {
-// }
-// };
-// ```
-//
-// then call this op with arguments:
-//
-// ```
-// address = "localhost:2345"
-// method = "MyService/MyMethod"
-// ```
-//
-// The `request` tensor is a string tensor representing serialized `MyRequestProto`
-// strings; and the output string tensor `response` will have the same shape
-// and contain (upon successful completion) corresponding serialized
-// `MyResponseProto` strings.
-//
-// For example, to send a single, empty, `MyRequestProto`, call
-// this op with `request = ""`. To send 5 **parallel** empty requests,
-// call this op with `request = ["", "", "", "", ""]`.
-//
-// More generally, one can create a batch of `MyRequestProto` serialized protos
-// from regular batched tensors using the `encode_proto` op, and convert
-// the response `MyResponseProto` serialized protos to batched tensors
-// using the `decode_proto` op.
-//
-// **NOTE** Working with serialized proto strings is faster than instantiating
-// actual proto objects in memory, so no performance degradation is expected
-// compared to writing custom kernels for this workflow.
-//
-// Unlike the standard `Rpc` op, if the connection fails or the remote worker
-// returns an error status, this op does **not** reraise the exception.
-// Instead, the `status_code` and `status_message` entry for the corresponding RPC
-// call is set with the error returned from the RPC call. The `response` tensor
-// will contain valid response values for those minibatch entries whose RPCs did
-// not fail; the rest of the entries will have empty strings.
-//
-// Arguments:
-// address: `0-D` or `1-D`. The address (i.e. host_name:port) of the RPC server.
-// If this tensor has more than 1 element, then multiple parallel rpc requests
-// are sent. This argument broadcasts with `method` and `request`.
-// method: `0-D` or `1-D`. The method address on the RPC server.
-// If this tensor has more than 1 element, then multiple parallel rpc requests
-// are sent. This argument broadcasts with `address` and `request`.
-// request: `0-D` or `1-D`. Serialized proto strings: the rpc request argument.
-// If this tensor has more than 1 element, then multiple parallel rpc requests
-// are sent. This argument broadcasts with `address` and `method`.
-//
-// Returns Same shape as `request`. Serialized proto strings: the rpc responses.Same shape as `request`. Values correspond to tensorflow Status enum codes.Same shape as `request`. Values correspond to Status messages
-// returned from the RPC calls.
-func TryRpc(scope *Scope, address tf.Output, method tf.Output, request tf.Output, optional ...TryRpcAttr) (response tf.Output, status_code tf.Output, status_message tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "TryRpc",
- Input: []tf.Input{
- address, method, request,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0), op.Output(1), op.Output(2)
-}
-
-// EnterAttr is an optional argument to Enter.
-type EnterAttr func(optionalAttr)
-
-// EnterIsConstant sets the optional is_constant attribute to value.
-//
-// value: If true, the output is constant within the child frame.
-// If not specified, defaults to false
-func EnterIsConstant(value bool) EnterAttr {
- return func(m optionalAttr) {
- m["is_constant"] = value
- }
-}
-
-// EnterParallelIterations sets the optional parallel_iterations attribute to value.
-//
-// value: The number of iterations allowed to run in parallel.
-// If not specified, defaults to 10
-func EnterParallelIterations(value int64) EnterAttr {
- return func(m optionalAttr) {
- m["parallel_iterations"] = value
- }
-}
-
-// Creates or finds a child frame, and makes `data` available to the child frame.
-//
-// This op is used together with `Exit` to create loops in the graph.
-// The unique `frame_name` is used by the `Executor` to identify frames. If
-// `is_constant` is true, `output` is a constant in the child frame; otherwise
-// it may be changed in the child frame. At most `parallel_iterations` iterations
-// are run in parallel in the child frame.
-//
-// Arguments:
-// data: The tensor to be made available to the child frame.
-// frame_name: The name of the child frame.
-//
-// Returns The same tensor as `data`.
-func Enter(scope *Scope, data tf.Output, frame_name string, optional ...EnterAttr) (output tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"frame_name": frame_name}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "Enter",
- Input: []tf.Input{
- data,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// Produce a string tensor that encodes the state of a Reader.
-//
-// Not all Readers support being serialized, so this can produce an
-// Unimplemented error.
-//
-// Arguments:
-// reader_handle: Handle to a Reader.
-func ReaderSerializeStateV2(scope *Scope, reader_handle tf.Output) (state tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "ReaderSerializeStateV2",
- Input: []tf.Input{
- reader_handle,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// Exits the current frame to its parent frame.
-//
-// Exit makes its input `data` available to the parent frame.
-//
-// Arguments:
-// data: The tensor to be made available to the parent frame.
-//
-// Returns The same tensor as `data`.
-func Exit(scope *Scope, data tf.Output) (output tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "Exit",
- Input: []tf.Input{
- data,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// Returns a copy of the input tensor.
-func Snapshot(scope *Scope, input tf.Output) (output tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "Snapshot",
- Input: []tf.Input{
- input,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// Returns a tensor of zeros with the same shape and type as x.
-//
-// Arguments:
-// x: a tensor of type T.
-//
-// Returns a tensor of the same shape and type as x but filled with zeros.
-func ZerosLike(scope *Scope, x tf.Output) (y tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "ZerosLike",
- Input: []tf.Input{
- x,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// AbortAttr is an optional argument to Abort.
-type AbortAttr func(optionalAttr)
-
-// AbortErrorMsg sets the optional error_msg attribute to value.
-//
-// value: A string which is the message associated with the exception.
-// If not specified, defaults to ""
-func AbortErrorMsg(value string) AbortAttr {
- return func(m optionalAttr) {
- m["error_msg"] = value
- }
-}
-
-// AbortExitWithoutError sets the optional exit_without_error attribute to value.
-// If not specified, defaults to false
-func AbortExitWithoutError(value bool) AbortAttr {
- return func(m optionalAttr) {
- m["exit_without_error"] = value
- }
-}
-
-// Raise a exception to abort the process when called.
-//
-// If exit_without_error is true, the process will exit normally,
-// otherwise it will exit with a SIGABORT signal.
-//
-// Returns nothing but an exception.
-//
-// Returns the created operation.
-func Abort(scope *Scope, optional ...AbortAttr) (o *tf.Operation) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "Abort",
-
- Attrs: attrs,
- }
- return scope.AddOperation(opspec)
-}
-
-// FixedUnigramCandidateSamplerAttr is an optional argument to FixedUnigramCandidateSampler.
-type FixedUnigramCandidateSamplerAttr func(optionalAttr)
-
-// FixedUnigramCandidateSamplerVocabFile sets the optional vocab_file attribute to value.
-//
-// value: Each valid line in this file (which should have a CSV-like format)
-// corresponds to a valid word ID. IDs are in sequential order, starting from
-// num_reserved_ids. The last entry in each line is expected to be a value
-// corresponding to the count or relative probability. Exactly one of vocab_file
-// and unigrams needs to be passed to this op.
-// If not specified, defaults to ""
-func FixedUnigramCandidateSamplerVocabFile(value string) FixedUnigramCandidateSamplerAttr {
- return func(m optionalAttr) {
- m["vocab_file"] = value
- }
-}
-
-// FixedUnigramCandidateSamplerDistortion sets the optional distortion attribute to value.
-//
-// value: The distortion is used to skew the unigram probability distribution.
-// Each weight is first raised to the distortion's power before adding to the
-// internal unigram distribution. As a result, distortion = 1.0 gives regular
-// unigram sampling (as defined by the vocab file), and distortion = 0.0 gives
-// a uniform distribution.
-// If not specified, defaults to 1
-func FixedUnigramCandidateSamplerDistortion(value float32) FixedUnigramCandidateSamplerAttr {
- return func(m optionalAttr) {
- m["distortion"] = value
- }
-}
-
-// FixedUnigramCandidateSamplerNumReservedIds sets the optional num_reserved_ids attribute to value.
-//
-// value: Optionally some reserved IDs can be added in the range [0,
-// ..., num_reserved_ids) by the users. One use case is that a special unknown
-// word token is used as ID 0. These IDs will have a sampling probability of 0.
-// If not specified, defaults to 0
-func FixedUnigramCandidateSamplerNumReservedIds(value int64) FixedUnigramCandidateSamplerAttr {
- return func(m optionalAttr) {
- m["num_reserved_ids"] = value
- }
-}
-
-// FixedUnigramCandidateSamplerNumShards sets the optional num_shards attribute to value.
-//
-// value: A sampler can be used to sample from a subset of the original range
-// in order to speed up the whole computation through parallelism. This parameter
-// (together with 'shard') indicates the number of partitions that are being
-// used in the overall computation.
-// If not specified, defaults to 1
-//
-// REQUIRES: value >= 1
-func FixedUnigramCandidateSamplerNumShards(value int64) FixedUnigramCandidateSamplerAttr {
- return func(m optionalAttr) {
- m["num_shards"] = value
- }
-}
-
-// FixedUnigramCandidateSamplerShard sets the optional shard attribute to value.
-//
-// value: A sampler can be used to sample from a subset of the original range
-// in order to speed up the whole computation through parallelism. This parameter
-// (together with 'num_shards') indicates the particular partition number of a
-// sampler op, when partitioning is being used.
-// If not specified, defaults to 0
-//
-// REQUIRES: value >= 0
-func FixedUnigramCandidateSamplerShard(value int64) FixedUnigramCandidateSamplerAttr {
- return func(m optionalAttr) {
- m["shard"] = value
- }
-}
-
-// FixedUnigramCandidateSamplerUnigrams sets the optional unigrams attribute to value.
-//
-// value: A list of unigram counts or probabilities, one per ID in sequential
-// order. Exactly one of vocab_file and unigrams should be passed to this op.
-// If not specified, defaults to <>
-func FixedUnigramCandidateSamplerUnigrams(value []float32) FixedUnigramCandidateSamplerAttr {
- return func(m optionalAttr) {
- m["unigrams"] = value
- }
-}
-
-// FixedUnigramCandidateSamplerSeed 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 FixedUnigramCandidateSamplerSeed(value int64) FixedUnigramCandidateSamplerAttr {
- return func(m optionalAttr) {
- m["seed"] = value
- }
-}
-
-// FixedUnigramCandidateSamplerSeed2 sets the optional seed2 attribute to value.
-//
-// value: An second seed to avoid seed collision.
-// If not specified, defaults to 0
-func FixedUnigramCandidateSamplerSeed2(value int64) FixedUnigramCandidateSamplerAttr {
- return func(m optionalAttr) {
- m["seed2"] = value
- }
-}
-
-// Generates labels for candidate sampling with a learned unigram distribution.
-//
-// A unigram sampler could use a fixed unigram distribution read from a
-// file or passed in as an in-memory array instead of building up the distribution
-// from data on the fly. There is also an option to skew the distribution by
-// applying a distortion power to the weights.
-//
-// The vocabulary file should be in CSV-like format, with the last field
-// being the weight associated with the word.
-//
-// 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 FixedUnigramCandidateSampler(scope *Scope, true_classes tf.Output, num_true int64, num_sampled int64, unique bool, range_max int64, optional ...FixedUnigramCandidateSamplerAttr) (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: "FixedUnigramCandidateSampler",
- Input: []tf.Input{
- true_classes,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0), op.Output(1), op.Output(2)
-}
diff --git a/tensorflow/java/maven/libtensorflow/pom.xml b/tensorflow/java/maven/libtensorflow/pom.xml
index 9fc6969..6b3e305 100644
--- a/tensorflow/java/maven/libtensorflow/pom.xml
+++ b/tensorflow/java/maven/libtensorflow/pom.xml
@@ -6,7 +6,7 @@
<parent>
<groupId>org.tensorflow</groupId>
<artifactId>parentpom</artifactId>
- <version>1.11.0-rc2</version>
+ <version>1.11.0</version>
<relativePath>../</relativePath>
</parent>
<artifactId>libtensorflow</artifactId>
diff --git a/tensorflow/java/maven/libtensorflow_jni/pom.xml b/tensorflow/java/maven/libtensorflow_jni/pom.xml
index 6871208..f130515 100644
--- a/tensorflow/java/maven/libtensorflow_jni/pom.xml
+++ b/tensorflow/java/maven/libtensorflow_jni/pom.xml
@@ -6,7 +6,7 @@
<parent>
<groupId>org.tensorflow</groupId>
<artifactId>parentpom</artifactId>
- <version>1.11.0-rc2</version>
+ <version>1.11.0</version>
<relativePath>../</relativePath>
</parent>
<artifactId>libtensorflow_jni</artifactId>
diff --git a/tensorflow/java/maven/libtensorflow_jni_gpu/pom.xml b/tensorflow/java/maven/libtensorflow_jni_gpu/pom.xml
index f031173..67ecc2d 100644
--- a/tensorflow/java/maven/libtensorflow_jni_gpu/pom.xml
+++ b/tensorflow/java/maven/libtensorflow_jni_gpu/pom.xml
@@ -6,7 +6,7 @@
<parent>
<groupId>org.tensorflow</groupId>
<artifactId>parentpom</artifactId>
- <version>1.11.0-rc2</version>
+ <version>1.11.0</version>
<relativePath>../</relativePath>
</parent>
<artifactId>libtensorflow_jni_gpu</artifactId>
diff --git a/tensorflow/java/maven/pom.xml b/tensorflow/java/maven/pom.xml
index 2cac279..8ba859d 100644
--- a/tensorflow/java/maven/pom.xml
+++ b/tensorflow/java/maven/pom.xml
@@ -6,7 +6,7 @@
<modelVersion>4.0.0</modelVersion>
<groupId>org.tensorflow</groupId>
<artifactId>parentpom</artifactId>
- <version>1.11.0-rc2</version>
+ <version>1.11.0</version>
<packaging>pom</packaging>
<url>https://www.tensorflow.org</url>
diff --git a/tensorflow/java/maven/proto/pom.xml b/tensorflow/java/maven/proto/pom.xml
index 8a93091..dcd654d 100644
--- a/tensorflow/java/maven/proto/pom.xml
+++ b/tensorflow/java/maven/proto/pom.xml
@@ -6,7 +6,7 @@
<parent>
<groupId>org.tensorflow</groupId>
<artifactId>parentpom</artifactId>
- <version>1.11.0-rc2</version>
+ <version>1.11.0</version>
<relativePath>../</relativePath>
</parent>
<artifactId>proto</artifactId>
diff --git a/tensorflow/java/maven/spark-tensorflow-connector/pom.xml b/tensorflow/java/maven/spark-tensorflow-connector/pom.xml
index 014bd8d..45214f8 100644
--- a/tensorflow/java/maven/spark-tensorflow-connector/pom.xml
+++ b/tensorflow/java/maven/spark-tensorflow-connector/pom.xml
@@ -6,7 +6,7 @@
<groupId>org.tensorflow</groupId>
<artifactId>spark-tensorflow-connector_2.11</artifactId>
<packaging>jar</packaging>
- <version>1.11.0-rc2</version>
+ <version>1.11.0</version>
<name>spark-tensorflow-connector</name>
<url>https://www.tensorflow.org</url>
<description>TensorFlow TFRecord connector for Apache Spark DataFrames</description>
diff --git a/tensorflow/java/maven/tensorflow-hadoop/pom.xml b/tensorflow/java/maven/tensorflow-hadoop/pom.xml
index d07c5fc..a8669ee 100644
--- a/tensorflow/java/maven/tensorflow-hadoop/pom.xml
+++ b/tensorflow/java/maven/tensorflow-hadoop/pom.xml
@@ -5,7 +5,7 @@
<groupId>org.tensorflow</groupId>
<artifactId>tensorflow-hadoop</artifactId>
<packaging>jar</packaging>
- <version>1.11.0-rc2</version>
+ <version>1.11.0</version>
<name>tensorflow-hadoop</name>
<url>https://www.tensorflow.org</url>
<description>TensorFlow TFRecord InputFormat/OutputFormat for Apache Hadoop</description>
diff --git a/tensorflow/java/maven/tensorflow/pom.xml b/tensorflow/java/maven/tensorflow/pom.xml
index af0c68a..67d628b 100644
--- a/tensorflow/java/maven/tensorflow/pom.xml
+++ b/tensorflow/java/maven/tensorflow/pom.xml
@@ -6,7 +6,7 @@
<parent>
<groupId>org.tensorflow</groupId>
<artifactId>parentpom</artifactId>
- <version>1.11.0-rc2</version>
+ <version>1.11.0</version>
<relativePath>../</relativePath>
</parent>
<artifactId>tensorflow</artifactId>
diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD
index 410b3a5..fe81254 100644
--- a/tensorflow/python/BUILD
+++ b/tensorflow/python/BUILD
@@ -1639,6 +1639,15 @@
)
tf_gen_op_wrapper_private_py(
+ name = "experimental_dataset_ops_gen",
+ visibility = [
+ "//learning/brain/python/ops:__pkg__",
+ "//tensorflow:__subpackages__",
+ "//tensorflow/python/kernel_tests:__pkg__",
+ ],
+)
+
+tf_gen_op_wrapper_private_py(
name = "image_ops_gen",
visibility = ["//learning/brain/python/ops:__pkg__"],
)
@@ -1731,6 +1740,14 @@
)
tf_gen_op_wrapper_private_py(
+ name = "stateless_random_ops_gen",
+ visibility = [
+ "//tensorflow/contrib/stateless:__pkg__",
+ "//tensorflow/python/data/experimental/ops:__pkg__",
+ ],
+)
+
+tf_gen_op_wrapper_private_py(
name = "list_ops_gen",
)
@@ -2008,6 +2025,7 @@
":array_ops",
":cond_v2_impl",
":constant_op",
+ ":control_flow_ops",
":control_flow_util",
":framework_ops",
":function_def_to_graph",
@@ -3292,9 +3310,11 @@
"training/checkpointable/**/*.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/saveable_object.py",
"training/saver.py",
+ "training/session_run_hook.py",
"training/training_util.py",
],
),
@@ -3302,6 +3322,7 @@
deps = [
":array_ops",
":array_ops_gen",
+ ":basic_session_run_hooks",
":checkpoint_management",
":checkpoint_ops_gen",
":client",
@@ -3326,6 +3347,7 @@
":saver",
":sdca_ops",
":session",
+ ":session_run_hook",
":sparse_ops",
":sparse_tensor",
":state_ops",
@@ -3370,6 +3392,28 @@
)
py_library(
+ name = "session_run_hook",
+ srcs = ["training/session_run_hook.py"],
+ srcs_version = "PY2AND3",
+ deps = [":util"],
+)
+
+py_library(
+ name = "basic_session_run_hooks",
+ srcs = ["training/basic_session_run_hooks.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ ":client",
+ ":framework",
+ ":platform",
+ ":protos_all_py",
+ ":session_run_hook",
+ ":training_util",
+ ":util",
+ ],
+)
+
+py_library(
name = "saver",
srcs = ["training/saver.py"],
srcs_version = "PY2AND3",
diff --git a/tensorflow/python/autograph/CONTRIBUTING.md b/tensorflow/python/autograph/CONTRIBUTING.md
index 1ded5ba..f3587a4 100644
--- a/tensorflow/python/autograph/CONTRIBUTING.md
+++ b/tensorflow/python/autograph/CONTRIBUTING.md
@@ -9,8 +9,6 @@
does not impact functionality, and AutoGraph will remain accessible under
`tensorflow.contrib.autograph` until `tensorflow.contrib` is retired.
-When
-
## TensorFlow Code of Conduct
Please review and follow the [TensorFlow Code of Conduct](../../CODE_OF_CONDUCT.md).
diff --git a/tensorflow/python/autograph/converters/BUILD b/tensorflow/python/autograph/converters/BUILD
index 7b029de..f06dc78 100644
--- a/tensorflow/python/autograph/converters/BUILD
+++ b/tensorflow/python/autograph/converters/BUILD
@@ -27,10 +27,10 @@
"decorators.py",
"directives.py",
"error_handlers.py",
+ "function_scopes.py",
"list_comprehensions.py",
"lists.py",
"logical_expressions.py",
- "name_scopes.py",
"return_statements.py",
"side_effect_guards.py",
"slices.py",
@@ -157,8 +157,8 @@
)
py_test(
- name = "name_scopes_test",
- srcs = ["name_scopes_test.py"],
+ name = "function_scopes_test",
+ srcs = ["function_scopes_test.py"],
deps = [
":converters",
"//tensorflow/python:client_testlib",
diff --git a/tensorflow/python/autograph/converters/name_scopes.py b/tensorflow/python/autograph/converters/function_scopes.py
similarity index 72%
rename from tensorflow/python/autograph/converters/name_scopes.py
rename to tensorflow/python/autograph/converters/function_scopes.py
index a9c55cc..284b5b3 100644
--- a/tensorflow/python/autograph/converters/name_scopes.py
+++ b/tensorflow/python/autograph/converters/function_scopes.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Wraps a function body with a `name_scope` of the function name."""
+"""Wraps the body of a converted function with auxiliary constructs."""
from __future__ import absolute_import
from __future__ import division
@@ -24,8 +24,8 @@
from tensorflow.python.autograph.pyct import templates
-class FunctionNameScopeTransformer(converter.Base):
- """Wrap a function body with a `name_scope` of the function name."""
+class FunctionBodyTransformer(converter.Base):
+ """Wraps function bodies around autograph-specific boilerplate."""
def _name_for_current_scope(self):
innermost = self.enclosing_entities[-1]
@@ -49,26 +49,28 @@
def visit_FunctionDef(self, node):
node = self.generic_visit(node)
- unscoped_body = []
- scoped_body = node.body
- if scoped_body:
- first = scoped_body[0]
- if isinstance(first, gast.Expr) and isinstance(first.value, gast.Str):
- # Skip any docstring.
- unscoped_body = scoped_body[:1]
- scoped_body = scoped_body[1:]
+ final_body = []
+ indented_body = node.body
+ if node.body:
+ first_statement = node.body[0]
+ # Skip the docstring, if any.
+ if (isinstance(first_statement, gast.Expr) and
+ isinstance(first_statement.value, gast.Str)):
+ indented_body = indented_body[1:]
+ final_body.append(first_statement)
template = """
- with tf.name_scope(scope_name):
+ with ag__.function_scope(scope_name):
body
"""
scoped_body = templates.replace(
template,
scope_name=gast.Str(self._name_for_current_scope()),
- body=scoped_body)
- node.body = unscoped_body + scoped_body
+ body=indented_body)
+ final_body.extend(scoped_body)
+ node.body = final_body
return node
def transform(node, ctx):
- return FunctionNameScopeTransformer(ctx).visit(node)
+ return FunctionBodyTransformer(ctx).visit(node)
diff --git a/tensorflow/python/autograph/converters/name_scopes_test.py b/tensorflow/python/autograph/converters/function_scopes_test.py
similarity index 71%
rename from tensorflow/python/autograph/converters/name_scopes_test.py
rename to tensorflow/python/autograph/converters/function_scopes_test.py
index 73933c1..e5ce03a 100644
--- a/tensorflow/python/autograph/converters/name_scopes_test.py
+++ b/tensorflow/python/autograph/converters/function_scopes_test.py
@@ -12,51 +12,51 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Tests for for_canonicalization module."""
+"""Tests for function_scopes module."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
-from tensorflow.python.autograph.converters import name_scopes
+from tensorflow.python.autograph.converters import function_scopes
from tensorflow.python.autograph.core import converter_testing
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import ops
from tensorflow.python.platform import test
-class FunctionNameScopeTransformer(converter_testing.TestCase):
+class FunctionBodyTransformerTest(converter_testing.TestCase):
def test_basic(self):
def test_fn(l):
- """This should stay here."""
+ """Docstring."""
a = 1
l += a
return l
- with self.converted(test_fn, name_scopes, {}, ops.name_scope) as result:
+ with self.converted(test_fn, function_scopes, {}) as result:
result_op = result.test_fn(constant_op.constant(1))
self.assertIn('test_fn/', result_op.op.name)
- self.assertEqual('This should stay here.', result.test_fn.__doc__)
+ self.assertEqual('Docstring.', result.test_fn.__doc__)
- def test_long_docstring(self):
+ def test_multiline_docstring(self):
- def test_fn(l):
- """Multi-line docstring.
+ tf = None
- Args:
- l: A thing.
- Returns:
- l
+ def test_fn():
+ """First sentence.
+
+ Second sentence.
"""
- return l + 1
+ return tf.constant(1)
- with self.converted(test_fn, name_scopes, {}, ops.name_scope) as result:
- result_op = result.test_fn(constant_op.constant(1))
+ with self.converted(test_fn, function_scopes, {},
+ constant_op.constant) as result:
+ result_op = result.test_fn()
self.assertIn('test_fn/', result_op.op.name)
- self.assertIn('Multi-line docstring.', result.test_fn.__doc__)
- self.assertIn('Returns:', result.test_fn.__doc__)
+ self.assertIn('First sentence.', result.test_fn.__doc__)
+ self.assertIn('Second sentence.', result.test_fn.__doc__)
def test_nested_functions(self):
@@ -68,7 +68,7 @@
l += 1
return l, inner_fn(l)
- with self.converted(test_fn, name_scopes, {}, ops.name_scope) as result:
+ with self.converted(test_fn, function_scopes, {}, ops.name_scope) as result:
first, second = result.test_fn(constant_op.constant(1))
self.assertIn('test_fn/', first.op.name)
self.assertNotIn('inner_fn', first.op.name)
@@ -88,7 +88,7 @@
ns = {'TestClass': TestClass}
node, ctx = self.prepare(TestClass, ns, owner_type=TestClass)
- node = name_scopes.transform(node, ctx)
+ node = function_scopes.transform(node, ctx)
with self.compiled(node, {}, ops.name_scope) as result:
first, second = result.TestClass().test_fn(constant_op.constant(1))
diff --git a/tensorflow/python/autograph/converters/return_statements.py b/tensorflow/python/autograph/converters/return_statements.py
index 62da045..496c99e 100644
--- a/tensorflow/python/autograph/converters/return_statements.py
+++ b/tensorflow/python/autograph/converters/return_statements.py
@@ -212,6 +212,7 @@
def __init__(self):
self.cant_return = False
+ self.function_level = 0
super(DetectReturnInUnsupportedControlFlow, self).__init__()
def visit_While(self, node):
@@ -229,6 +230,12 @@
self.generic_visit(node)
self.cant_return = False
+ def visit_FunctionDef(self, node):
+ if not self.function_level:
+ self.function_level += 1
+ self.generic_visit(node)
+ self.function_level -= 1
+
def visit_Return(self, node):
if self.cant_return:
raise ValueError(
@@ -242,6 +249,7 @@
def __init__(self):
self.cant_return = False
+ self.function_level = 0
super(DetectReturnInConditional, self).__init__()
def visit_If(self, node):
@@ -249,6 +257,12 @@
self.generic_visit(node)
self.cant_return = False
+ def visit_FunctionDef(self, node):
+ if not self.function_level:
+ self.function_level += 1
+ self.generic_visit(node)
+ self.function_level -= 1
+
def visit_Return(self, node):
if self.cant_return:
raise ValueError(
diff --git a/tensorflow/python/autograph/converters/return_statements_test.py b/tensorflow/python/autograph/converters/return_statements_test.py
index 01dd03d..762fbc6 100644
--- a/tensorflow/python/autograph/converters/return_statements_test.py
+++ b/tensorflow/python/autograph/converters/return_statements_test.py
@@ -151,6 +151,18 @@
self.assertTransformedEquivalent(test_fn, 2)
self.assertTransformedEquivalent(test_fn, -2)
+ def test_nested_functions_in_control_flow(self):
+
+ def test_fn(x):
+
+ if x:
+ def inner_fn(y):
+ return y
+ inner_fn(x)
+
+ self.assertTransformedEquivalent(test_fn, 2)
+ self.assertTransformedEquivalent(test_fn, -2)
+
def test_loop(self):
def test_fn(x):
diff --git a/tensorflow/python/autograph/core/BUILD b/tensorflow/python/autograph/core/BUILD
index 85fecf0..843e381 100644
--- a/tensorflow/python/autograph/core/BUILD
+++ b/tensorflow/python/autograph/core/BUILD
@@ -20,11 +20,13 @@
"config.py",
"converter.py",
"errors.py",
+ "function_wrapping.py",
"naming.py",
],
srcs_version = "PY2AND3",
visibility = ["//tensorflow:__subpackages__"],
deps = [
+ "//tensorflow/python:framework_ops",
"//tensorflow/python/autograph/pyct",
"//tensorflow/python/autograph/pyct/static_analysis",
"//tensorflow/python/autograph/utils",
@@ -47,6 +49,16 @@
)
py_test(
+ name = "function_wrapping_test",
+ srcs = ["function_wrapping_test.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ ":core",
+ "//tensorflow/python:client_testlib",
+ ],
+)
+
+py_test(
name = "naming_test",
srcs = ["naming_test.py"],
srcs_version = "PY2AND3",
diff --git a/tensorflow/python/autograph/core/converter_testing.py b/tensorflow/python/autograph/core/converter_testing.py
index 7ce1b7c..dc2d419 100644
--- a/tensorflow/python/autograph/core/converter_testing.py
+++ b/tensorflow/python/autograph/core/converter_testing.py
@@ -29,6 +29,7 @@
from tensorflow.python.autograph.core import config
from tensorflow.python.autograph.core import converter
from tensorflow.python.autograph.core import errors
+from tensorflow.python.autograph.core import function_wrapping
from tensorflow.python.autograph.pyct import compiler
from tensorflow.python.autograph.pyct import parser
from tensorflow.python.autograph.pyct import pretty_printer
@@ -112,6 +113,7 @@
fake_ag.__dict__['utils'] = utils
fake_ag.__dict__['rewrite_graph_construction_error'] = (
errors.rewrite_graph_construction_error)
+ fake_ag.__dict__['function_scope'] = function_wrapping.function_scope
result.__dict__['ag__'] = fake_ag
for k, v in namespace.items():
result.__dict__[k] = v
diff --git a/tensorflow/contrib/data/python/ops/contrib_op_loader.py b/tensorflow/python/autograph/core/function_wrapping.py
similarity index 65%
copy from tensorflow/contrib/data/python/ops/contrib_op_loader.py
copy to tensorflow/python/autograph/core/function_wrapping.py
index 8f495a9..21b66ef 100644
--- a/tensorflow/contrib/data/python/ops/contrib_op_loader.py
+++ b/tensorflow/python/autograph/core/function_wrapping.py
@@ -1,4 +1,4 @@
-# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+# 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.
@@ -12,13 +12,19 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Python helper for loading contrib ops and kernels."""
+"""Support for wrapping converted functions bodies with auxiliary logic."""
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
-from tensorflow.contrib.util import loader
-from tensorflow.python.platform import resource_loader
+import contextlib
-_dataset_ops = loader.load_op_library(
- resource_loader.get_path_to_datafile("../../_dataset_ops.so"))
+from tensorflow.python.framework import ops
+
+
+@contextlib.contextmanager
+def function_scope(function_name):
+ """Returns a context manager for the converted body of a function."""
+ with ops.name_scope(function_name):
+ yield
diff --git a/tensorflow/contrib/data/python/ops/contrib_op_loader.py b/tensorflow/python/autograph/core/function_wrapping_test.py
similarity index 63%
copy from tensorflow/contrib/data/python/ops/contrib_op_loader.py
copy to tensorflow/python/autograph/core/function_wrapping_test.py
index 8f495a9..5e21705 100644
--- a/tensorflow/contrib/data/python/ops/contrib_op_loader.py
+++ b/tensorflow/python/autograph/core/function_wrapping_test.py
@@ -12,13 +12,23 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Python helper for loading contrib ops and kernels."""
+"""Tests for function_wrapping module."""
+
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
-from tensorflow.contrib.util import loader
-from tensorflow.python.platform import resource_loader
+from tensorflow.python.autograph.core import function_wrapping
+from tensorflow.python.framework import constant_op
+from tensorflow.python.platform import test
-_dataset_ops = loader.load_op_library(
- resource_loader.get_path_to_datafile("../../_dataset_ops.so"))
+
+class FunctionWrappingTest(test.TestCase):
+
+ def test_function_scope_name(self):
+ with function_wrapping.function_scope('test_name'):
+ t = constant_op.constant(1)
+ self.assertIn('test_name', t.name)
+
+if __name__ == '__main__':
+ test.main()
diff --git a/tensorflow/python/autograph/impl/conversion.py b/tensorflow/python/autograph/impl/conversion.py
index a0d13c8..52abd40 100644
--- a/tensorflow/python/autograph/impl/conversion.py
+++ b/tensorflow/python/autograph/impl/conversion.py
@@ -34,15 +34,16 @@
from tensorflow.python.autograph.converters import decorators
from tensorflow.python.autograph.converters import directives
from tensorflow.python.autograph.converters import error_handlers
+from tensorflow.python.autograph.converters import function_scopes
from tensorflow.python.autograph.converters import lists
from tensorflow.python.autograph.converters import logical_expressions
-from tensorflow.python.autograph.converters import name_scopes
from tensorflow.python.autograph.converters import return_statements
from tensorflow.python.autograph.converters import side_effect_guards
from tensorflow.python.autograph.converters import slices
from tensorflow.python.autograph.core import config
from tensorflow.python.autograph.core import converter
from tensorflow.python.autograph.core import errors
+from tensorflow.python.autograph.core import function_wrapping
from tensorflow.python.autograph.pyct import ast_util
from tensorflow.python.autograph.pyct import inspect_utils
from tensorflow.python.autograph.pyct import origin_info
@@ -257,6 +258,7 @@
ag_internal.converted_call = autograph_module.converted_call
ag_internal.ConversionOptions = autograph_module.ConversionOptions
ag_internal.utils = utils
+ ag_internal.function_scope = function_wrapping.function_scope
ag_internal.rewrite_graph_construction_error = (
errors.rewrite_graph_construction_error)
# TODO(mdan): Add safeguards against name clashes.
@@ -346,7 +348,7 @@
node = converter.apply_(node, context, conditional_expressions)
node = converter.apply_(node, context, logical_expressions)
node = converter.apply_(node, context, side_effect_guards)
- node = converter.apply_(node, context, name_scopes)
+ node = converter.apply_(node, context, function_scopes)
if rewrite_errors:
node = converter.apply_(node, context, error_handlers)
return node
diff --git a/tensorflow/python/autograph/pyct/templates.py b/tensorflow/python/autograph/pyct/templates.py
index 1bf0515..1af8fca 100644
--- a/tensorflow/python/autograph/pyct/templates.py
+++ b/tensorflow/python/autograph/pyct/templates.py
@@ -123,6 +123,8 @@
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)
diff --git a/tensorflow/python/autograph/pyct/templates_test.py b/tensorflow/python/autograph/pyct/templates_test.py
index 078d9a1..3032241 100644
--- a/tensorflow/python/autograph/pyct/templates_test.py
+++ b/tensorflow/python/autograph/pyct/templates_test.py
@@ -158,6 +158,18 @@
self.assertIsInstance(function_call_arg.elts[0].elts[0].ctx, gast.Load)
self.assertIsInstance(function_call_arg.elts[0].elts[1].ctx, gast.Load)
+ def test_replace_index(self):
+ template = """
+ def test_fn(foo):
+ foo = 0
+ """
+
+ node = templates.replace(
+ template, foo=parser.parse_expression('foo(a[b]).bar'))[0]
+ function_call_arg = node.body[0].targets[0].value.args[0]
+ self.assertIsInstance(function_call_arg.ctx, gast.Load)
+ self.assertIsInstance(function_call_arg.slice.value.ctx, gast.Load)
+
def test_replace_call_keyword(self):
template = """
def test_fn():
diff --git a/tensorflow/python/client/session_test.py b/tensorflow/python/client/session_test.py
index f576435..347833c 100644
--- a/tensorflow/python/client/session_test.py
+++ b/tensorflow/python/client/session_test.py
@@ -120,11 +120,17 @@
inp = constant_op.constant(10.0, name='W1')
self.assertAllEqual(inp.eval(), 10.0)
- devices = sess.list_devices()
- self.assertEqual(2, len(devices))
- for device in devices:
- self.assertEqual('CPU', framework_device_lib.DeviceSpec.from_string(
- device.name).device_type)
+ num_cpu_devices = 0
+ num_gpu_devices = 0
+ for device in sess.list_devices():
+ device_type = framework_device_lib.DeviceSpec.from_string(
+ device.name).device_type
+ if device_type == 'CPU':
+ num_cpu_devices += 1
+ elif device_type == 'GPU':
+ num_gpu_devices += 1
+ self.assertEqual(2, num_cpu_devices)
+ self.assertEqual(0, num_gpu_devices)
def testPerSessionThreads(self):
with session.Session(
@@ -1022,7 +1028,7 @@
with session.Session():
a = constant_op.constant(1.0, shape=[1, 2])
b = constant_op.constant(2.0, shape=[1, 2], name='b')
- v = variables.Variable(a, a.dtype)
+ v = variables.VariableV1(a, a.dtype)
assign_a_to_v = state_ops.assign(v, a)
assign_a_to_v.eval()
diff --git a/tensorflow/python/compat/compat.py b/tensorflow/python/compat/compat.py
index ce23073..bea5aa9 100644
--- a/tensorflow/python/compat/compat.py
+++ b/tensorflow/python/compat/compat.py
@@ -26,7 +26,7 @@
from tensorflow.python.util import tf_contextlib
from tensorflow.python.util.tf_export import tf_export
-_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2018, 9, 26)
+_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2018, 10, 1)
@tf_export("compat.forward_compatible")
diff --git a/tensorflow/python/data/BUILD b/tensorflow/python/data/BUILD
index 138141f..e32eeec 100644
--- a/tensorflow/python/data/BUILD
+++ b/tensorflow/python/data/BUILD
@@ -10,6 +10,7 @@
srcs_version = "PY2AND3",
deps = [
"//tensorflow/python:util",
+ "//tensorflow/python/data/experimental",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/ops:iterator_ops",
"//tensorflow/python/data/ops:multi_device_iterator_ops",
diff --git a/tensorflow/python/data/__init__.py b/tensorflow/python/data/__init__.py
index f8b5612..7536ba6 100644
--- a/tensorflow/python/data/__init__.py
+++ b/tensorflow/python/data/__init__.py
@@ -22,6 +22,7 @@
from __future__ import print_function
# pylint: disable=unused-import
+from tensorflow.python.data import experimental
from tensorflow.python.data.ops.dataset_ops import Dataset
from tensorflow.python.data.ops.iterator_ops import Iterator
from tensorflow.python.data.ops.readers import FixedLengthRecordDataset
diff --git a/tensorflow/python/data/experimental/BUILD b/tensorflow/python/data/experimental/BUILD
new file mode 100644
index 0000000..84e761d
--- /dev/null
+++ b/tensorflow/python/data/experimental/BUILD
@@ -0,0 +1,16 @@
+package(default_visibility = ["//tensorflow:internal"])
+
+licenses(["notice"]) # Apache 2.0
+
+exports_files(["LICENSE"])
+
+py_library(
+ name = "experimental",
+ srcs = ["__init__.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ "//tensorflow/python:util",
+ "//tensorflow/python/data/experimental/ops:dataset_ops",
+ "//tensorflow/python/data/experimental/ops:iterator_ops",
+ ],
+)
diff --git a/tensorflow/python/data/experimental/__init__.py b/tensorflow/python/data/experimental/__init__.py
new file mode 100644
index 0000000..2ac159d
--- /dev/null
+++ b/tensorflow/python/data/experimental/__init__.py
@@ -0,0 +1,109 @@
+# 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.
+# ==============================================================================
+"""Experimental API for building input pipelines.
+
+This module contains experimental `Dataset` sources and transformations that can
+be used in conjunction with the `tf.data.Dataset` API. Note that the
+`tf.data.experimental` API is not subject to the same backwards compatibility
+guarantees as `tf.data`, but we will provide deprecation advice in advance of
+removing existing functionality.
+
+See [Importing Data](https://tensorflow.org/guide/datasets) for an overview.
+
+@@Counter
+@@CheckpointInputPipelineHook
+@@CsvDataset
+@@Optional
+@@RandomDataset
+@@Reducer
+@@SqlDataset
+@@TFRecordWriter
+
+@@bucket_by_sequence_length
+@@choose_from_datasets
+@@copy_to_device
+@@dense_to_sparse_batch
+@@enumerate_dataset
+@@get_next_as_optional
+@@get_single_element
+@@group_by_reducer
+@@group_by_window
+@@ignore_errors
+@@latency_stats
+@@make_batched_features_dataset
+@@make_csv_dataset
+@@make_saveable_from_iterator
+@@map_and_batch
+@@parallel_interleave
+@@parse_example_dataset
+@@prefetch_to_device
+@@rejection_resample
+@@sample_from_datasets
+@@scan
+@@set_stats_aggregator
+@@shuffle_and_repeat
+@@StatsAggregator
+@@unbatch
+@@unique
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+# pylint: disable=unused-import
+
+from tensorflow.python.data.experimental.ops.batching import dense_to_sparse_batch
+from tensorflow.python.data.experimental.ops.batching import map_and_batch
+from tensorflow.python.data.experimental.ops.batching import unbatch
+from tensorflow.python.data.experimental.ops.counter import Counter
+from tensorflow.python.data.experimental.ops.enumerate_ops import enumerate_dataset
+from tensorflow.python.data.experimental.ops.error_ops import ignore_errors
+from tensorflow.python.data.experimental.ops.get_single_element import get_single_element
+from tensorflow.python.data.experimental.ops.grouping import bucket_by_sequence_length
+from tensorflow.python.data.experimental.ops.grouping import group_by_reducer
+from tensorflow.python.data.experimental.ops.grouping import group_by_window
+from tensorflow.python.data.experimental.ops.grouping import Reducer
+from tensorflow.python.data.experimental.ops.interleave_ops import choose_from_datasets
+from tensorflow.python.data.experimental.ops.interleave_ops import parallel_interleave
+from tensorflow.python.data.experimental.ops.interleave_ops import sample_from_datasets
+from tensorflow.python.data.experimental.ops.iterator_ops import CheckpointInputPipelineHook
+from tensorflow.python.data.experimental.ops.iterator_ops import make_saveable_from_iterator
+
+# Optimization constant that can be used to enable auto-tuning.
+from tensorflow.python.data.experimental.ops.optimization import AUTOTUNE
+
+from tensorflow.python.data.experimental.ops.parsing_ops import parse_example_dataset
+from tensorflow.python.data.experimental.ops.prefetching_ops import copy_to_device
+from tensorflow.python.data.experimental.ops.prefetching_ops import prefetch_to_device
+from tensorflow.python.data.experimental.ops.random_ops import RandomDataset
+from tensorflow.python.data.experimental.ops.readers import CsvDataset
+from tensorflow.python.data.experimental.ops.readers import make_batched_features_dataset
+from tensorflow.python.data.experimental.ops.readers import make_csv_dataset
+from tensorflow.python.data.experimental.ops.readers import SqlDataset
+from tensorflow.python.data.experimental.ops.resampling import rejection_resample
+from tensorflow.python.data.experimental.ops.scan_ops import scan
+from tensorflow.python.data.experimental.ops.shuffle_ops import shuffle_and_repeat
+from tensorflow.python.data.experimental.ops.stats_ops import latency_stats
+from tensorflow.python.data.experimental.ops.stats_ops import set_stats_aggregator
+from tensorflow.python.data.experimental.ops.stats_ops import StatsAggregator
+from tensorflow.python.data.experimental.ops.unique import unique
+from tensorflow.python.data.experimental.ops.writers import TFRecordWriter
+from tensorflow.python.data.ops.iterator_ops import get_next_as_optional
+from tensorflow.python.data.ops.optional_ops import Optional
+# pylint: enable=unused-import
+
+from tensorflow.python.util.all_util import remove_undocumented
+remove_undocumented(__name__)
diff --git a/tensorflow/python/data/experimental/kernel_tests/BUILD b/tensorflow/python/data/experimental/kernel_tests/BUILD
new file mode 100644
index 0000000..a46c30e
--- /dev/null
+++ b/tensorflow/python/data/experimental/kernel_tests/BUILD
@@ -0,0 +1,569 @@
+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 = "batch_dataset_op_test",
+ size = "medium",
+ srcs = ["batch_dataset_op_test.py"],
+ srcs_version = "PY2AND3",
+ tags = [
+ "no_oss", # (b/79552534)
+ "no_pip",
+ ],
+ deps = [
+ "//tensorflow/python:array_ops",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:constant_op",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:errors",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python:math_ops",
+ "//tensorflow/python:script_ops",
+ "//tensorflow/python:session",
+ "//tensorflow/python:sparse_tensor",
+ "//tensorflow/python:string_ops",
+ "//tensorflow/python:tensor_shape",
+ "//tensorflow/python:util",
+ "//tensorflow/python/data/experimental/ops:batching",
+ "//tensorflow/python/data/kernel_tests:test_base",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//third_party/py/numpy",
+ "@absl_py//absl/testing:parameterized",
+ ],
+)
+
+py_test(
+ name = "bucketing_test",
+ size = "medium",
+ srcs = ["bucketing_test.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ "//tensorflow/python:array_ops",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:constant_op",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:errors",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python:math_ops",
+ "//tensorflow/python:sparse_tensor",
+ "//tensorflow/python:string_ops",
+ "//tensorflow/python:tensor_shape",
+ "//tensorflow/python/data/experimental/ops:grouping",
+ "//tensorflow/python/data/kernel_tests:test_base",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//third_party/py/numpy",
+ ],
+)
+
+py_test(
+ name = "csv_dataset_op_test",
+ size = "medium",
+ srcs = ["csv_dataset_op_test.py"],
+ srcs_version = "PY2AND3",
+ tags = ["no_pip"],
+ deps = [
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:constant_op",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:errors",
+ "//tensorflow/python:framework_test_lib",
+ "//tensorflow/python:parsing_ops",
+ "//tensorflow/python:platform",
+ "//tensorflow/python:platform_test",
+ "//tensorflow/python:session",
+ "//tensorflow/python/data/experimental/ops:error_ops",
+ "//tensorflow/python/data/experimental/ops:readers",
+ "//tensorflow/python/data/kernel_tests:test_base",
+ "//tensorflow/python/data/ops:readers",
+ "//tensorflow/python/eager:context",
+ "//third_party/py/numpy",
+ ],
+)
+
+py_test(
+ name = "dataset_constructor_op_test",
+ size = "medium",
+ srcs = ["dataset_constructor_op_test.py"],
+ srcs_version = "PY2AND3",
+ tags = [
+ "manual",
+ "nomac", # b/62040583
+ ],
+ deps = [
+ "//tensorflow/python:array_ops",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python/data/experimental/ops:batching",
+ "//tensorflow/python/data/kernel_tests:test_base",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/util:nest",
+ ],
+)
+
+py_test(
+ name = "directed_interleave_dataset_test",
+ size = "medium",
+ srcs = ["directed_interleave_dataset_test.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:errors",
+ "//tensorflow/python:random_seed",
+ "//tensorflow/python/data/experimental/ops:interleave_ops",
+ "//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"],
+ deps = [
+ "//tensorflow/python:array_ops",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:constant_op",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:errors",
+ "//tensorflow/python:sparse_tensor",
+ "//tensorflow/python/data/experimental/ops:get_single_element",
+ "//tensorflow/python/data/experimental/ops:grouping",
+ "//tensorflow/python/data/kernel_tests:test_base",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "@absl_py//absl/testing:parameterized",
+ ],
+)
+
+py_test(
+ name = "indexed_dataset_ops_test",
+ srcs = ["indexed_dataset_ops_test.py"],
+ deps = [
+ "//tensorflow/python:array_ops",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:experimental_dataset_ops_gen",
+ "//tensorflow/python/data/experimental/ops:indexed_dataset_ops",
+ "//tensorflow/python/data/kernel_tests:test_base",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//third_party/py/numpy",
+ ],
+)
+
+py_test(
+ name = "interleave_dataset_op_test",
+ size = "medium",
+ srcs = ["interleave_dataset_op_test.py"],
+ srcs_version = "PY2AND3",
+ tags = [
+ "no_oss",
+ "no_pip",
+ "notap",
+ ],
+ deps = [
+ "//tensorflow/python:array_ops",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:errors",
+ "//tensorflow/python:math_ops",
+ "//tensorflow/python:script_ops",
+ "//tensorflow/python:sparse_ops",
+ "//tensorflow/python:sparse_tensor",
+ "//tensorflow/python/data/experimental/ops:interleave_ops",
+ "//tensorflow/python/data/kernel_tests:test_base",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "@six_archive//:six",
+ ],
+)
+
+py_test(
+ name = "iterator_ops_test",
+ size = "small",
+ srcs = ["iterator_ops_test.py"],
+ srcs_version = "PY2AND3",
+ tags = ["no_pip"],
+ deps = [
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:constant_op",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python:training",
+ "//tensorflow/python:variables",
+ "//tensorflow/python/data/experimental/ops:iterator_ops",
+ "//tensorflow/python/data/kernel_tests:test_base",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/estimator:estimator_py",
+ ],
+)
+
+py_test(
+ name = "map_dataset_op_test",
+ size = "medium",
+ srcs = ["map_dataset_op_test.py"],
+ srcs_version = "PY2AND3",
+ tags = [
+ "no_pip",
+ "noasan", # times out
+ "optonly",
+ ],
+ 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:batching",
+ "//tensorflow/python/data/experimental/ops:error_ops",
+ "//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 = "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 = "map_defun_op_test",
+ size = "small",
+ srcs = ["map_defun_op_test.py"],
+ srcs_version = "PY2AND3",
+ tags = ["no_pip"],
+ deps = [
+ "//tensorflow/python:array_ops",
+ "//tensorflow/python:check_ops",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:constant_op",
+ "//tensorflow/python:data_flow_ops",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python:function",
+ "//tensorflow/python:functional_ops",
+ "//tensorflow/python:math_ops",
+ "//tensorflow/python:session",
+ "//tensorflow/python/data/experimental/ops:map_defun",
+ "//tensorflow/python/data/kernel_tests:test_base",
+ ],
+)
+
+py_test(
+ name = "parsing_ops_test",
+ size = "small",
+ srcs = ["parsing_ops_test.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ "//tensorflow/core:protos_all_py",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:errors",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python:parsing_ops",
+ "//tensorflow/python:platform",
+ "//tensorflow/python:sparse_tensor",
+ "//tensorflow/python/data/experimental/ops:parsing_ops",
+ "//tensorflow/python/data/kernel_tests:test_base",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/util:nest",
+ "//third_party/py/numpy",
+ ],
+)
+
+cuda_py_test(
+ name = "prefetching_ops_test",
+ size = "small",
+ srcs = ["prefetching_ops_test.py"],
+ additional_deps = [
+ "//tensorflow/python/data/experimental/ops:prefetching_ops",
+ "//tensorflow/core:protos_all_py",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:constant_op",
+ "//tensorflow/python/data/kernel_tests:test_base",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python:framework_test_lib",
+ "//tensorflow/python:function",
+ "//tensorflow/python:resource_variable_ops",
+ "//tensorflow/python/compat:compat",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/ops:iterator_ops",
+ ],
+ tags = ["no_windows_gpu"],
+)
+
+py_test(
+ name = "range_dataset_op_test",
+ size = "small",
+ srcs = ["range_dataset_op_test.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:constant_op",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:errors",
+ "//tensorflow/python:tensor_shape",
+ "//tensorflow/python/data/experimental/ops:counter",
+ "//tensorflow/python/data/experimental/ops:enumerate_ops",
+ "//tensorflow/python/data/kernel_tests:test_base",
+ "//tensorflow/python/data/ops:dataset_ops",
+ ],
+)
+
+py_library(
+ name = "reader_dataset_ops_test_base",
+ testonly = 1,
+ srcs = [
+ "reader_dataset_ops_test_base.py",
+ ],
+ srcs_version = "PY2AND3",
+ visibility = [
+ "//tensorflow/python/data/experimental/kernel_tests:__pkg__",
+ "//tensorflow/python/data/experimental/kernel_tests/serialization:__pkg__",
+ ],
+ deps = [
+ "//tensorflow/core:protos_all_py",
+ "//tensorflow/python:array_ops",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:constant_op",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:lib",
+ "//tensorflow/python:parsing_ops",
+ "//tensorflow/python:util",
+ "//tensorflow/python/data/experimental/ops:readers",
+ "//tensorflow/python/data/kernel_tests:test_base",
+ "//tensorflow/python/data/ops:iterator_ops",
+ "//tensorflow/python/data/ops:readers",
+ ],
+)
+
+py_test(
+ name = "reader_dataset_ops_test",
+ size = "medium",
+ srcs = ["reader_dataset_ops_test.py"],
+ srcs_version = "PY2AND3",
+ tags = ["no_pip"],
+ deps = [
+ ":reader_dataset_ops_test_base",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:constant_op",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:errors",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python:parsing_ops",
+ "//tensorflow/python:string_ops",
+ "//tensorflow/python/data/experimental/ops:readers",
+ "//tensorflow/python/data/ops:readers",
+ "//tensorflow/python/data/util:nest",
+ "//third_party/py/numpy",
+ ],
+)
+
+py_test(
+ name = "resample_test",
+ size = "medium",
+ srcs = ["resample_test.py"],
+ shard_count = 2,
+ srcs_version = "PY2AND3",
+ tags = [
+ "noasan",
+ "optonly",
+ ],
+ deps = [
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:errors",
+ "//tensorflow/python:math_ops",
+ "//tensorflow/python:random_ops",
+ "//tensorflow/python:string_ops",
+ "//tensorflow/python:util",
+ "//tensorflow/python/data/experimental/ops:resampling",
+ "//tensorflow/python/data/kernel_tests:test_base",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//third_party/py/numpy",
+ "@absl_py//absl/testing:parameterized",
+ "@six_archive//:six",
+ ],
+)
+
+py_test(
+ name = "scan_dataset_op_test",
+ size = "small",
+ srcs = ["scan_dataset_op_test.py"],
+ srcs_version = "PY2AND3",
+ tags = ["no_pip"],
+ deps = [
+ "//tensorflow/python:array_ops",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:constant_op",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:errors",
+ "//tensorflow/python:framework_test_lib",
+ "//tensorflow/python:sparse_tensor",
+ "//tensorflow/python/data/experimental/ops:scan_ops",
+ "//tensorflow/python/data/kernel_tests:test_base",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/eager:context",
+ "//third_party/py/numpy",
+ ],
+)
+
+py_test(
+ name = "shuffle_dataset_op_test",
+ size = "medium",
+ srcs = ["shuffle_dataset_op_test.py"],
+ srcs_version = "PY2AND3",
+ tags = [
+ "no_pip",
+ "optonly",
+ ],
+ deps = [
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:errors",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python/data/experimental/ops:shuffle_ops",
+ "//tensorflow/python/data/kernel_tests:test_base",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//third_party/py/numpy",
+ ],
+)
+
+py_library(
+ name = "sql_dataset_op_test_base",
+ srcs = ["sql_dataset_op_test_base.py"],
+ srcs_version = "PY2AND3",
+ visibility = [
+ "//tensorflow/python/data/experimental/kernel_tests:__pkg__",
+ "//tensorflow/python/data/experimental/kernel_tests/serialization:__pkg__",
+ ],
+ deps = [
+ "//tensorflow/python:array_ops",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python/data/experimental/ops:readers",
+ "//tensorflow/python/data/kernel_tests:test_base",
+ "@org_sqlite//:python",
+ ],
+)
+
+py_test(
+ name = "sql_dataset_op_test",
+ size = "small",
+ srcs = ["sql_dataset_op_test.py"],
+ srcs_version = "PY2AND3",
+ tags = ["no_pip"],
+ deps = [
+ ":sql_dataset_op_test_base",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:errors",
+ ],
+)
+
+py_test(
+ name = "stats_dataset_ops_test",
+ size = "medium",
+ srcs = ["stats_dataset_ops_test.py"],
+ srcs_version = "PY2AND3",
+ tags = ["no_pip"],
+ deps = [
+ ":reader_dataset_ops_test_base",
+ ":stats_dataset_test_base",
+ "//tensorflow/python:array_ops",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:errors",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python/data/experimental/ops:stats_ops",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//third_party/py/numpy",
+ ],
+)
+
+py_library(
+ name = "stats_dataset_test_base",
+ srcs = ["stats_dataset_test_base.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ "//tensorflow/core:protos_all_py",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python/data/kernel_tests:test_base",
+ ],
+)
+
+py_test(
+ name = "threadpool_dataset_ops_test",
+ size = "small",
+ srcs = ["threadpool_dataset_ops_test.py"],
+ srcs_version = "PY2AND3",
+ tags = ["no_pip"],
+ deps = [
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:errors",
+ "//tensorflow/python:script_ops",
+ "//tensorflow/python/data/experimental/ops:threadpool",
+ "//tensorflow/python/data/experimental/ops:unique",
+ "//tensorflow/python/data/kernel_tests:test_base",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//third_party/py/numpy",
+ "@absl_py//absl/testing:parameterized",
+ ],
+)
+
+py_test(
+ name = "unique_dataset_op_test",
+ size = "small",
+ srcs = ["unique_dataset_op_test.py"],
+ srcs_version = "PY2AND3",
+ tags = ["no_pip"],
+ deps = [
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:errors",
+ "//tensorflow/python:util",
+ "//tensorflow/python/data/experimental/ops:unique",
+ "//tensorflow/python/data/kernel_tests:test_base",
+ "//tensorflow/python/data/ops:dataset_ops",
+ ],
+)
+
+py_test(
+ name = "writer_ops_test",
+ size = "small",
+ srcs = ["writer_ops_test.py"],
+ deps = [
+ "//tensorflow/python:array_ops",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:lib",
+ "//tensorflow/python:util",
+ "//tensorflow/python/data/experimental/ops:writers",
+ "//tensorflow/python/data/kernel_tests:test_base",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/ops:readers",
+ ],
+)
diff --git a/tensorflow/contrib/data/python/kernel_tests/batch_dataset_op_test.py b/tensorflow/python/data/experimental/kernel_tests/batch_dataset_op_test.py
similarity index 66%
rename from tensorflow/contrib/data/python/kernel_tests/batch_dataset_op_test.py
rename to tensorflow/python/data/experimental/kernel_tests/batch_dataset_op_test.py
index e2508de..8703b28 100644
--- a/tensorflow/contrib/data/python/kernel_tests/batch_dataset_op_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/batch_dataset_op_test.py
@@ -23,15 +23,15 @@
from absl.testing import parameterized
import numpy as np
-from tensorflow.contrib.data.python.ops import batching
from tensorflow.python.client import session
+from tensorflow.python.data.experimental.ops import batching
+from tensorflow.python.data.kernel_tests import test_base
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 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 math_ops
from tensorflow.python.ops import script_ops
@@ -40,12 +40,7 @@
from tensorflow.python.util import compat
-class BatchDatasetTest(test.TestCase, parameterized.TestCase):
-
- def assertSparseValuesEqual(self, a, b):
- self.assertAllEqual(a.indices, b.indices)
- self.assertAllEqual(a.values, b.values)
- self.assertAllEqual(a.dense_shape, b.dense_shape)
+class BatchDatasetTest(test_base.DatasetTestBase, parameterized.TestCase):
def testDenseToSparseBatchDataset(self):
components = np.random.randint(12, size=(100,)).astype(np.int32)
@@ -305,128 +300,6 @@
with self.assertRaises(errors.InvalidArgumentError):
sess.run(next_element)
- def testBatchAndDropRemainder(self):
- components = (np.arange(7),
- np.array([[1, 2, 3]]) * np.arange(7)[:, np.newaxis],
- np.array(37.0) * np.arange(7))
-
- batch_size = array_ops.placeholder(dtypes.int64, shape=[])
-
- iterator = (
- dataset_ops.Dataset.from_tensor_slices(components).apply(
- batching.batch_and_drop_remainder(batch_size))
- .make_initializable_iterator())
-
- next_element = iterator.get_next()
-
- with self.cached_session() as sess:
- for test_batch_size in [1, 3, 7, 10]:
- sess.run(iterator.initializer, feed_dict={batch_size: test_batch_size})
- num_batches = 7 // test_batch_size
- for i in range(num_batches):
- result = sess.run(next_element)
- for component, result_component in zip(components, result):
- for j in range(test_batch_size):
- self.assertAllEqual(component[(i * test_batch_size + j)],
- result_component[j])
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(next_element)
-
- def testBatchAndDropRemainderSparse(self):
-
- def _sparse(i):
- return sparse_tensor.SparseTensorValue(
- indices=[[0]], values=(i * [1]), dense_shape=[1])
-
- iterator = dataset_ops.Dataset.range(12).map(_sparse).apply(
- batching.batch_and_drop_remainder(5)).make_initializable_iterator()
- init_op = iterator.initializer
- get_next = iterator.get_next()
-
- with self.cached_session() as sess:
- sess.run(init_op)
- for i in range(2):
- actual = sess.run(get_next)
- expected = sparse_tensor.SparseTensorValue(
- indices=[[0, 0], [1, 0], [2, 0], [3, 0], [4, 0]],
- values=[i * 5, i * 5 + 1, i * 5 + 2, i * 5 + 3, i * 5 + 4],
- dense_shape=[5, 1])
- self.assertTrue(sparse_tensor.is_sparse(actual))
- self.assertSparseValuesEqual(actual, expected)
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
-
- def testPaddedBatchAndDropRemainder(self):
- els = []
- for length in [3, 6, 9, 4, 12, 10, 2]:
- els.append((np.array(length), np.arange(length) + 1,
- np.array(length * 2)))
-
- dataset = dataset_ops.Dataset.from_tensors(els[0])
- for el in els[1:]:
- dataset = dataset.concatenate(dataset_ops.Dataset.from_tensors(el))
-
- batch_size = array_ops.placeholder(dtypes.int64, shape=[])
- iterator = (
- dataset.apply(
- batching.padded_batch_and_drop_remainder(
- batch_size, ([], [None], []))).make_initializable_iterator())
-
- next_element = iterator.get_next()
-
- with self.cached_session() as sess:
- for test_batch_size in [1, 3, 7, 10]:
- sess.run(iterator.initializer, feed_dict={batch_size: test_batch_size})
- num_batches = 7 // test_batch_size
- for i in range(num_batches):
- result = sess.run(next_element)
- for component_idx, result_component in enumerate(result):
- for j in range(test_batch_size):
- data_idx = i * test_batch_size + j
- comp = result_component[j]
- unpadded = comp[comp > 0]
- if np.isscalar(comp):
- # The boolean mask indexing above adds a dim back. Rm it.
- unpadded = unpadded[0]
- self.assertAllEqual(els[data_idx][component_idx], unpadded)
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(next_element)
-
- def testPaddedBatchAndDropRemainderSparseError(self):
-
- def _map_fn(i):
- return sparse_tensor.SparseTensorValue(
- indices=[[0, 0]], values=(i * [1]), dense_shape=[1, 1]), i
-
- with self.assertRaises(TypeError):
- _ = dataset_ops.Dataset.range(10).map(_map_fn).apply(
- batching.padded_batch_and_drop_remainder(5))
-
- def testBatchAndDropRemainderShapeInference(self):
- components = (array_ops.placeholder(dtypes.int32),
- (array_ops.placeholder(dtypes.int32, shape=[None]),
- array_ops.placeholder(dtypes.int32, shape=[20, 30])))
-
- # Test with a statically known batch size.
- dataset = (
- dataset_ops.Dataset.from_tensor_slices(components).apply(
- batching.batch_and_drop_remainder(128)))
-
- self.assertIs(None, dataset.output_shapes[0].ndims)
- self.assertEqual([128], dataset.output_shapes[1][0].as_list())
- self.assertEqual([128, 30], dataset.output_shapes[1][1].as_list())
-
- # Test with a dynamic batch size: the static shape will be unknown, because
- # `batch_size` is a placeholder.
- batch_size = array_ops.placeholder(dtypes.int64)
- dataset = (
- dataset_ops.Dataset.from_tensor_slices(components).apply(
- batching.batch_and_drop_remainder(batch_size)))
-
- self.assertIs(None, dataset.output_shapes[0].ndims)
- self.assertEqual([None], dataset.output_shapes[1][0].as_list())
- self.assertEqual([None, 30], dataset.output_shapes[1][1].as_list())
-
@parameterized.named_parameters(
("Default", None, None),
("SequentialCalls", 1, None),
@@ -723,197 +596,6 @@
self.assertAllEqual([element for _ in range(10)], sess.run(get_next))
-class RestructuredDatasetTest(test.TestCase):
-
- def test_assert_element_shape(self):
-
- def create_dataset(_):
- return (array_ops.ones(2, dtype=dtypes.float32),
- array_ops.zeros((3, 4), dtype=dtypes.int32))
-
- dataset = dataset_ops.Dataset.range(5).map(create_dataset)
- expected_shapes = (tensor_shape.TensorShape(2),
- tensor_shape.TensorShape((3, 4)))
- self.assertEqual(expected_shapes, dataset.output_shapes)
-
- result = dataset.apply(batching.assert_element_shape(expected_shapes))
- self.assertEqual(expected_shapes, result.output_shapes)
-
- iterator = result.make_initializable_iterator()
- init_op = iterator.initializer
- get_next = iterator.get_next()
- with self.cached_session() as sess:
- sess.run(init_op)
- for _ in range(5):
- sess.run(get_next)
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
-
- def test_assert_wrong_element_shape(self):
-
- def create_dataset(_):
- return (array_ops.ones(2, dtype=dtypes.float32),
- array_ops.zeros((3, 4), dtype=dtypes.int32))
-
- dataset = dataset_ops.Dataset.range(3).map(create_dataset)
- wrong_shapes = (tensor_shape.TensorShape(2),
- tensor_shape.TensorShape((3, 10)))
- with self.assertRaises(ValueError):
- dataset.apply(batching.assert_element_shape(wrong_shapes))
-
- def test_assert_element_shape_on_unknown_shape_dataset(self):
-
- def create_unknown_shape_dataset(x):
- return script_ops.py_func(
- lambda _: ( # pylint: disable=g-long-lambda
- np.ones(2, dtype=np.float32),
- np.zeros((3, 4), dtype=np.int32)),
- [x],
- [dtypes.float32, dtypes.int32])
-
- dataset = dataset_ops.Dataset.range(5).map(create_unknown_shape_dataset)
- unknown_shapes = (tensor_shape.TensorShape(None),
- tensor_shape.TensorShape(None))
- self.assertEqual(unknown_shapes, dataset.output_shapes)
-
- expected_shapes = (tensor_shape.TensorShape(2),
- tensor_shape.TensorShape((3, 4)))
- result = dataset.apply(batching.assert_element_shape(expected_shapes))
- self.assertEqual(expected_shapes, result.output_shapes)
-
- iterator = result.make_initializable_iterator()
- init_op = iterator.initializer
- get_next = iterator.get_next()
- with self.cached_session() as sess:
- sess.run(init_op)
- for _ in range(5):
- sess.run(get_next)
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
-
- def test_assert_wrong_element_shape_on_unknown_shape_dataset(self):
-
- def create_unknown_shape_dataset(x):
- return script_ops.py_func(
- lambda _: ( # pylint: disable=g-long-lambda
- np.ones(2, dtype=np.float32),
- np.zeros((3, 4), dtype=np.int32)),
- [x],
- [dtypes.float32, dtypes.int32])
-
- dataset = dataset_ops.Dataset.range(3).map(create_unknown_shape_dataset)
- unknown_shapes = (tensor_shape.TensorShape(None),
- tensor_shape.TensorShape(None))
- self.assertEqual(unknown_shapes, dataset.output_shapes)
-
- wrong_shapes = (tensor_shape.TensorShape(2),
- tensor_shape.TensorShape((3, 10)))
- iterator = (
- dataset.apply(batching.assert_element_shape(wrong_shapes))
- .make_initializable_iterator())
- init_op = iterator.initializer
- get_next = iterator.get_next()
- with self.cached_session() as sess:
- sess.run(init_op)
- with self.assertRaises(errors.InvalidArgumentError):
- sess.run(get_next)
-
- def test_assert_partial_element_shape(self):
-
- def create_dataset(_):
- return (array_ops.ones(2, dtype=dtypes.float32),
- array_ops.zeros((3, 4), dtype=dtypes.int32))
-
- dataset = dataset_ops.Dataset.range(5).map(create_dataset)
- partial_expected_shape = (tensor_shape.TensorShape(None), # Unknown shape
- tensor_shape.TensorShape((None, 4))) # Partial shape
- result = dataset.apply(
- batching.assert_element_shape(partial_expected_shape))
- # Partial shapes are merged with actual shapes:
- actual_shapes = (tensor_shape.TensorShape(2),
- tensor_shape.TensorShape((3, 4)))
- self.assertEqual(actual_shapes, result.output_shapes)
-
- iterator = result.make_initializable_iterator()
- init_op = iterator.initializer
- get_next = iterator.get_next()
- with self.cached_session() as sess:
- sess.run(init_op)
- for _ in range(5):
- sess.run(get_next)
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
-
- def test_assert_wrong_partial_element_shape(self):
-
- def create_dataset(_):
- return (array_ops.ones(2, dtype=dtypes.float32),
- array_ops.zeros((3, 4), dtype=dtypes.int32))
-
- dataset = dataset_ops.Dataset.range(3).map(create_dataset)
- wrong_shapes = (tensor_shape.TensorShape(2),
- tensor_shape.TensorShape((None, 10)))
- with self.assertRaises(ValueError):
- dataset.apply(batching.assert_element_shape(wrong_shapes))
-
- def test_assert_partial_element_shape_on_unknown_shape_dataset(self):
-
- def create_unknown_shape_dataset(x):
- return script_ops.py_func(
- lambda _: ( # pylint: disable=g-long-lambda
- np.ones(2, dtype=np.float32),
- np.zeros((3, 4), dtype=np.int32)),
- [x],
- [dtypes.float32, dtypes.int32])
-
- dataset = dataset_ops.Dataset.range(5).map(create_unknown_shape_dataset)
- unknown_shapes = (tensor_shape.TensorShape(None),
- tensor_shape.TensorShape(None))
- self.assertEqual(unknown_shapes, dataset.output_shapes)
-
- expected_shapes = (tensor_shape.TensorShape(2),
- tensor_shape.TensorShape((None, 4)))
- result = dataset.apply(batching.assert_element_shape(expected_shapes))
- self.assertEqual(expected_shapes, result.output_shapes)
-
- iterator = result.make_initializable_iterator()
- init_op = iterator.initializer
- get_next = iterator.get_next()
- with self.cached_session() as sess:
- sess.run(init_op)
- for _ in range(5):
- sess.run(get_next)
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
-
- def test_assert_wrong_partial_element_shape_on_unknown_shape_dataset(self):
-
- def create_unknown_shape_dataset(x):
- return script_ops.py_func(
- lambda _: ( # pylint: disable=g-long-lambda
- np.ones(2, dtype=np.float32),
- np.zeros((3, 4), dtype=np.int32)),
- [x],
- [dtypes.float32, dtypes.int32])
-
- dataset = dataset_ops.Dataset.range(3).map(create_unknown_shape_dataset)
- unknown_shapes = (tensor_shape.TensorShape(None),
- tensor_shape.TensorShape(None))
- self.assertEqual(unknown_shapes, dataset.output_shapes)
-
- wrong_shapes = (tensor_shape.TensorShape(2),
- tensor_shape.TensorShape((None, 10)))
- iterator = (
- dataset.apply(batching.assert_element_shape(wrong_shapes))
- .make_initializable_iterator())
- init_op = iterator.initializer
- get_next = iterator.get_next()
- with self.cached_session() as sess:
- sess.run(init_op)
- with self.assertRaises(errors.InvalidArgumentError):
- sess.run(get_next)
-
-
class UnbatchDatasetBenchmark(test.Benchmark):
def benchmarkNativeUnbatch(self):
diff --git a/tensorflow/contrib/data/python/kernel_tests/bucketing_test.py b/tensorflow/python/data/experimental/kernel_tests/bucketing_test.py
similarity index 98%
rename from tensorflow/contrib/data/python/kernel_tests/bucketing_test.py
rename to tensorflow/python/data/experimental/kernel_tests/bucketing_test.py
index 48971f2..153a039 100644
--- a/tensorflow/contrib/data/python/kernel_tests/bucketing_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/bucketing_test.py
@@ -21,7 +21,8 @@
import numpy as np
-from tensorflow.contrib.data.python.ops import grouping
+from tensorflow.python.data.experimental.ops import grouping
+from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
@@ -35,7 +36,7 @@
from tensorflow.python.platform import test
-class GroupByReducerTest(test.TestCase):
+class GroupByReducerTest(test_base.DatasetTestBase):
def checkResults(self, dataset, shapes, values):
self.assertEqual(shapes, dataset.output_shapes)
@@ -198,7 +199,7 @@
self.assertEqual(y, 45)
-class GroupByWindowTest(test.TestCase):
+class GroupByWindowTest(test_base.DatasetTestBase):
def testSimple(self):
components = np.random.randint(100, size=(200,)).astype(np.int64)
@@ -345,7 +346,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.
-class BucketTest(test.TestCase):
+class BucketTest(test_base.DatasetTestBase):
def _dynamicPad(self, bucket, window, window_size):
# TODO(mrry): To match `tf.contrib.training.bucket()`, implement a
@@ -570,7 +571,7 @@
return tensor_shape.TensorShape([None])
-class BucketBySequenceLength(test.TestCase):
+class BucketBySequenceLength(test_base.DatasetTestBase):
def testBucket(self):
diff --git a/tensorflow/contrib/data/python/kernel_tests/csv_dataset_op_test.py b/tensorflow/python/data/experimental/kernel_tests/csv_dataset_op_test.py
similarity index 93%
rename from tensorflow/contrib/data/python/kernel_tests/csv_dataset_op_test.py
rename to tensorflow/python/data/experimental/kernel_tests/csv_dataset_op_test.py
index f8e74e4..4ee1779 100644
--- a/tensorflow/contrib/data/python/kernel_tests/csv_dataset_op_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/csv_dataset_op_test.py
@@ -27,9 +27,10 @@
import numpy as np
-from tensorflow.contrib.data.python.ops import error_ops
-from tensorflow.contrib.data.python.ops import readers
from tensorflow.python.client import session
+from tensorflow.python.data.experimental.ops import error_ops
+from tensorflow.python.data.experimental.ops import readers
+from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import readers as core_readers
from tensorflow.python.eager import context
from tensorflow.python.framework import constant_op
@@ -43,37 +44,7 @@
@test_util.run_all_in_graph_and_eager_modes
-class CsvDatasetOpTest(test.TestCase):
-
- def _get_next(self, dataset):
- # Returns a no argument function whose result is fed to self.evaluate to
- # yield the next element
- it = dataset.make_one_shot_iterator()
- if context.executing_eagerly():
- return it.get_next
- else:
- get_next = it.get_next()
- return lambda: get_next
-
- def _assert_datasets_equal(self, ds1, ds2):
- assert ds1.output_shapes == ds2.output_shapes, ('output_shapes differ: %s, '
- '%s') % (ds1.output_shapes,
- ds2.output_shapes)
- assert ds1.output_types == ds2.output_types
- assert ds1.output_classes == ds2.output_classes
- next1 = self._get_next(ds1)
- next2 = self._get_next(ds2)
- # Run through datasets and check that outputs match, or errors match.
- while True:
- try:
- op1 = self.evaluate(next1())
- except (errors.OutOfRangeError, ValueError) as e:
- # If op1 throws an exception, check that op2 throws same exception.
- with self.assertRaises(type(e)):
- self.evaluate(next2())
- break
- op2 = self.evaluate(next2())
- self.assertAllEqual(op1, op2)
+class CsvDatasetOpTest(test_base.DatasetTestBase):
def _setup_files(self, inputs, linebreak='\n', compression_type=None):
filenames = []
@@ -108,7 +79,7 @@
"""Checks that CsvDataset is equiv to TextLineDataset->map(decode_csv)."""
dataset_actual, dataset_expected = self._make_test_datasets(
inputs, **kwargs)
- self._assert_datasets_equal(dataset_actual, dataset_expected)
+ self.assertDatasetsEqual(dataset_actual, dataset_expected)
def _verify_output_or_err(self,
dataset,
@@ -116,7 +87,7 @@
expected_err_re=None):
if expected_err_re is None:
# Verify that output is expected, without errors
- nxt = self._get_next(dataset)
+ nxt = self.getNext(dataset)
expected_output = [[
v.encode('utf-8') if isinstance(v, str) else v for v in op
] for op in expected_output]
@@ -128,7 +99,7 @@
else:
# Verify that OpError is produced as expected
with self.assertRaisesOpError(expected_err_re):
- nxt = self._get_next(dataset)
+ nxt = self.getNext(dataset)
while True:
try:
self.evaluate(nxt())
@@ -354,7 +325,7 @@
inputs = [['1,,3,4', '5,6,,8']]
ds_actual, ds_expected = self._make_test_datasets(
inputs, record_defaults=record_defaults)
- self._assert_datasets_equal(
+ self.assertDatasetsEqual(
ds_actual.repeat(5).prefetch(1),
ds_expected.repeat(5).prefetch(1))
@@ -377,7 +348,7 @@
ds = readers.make_csv_dataset(
file_path, batch_size=1, shuffle=False, num_epochs=1)
- nxt = self._get_next(ds)
+ nxt = self.getNext(ds)
result = list(self.evaluate(nxt()).values())
diff --git a/tensorflow/contrib/data/python/kernel_tests/dataset_constructor_op_test.py b/tensorflow/python/data/experimental/kernel_tests/dataset_constructor_op_test.py
similarity index 93%
rename from tensorflow/contrib/data/python/kernel_tests/dataset_constructor_op_test.py
rename to tensorflow/python/data/experimental/kernel_tests/dataset_constructor_op_test.py
index a2ab3de..3fc7157 100644
--- a/tensorflow/contrib/data/python/kernel_tests/dataset_constructor_op_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/dataset_constructor_op_test.py
@@ -17,7 +17,8 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.contrib.data.python.ops import batching
+from tensorflow.python.data.experimental.ops import batching
+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 dtypes
@@ -25,7 +26,7 @@
from tensorflow.python.platform import test
-class DatasetConstructorTest(test.TestCase):
+class DatasetConstructorTest(test_base.DatasetTestBase):
def testRestructureDataset(self):
components = (array_ops.placeholder(dtypes.int32),
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/dataset_serialization_test_base.py b/tensorflow/python/data/experimental/kernel_tests/dataset_serialization_test_base.py
similarity index 99%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/dataset_serialization_test_base.py
rename to tensorflow/python/data/experimental/kernel_tests/dataset_serialization_test_base.py
index 595cece..7f435b8 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/dataset_serialization_test_base.py
+++ b/tensorflow/python/data/experimental/kernel_tests/dataset_serialization_test_base.py
@@ -22,7 +22,7 @@
import numpy as np
-from tensorflow.contrib.data.python.ops import iterator_ops as contrib_iterator_ops
+from tensorflow.python.data.experimental.ops import iterator_ops as contrib_iterator_ops
from tensorflow.python.data.ops import iterator_ops
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
diff --git a/tensorflow/contrib/data/python/kernel_tests/directed_interleave_dataset_test.py b/tensorflow/python/data/experimental/kernel_tests/directed_interleave_dataset_test.py
similarity index 95%
rename from tensorflow/contrib/data/python/kernel_tests/directed_interleave_dataset_test.py
rename to tensorflow/python/data/experimental/kernel_tests/directed_interleave_dataset_test.py
index eb11032..796a692 100644
--- a/tensorflow/contrib/data/python/kernel_tests/directed_interleave_dataset_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/directed_interleave_dataset_test.py
@@ -19,14 +19,15 @@
import numpy as np
-from tensorflow.contrib.data.python.ops import interleave_ops
+from tensorflow.python.data.experimental.ops import interleave_ops
+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 random_seed
from tensorflow.python.platform import test
-class DirectedInterleaveDatasetTest(test.TestCase):
+class DirectedInterleaveDatasetTest(test_base.DatasetTestBase):
def testBasic(self):
selector_dataset = dataset_ops.Dataset.range(10).repeat(100)
@@ -83,7 +84,7 @@
# Use chi-squared test to assert that the observed distribution matches the
# expected distribution. Based on the implementation in
- # "tensorflow/python/kernel_tests/multinomial_op_test.py".
+ # "third_party/tensorflow/python/kernel_tests/multinomial_op_test.py".
for probs in [[.85, .05, .1], rand_probs, [1.]]:
probs = np.asarray(probs)
classes = len(probs)
diff --git a/tensorflow/contrib/data/python/kernel_tests/filter_dataset_op_test.py b/tensorflow/python/data/experimental/kernel_tests/filter_dataset_op_test.py
similarity index 97%
rename from tensorflow/contrib/data/python/kernel_tests/filter_dataset_op_test.py
rename to tensorflow/python/data/experimental/kernel_tests/filter_dataset_op_test.py
index 6d01bf5..c6ee88c 100644
--- a/tensorflow/contrib/data/python/kernel_tests/filter_dataset_op_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/filter_dataset_op_test.py
@@ -21,8 +21,8 @@
import numpy as np
-from tensorflow.contrib.data.python.ops import optimization
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
diff --git a/tensorflow/contrib/data/python/kernel_tests/get_single_element_test.py b/tensorflow/python/data/experimental/kernel_tests/get_single_element_test.py
similarity index 73%
rename from tensorflow/contrib/data/python/kernel_tests/get_single_element_test.py
rename to tensorflow/python/data/experimental/kernel_tests/get_single_element_test.py
index f3968cd..8c07afb 100644
--- a/tensorflow/contrib/data/python/kernel_tests/get_single_element_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/get_single_element_test.py
@@ -18,10 +18,9 @@
from __future__ import print_function
from absl.testing import parameterized
-import numpy as np
-from tensorflow.contrib.data.python.ops import get_single_element
-from tensorflow.contrib.data.python.ops import grouping
+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
@@ -30,7 +29,7 @@
from tensorflow.python.platform import test
-class GetSingleElementTest(test.TestCase, parameterized.TestCase):
+class GetSingleElementTest(test_base.DatasetTestBase, parameterized.TestCase):
@parameterized.named_parameters(
("Zero", 0, 1),
@@ -68,32 +67,6 @@
with self.assertRaisesRegexp(error, error_msg):
sess.run(element, feed_dict={skip_t: skip, take_t: take})
- @parameterized.named_parameters(
- ("SumZero", 0),
- ("SumOne", 1),
- ("SumFive", 5),
- ("SumTen", 10),
- )
- def testReduceDataset(self, stop):
- def init_fn(_):
- return np.int64(0)
-
- def reduce_fn(state, value):
- return state + value
-
- def finalize_fn(state):
- return state
-
- sum_reducer = grouping.Reducer(init_fn, reduce_fn, finalize_fn)
-
- stop_t = array_ops.placeholder(dtypes.int64, shape=[])
- dataset = dataset_ops.Dataset.range(stop_t)
- element = get_single_element.reduce_dataset(dataset, sum_reducer)
-
- with self.cached_session() as sess:
- value = sess.run(element, feed_dict={stop_t: stop})
- self.assertEqual(stop * (stop - 1) / 2, value)
-
if __name__ == "__main__":
test.main()
diff --git a/tensorflow/contrib/data/python/kernel_tests/indexed_dataset_ops_test.py b/tensorflow/python/data/experimental/kernel_tests/indexed_dataset_ops_test.py
similarity index 82%
rename from tensorflow/contrib/data/python/kernel_tests/indexed_dataset_ops_test.py
rename to tensorflow/python/data/experimental/kernel_tests/indexed_dataset_ops_test.py
index 9c508d6..c93a835 100644
--- a/tensorflow/contrib/data/python/kernel_tests/indexed_dataset_ops_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/indexed_dataset_ops_test.py
@@ -19,29 +19,30 @@
import unittest
-from tensorflow.contrib.data.python.ops import contrib_op_loader # pylint: disable=unused-import
-from tensorflow.contrib.data.python.ops import gen_dataset_ops
-from tensorflow.contrib.data.python.ops import indexed_dataset_ops
+from tensorflow.python.data.experimental.ops import indexed_dataset_ops
+from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
from tensorflow.python.platform import test
-class IndexedDatasetOpsTest(test.TestCase):
+class IndexedDatasetOpsTest(test_base.DatasetTestBase):
def testLowLevelIndexedDatasetOps(self):
- identity = gen_dataset_ops.identity_indexed_dataset(
+ identity = ged_ops.experimental_identity_indexed_dataset(
ops.convert_to_tensor(16, dtype=dtypes.uint64))
- handle = gen_dataset_ops.materialized_index_dataset_handle(
+ handle = ged_ops.experimental_materialized_index_dataset_handle(
container="",
shared_name="",
output_types=[dtypes.uint64],
output_shapes=[[]])
- materialize = gen_dataset_ops.indexed_dataset_materialize(identity, handle)
+ materialize = ged_ops.experimental_indexed_dataset_materialize(
+ identity, handle)
index = array_ops.placeholder(dtypes.uint64)
- get_op = gen_dataset_ops.indexed_dataset_get(
+ get_op = ged_ops.experimental_indexed_dataset_get(
handle, index, output_types=[dtypes.uint64], output_shapes=[[]])
with self.cached_session() as sess:
diff --git a/tensorflow/contrib/data/python/kernel_tests/interleave_dataset_op_test.py b/tensorflow/python/data/experimental/kernel_tests/interleave_dataset_op_test.py
similarity index 99%
rename from tensorflow/contrib/data/python/kernel_tests/interleave_dataset_op_test.py
rename to tensorflow/python/data/experimental/kernel_tests/interleave_dataset_op_test.py
index b9e74df..560902c 100644
--- a/tensorflow/contrib/data/python/kernel_tests/interleave_dataset_op_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/interleave_dataset_op_test.py
@@ -24,7 +24,8 @@
from six.moves import zip_longest
-from tensorflow.contrib.data.python.ops import interleave_ops
+from tensorflow.python.data.experimental.ops import interleave_ops
+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
@@ -36,7 +37,7 @@
from tensorflow.python.platform import test
-class ParallelInterleaveDatasetTest(test.TestCase):
+class ParallelInterleaveDatasetTest(test_base.DatasetTestBase):
def setUp(self):
diff --git a/tensorflow/contrib/data/python/kernel_tests/iterator_ops_test.py b/tensorflow/python/data/experimental/kernel_tests/iterator_ops_test.py
similarity index 94%
rename from tensorflow/contrib/data/python/kernel_tests/iterator_ops_test.py
rename to tensorflow/python/data/experimental/kernel_tests/iterator_ops_test.py
index 704c0d1..94393d6 100644
--- a/tensorflow/contrib/data/python/kernel_tests/iterator_ops_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/iterator_ops_test.py
@@ -18,7 +18,8 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.contrib.data.python.ops import iterator_ops
+from tensorflow.python.data.experimental.ops import iterator_ops
+from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.estimator import estimator
from tensorflow.python.estimator import model_fn
@@ -33,7 +34,7 @@
from tensorflow.python.training import training_util
-class CheckpointInputPipelineHookTest(test.TestCase):
+class CheckpointInputPipelineHookTest(test_base.DatasetTestBase):
@staticmethod
def _model_fn(features, labels, mode, config):
@@ -42,7 +43,7 @@
del config
global_step = training_util.get_or_create_global_step()
update_global_step_op = global_step.assign_add(1)
- latest_feature = variables.Variable(
+ latest_feature = variables.VariableV1(
0, name='latest_feature', dtype=dtypes.int64)
store_latest_feature_op = latest_feature.assign(features)
ops.add_to_collection('my_vars', global_step)
diff --git a/tensorflow/contrib/data/python/kernel_tests/map_dataset_op_test.py b/tensorflow/python/data/experimental/kernel_tests/map_dataset_op_test.py
similarity index 97%
rename from tensorflow/contrib/data/python/kernel_tests/map_dataset_op_test.py
rename to tensorflow/python/data/experimental/kernel_tests/map_dataset_op_test.py
index e851938..2f0bd14 100644
--- a/tensorflow/contrib/data/python/kernel_tests/map_dataset_op_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/map_dataset_op_test.py
@@ -24,11 +24,12 @@
import numpy as np
-from tensorflow.contrib.data.python.ops import batching
-from tensorflow.contrib.data.python.ops import error_ops
-from tensorflow.contrib.data.python.ops import optimization
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 import error_ops
+from tensorflow.python.data.experimental.ops import optimization
+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
@@ -41,7 +42,7 @@
_NUMPY_RANDOM_SEED = 42
-class MapDatasetTest(test.TestCase):
+class MapDatasetTest(test_base.DatasetTestBase):
def testMapIgnoreError(self):
components = np.array([1., 2., 3., np.nan, 5.]).astype(np.float32)
diff --git a/tensorflow/contrib/data/python/kernel_tests/map_defun_op_test.py b/tensorflow/python/data/experimental/kernel_tests/map_defun_op_test.py
similarity index 97%
rename from tensorflow/contrib/data/python/kernel_tests/map_defun_op_test.py
rename to tensorflow/python/data/experimental/kernel_tests/map_defun_op_test.py
index 25aea03..612ee33 100644
--- a/tensorflow/contrib/data/python/kernel_tests/map_defun_op_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/map_defun_op_test.py
@@ -19,8 +19,9 @@
import time
-from tensorflow.contrib.data.python.ops import map_defun
from tensorflow.python.client import session
+from tensorflow.python.data.experimental.ops import map_defun
+from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
@@ -33,7 +34,8 @@
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test
-class MapDefunTest(test.TestCase):
+
+class MapDefunTest(test_base.DatasetTestBase):
def testMapDefunSimple(self):
diff --git a/tensorflow/contrib/data/python/kernel_tests/optimization/BUILD b/tensorflow/python/data/experimental/kernel_tests/optimization/BUILD
similarity index 63%
rename from tensorflow/contrib/data/python/kernel_tests/optimization/BUILD
rename to tensorflow/python/data/experimental/kernel_tests/optimization/BUILD
index a2fc244..68f73bd 100644
--- a/tensorflow/contrib/data/python/kernel_tests/optimization/BUILD
+++ b/tensorflow/python/data/experimental/kernel_tests/optimization/BUILD
@@ -12,24 +12,43 @@
srcs = ["assert_next_dataset_op_test.py"],
srcs_version = "PY2AND3",
deps = [
- "//tensorflow/contrib/data/python/ops:optimization",
"//tensorflow/python:client_testlib",
"//tensorflow/python:errors",
+ "//tensorflow/python/data/experimental/ops:optimization",
+ "//tensorflow/python/data/kernel_tests:test_base",
"//tensorflow/python/data/ops:dataset_ops",
],
)
py_test(
+ name = "hoist_random_uniform_test",
+ size = "small",
+ srcs = ["hoist_random_uniform_test.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:constant_op",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:errors",
+ "//tensorflow/python:math_ops",
+ "//tensorflow/python/data/experimental/ops:optimization",
+ "//tensorflow/python/data/kernel_tests:test_base",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "@absl_py//absl/testing:parameterized",
+ ],
+)
+
+py_test(
name = "latency_all_edges_test",
size = "small",
srcs = ["latency_all_edges_test.py"],
srcs_version = "PY2AND3",
deps = [
- "//tensorflow/contrib/data/python/kernel_tests:stats_dataset_test_base",
- "//tensorflow/contrib/data/python/ops:optimization",
- "//tensorflow/contrib/data/python/ops:stats_ops",
"//tensorflow/python:client_testlib",
"//tensorflow/python:errors",
+ "//tensorflow/python/data/experimental/kernel_tests:stats_dataset_test_base",
+ "//tensorflow/python/data/experimental/ops:optimization",
+ "//tensorflow/python/data/experimental/ops:stats_ops",
"//tensorflow/python/data/ops:dataset_ops",
],
)
@@ -40,8 +59,6 @@
srcs = ["map_vectorization_test.py"],
srcs_version = "PY2AND3",
deps = [
- "//tensorflow/contrib/data/python/kernel_tests:test_utils",
- "//tensorflow/contrib/data/python/ops:optimization",
"//tensorflow/python:check_ops",
"//tensorflow/python:client_testlib",
"//tensorflow/python:constant_op",
@@ -50,6 +67,8 @@
"//tensorflow/python:framework_ops",
"//tensorflow/python:math_ops",
"//tensorflow/python:session",
+ "//tensorflow/python/data/experimental/ops:optimization",
+ "//tensorflow/python/data/kernel_tests:test_base",
"//tensorflow/python/data/ops:dataset_ops",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
@@ -62,12 +81,13 @@
srcs = ["map_and_filter_fusion_test.py"],
srcs_version = "PY2AND3",
deps = [
- "//tensorflow/contrib/data/python/ops:optimization",
"//tensorflow/python:client_testlib",
"//tensorflow/python:constant_op",
"//tensorflow/python:dtypes",
"//tensorflow/python:errors",
"//tensorflow/python:math_ops",
+ "//tensorflow/python/data/experimental/ops:optimization",
+ "//tensorflow/python/data/kernel_tests:test_base",
"//tensorflow/python/data/ops:dataset_ops",
"@absl_py//absl/testing:parameterized",
],
@@ -79,12 +99,13 @@
srcs = ["map_parallelization_test.py"],
srcs_version = "PY2AND3",
deps = [
- "//tensorflow/contrib/data/python/ops:optimization",
"//tensorflow/python:client_testlib",
"//tensorflow/python:constant_op",
"//tensorflow/python:dtypes",
"//tensorflow/python:errors",
"//tensorflow/python:math_ops",
+ "//tensorflow/python/data/experimental/ops:optimization",
+ "//tensorflow/python/data/kernel_tests:test_base",
"//tensorflow/python/data/ops:dataset_ops",
"@absl_py//absl/testing:parameterized",
],
@@ -99,11 +120,12 @@
"optonly",
],
deps = [
- "//tensorflow/contrib/data/python/ops:batching",
- "//tensorflow/contrib/data/python/ops:interleave_ops",
- "//tensorflow/contrib/data/python/ops:optimization",
"//tensorflow/python:client_testlib",
"//tensorflow/python:errors",
+ "//tensorflow/python/data/experimental/ops:batching",
+ "//tensorflow/python/data/experimental/ops:interleave_ops",
+ "//tensorflow/python/data/experimental/ops:optimization",
+ "//tensorflow/python/data/kernel_tests:test_base",
"//tensorflow/python/data/ops:dataset_ops",
"//third_party/py/numpy",
],
@@ -115,11 +137,12 @@
srcs = ["noop_elimination_test.py"],
srcs_version = "PY2AND3",
deps = [
- "//tensorflow/contrib/data/python/ops:batching",
- "//tensorflow/contrib/data/python/ops:interleave_ops",
- "//tensorflow/contrib/data/python/ops:optimization",
"//tensorflow/python:client_testlib",
"//tensorflow/python:errors",
+ "//tensorflow/python/data/experimental/ops:batching",
+ "//tensorflow/python/data/experimental/ops:interleave_ops",
+ "//tensorflow/python/data/experimental/ops:optimization",
+ "//tensorflow/python/data/kernel_tests:test_base",
"//tensorflow/python/data/ops:dataset_ops",
"//third_party/py/numpy",
],
@@ -131,9 +154,10 @@
srcs = ["optimize_dataset_op_test.py"],
srcs_version = "PY2AND3",
deps = [
- "//tensorflow/contrib/data/python/ops:optimization",
"//tensorflow/python:client_testlib",
"//tensorflow/python:errors",
+ "//tensorflow/python/data/experimental/ops:optimization",
+ "//tensorflow/python/data/kernel_tests:test_base",
"//tensorflow/python/data/ops:dataset_ops",
"//third_party/py/numpy",
],
diff --git a/tensorflow/contrib/data/python/kernel_tests/optimization/assert_next_dataset_op_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/assert_next_dataset_op_test.py
similarity index 92%
rename from tensorflow/contrib/data/python/kernel_tests/optimization/assert_next_dataset_op_test.py
rename to tensorflow/python/data/experimental/kernel_tests/optimization/assert_next_dataset_op_test.py
index d10da80..45b77b5 100644
--- a/tensorflow/contrib/data/python/kernel_tests/optimization/assert_next_dataset_op_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/optimization/assert_next_dataset_op_test.py
@@ -17,13 +17,14 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.contrib.data.python.ops import optimization
+from tensorflow.python.data.experimental.ops import optimization
+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.platform import test
-class AssertNextDatasetTest(test.TestCase):
+class AssertNextDatasetTest(test_base.DatasetTestBase):
def testAssertNext(self):
dataset = dataset_ops.Dataset.from_tensors(0).apply(
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
new file mode 100644
index 0000000..3cd9753
--- /dev/null
+++ b/tensorflow/python/data/experimental/kernel_tests/optimization/hoist_random_uniform_test.py
@@ -0,0 +1,103 @@
+# 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 HostState optimization."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.experimental.ops import optimization
+from tensorflow.python.data.kernel_tests import test_base
+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 ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.platform import test
+
+
+class HoistRandomUniformTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+ @staticmethod
+ def map_functions():
+ plus_one = lambda x: x + 1
+
+ def random(_):
+ return random_ops.random_uniform([],
+ minval=1,
+ maxval=10,
+ dtype=dtypes.float32,
+ seed=42)
+
+ def random_with_assert(x):
+ y = random(x)
+ assert_op = control_flow_ops.Assert(math_ops.greater_equal(y, 1), [y])
+ with ops.control_dependencies([assert_op]):
+ return y
+
+ twice_random = lambda x: (random(x) + random(x)) / 2.
+
+ tests = [("PlusOne", plus_one, False), ("RandomUniform", random, True),
+ ("RandomWithAssert", random_with_assert, True),
+ ("TwiceRandom", twice_random, False)]
+ return tuple(tests)
+
+ @parameterized.named_parameters(*map_functions.__func__())
+ def testHoisting(self, function, will_optimize):
+ dataset = dataset_ops.Dataset.range(5).apply(
+ optimization.assert_next(
+ ["Zip[0]", "Map"] if will_optimize else ["Map"])).map(function)
+
+ dataset = dataset.apply(optimization.optimize(["hoist_random_uniform"]))
+ self._testDataset(dataset)
+
+ def testAdditionalInputs(self):
+ a = constant_op.constant(1, dtype=dtypes.float32)
+ b = constant_op.constant(0, dtype=dtypes.float32)
+ some_tensor = math_ops.mul(a, b)
+
+ def random_with_capture(_):
+ return some_tensor + random_ops.random_uniform(
+ [], minval=1, maxval=10, dtype=dtypes.float32, seed=42)
+
+ dataset = dataset_ops.Dataset.range(5).apply(
+ optimization.assert_next(
+ ["Zip[0]", "Map"])).map(random_with_capture).apply(
+ optimization.optimize(["hoist_random_uniform"]))
+ self._testDataset(dataset)
+
+ def _testDataset(self, dataset):
+ iterator = dataset.make_one_shot_iterator()
+ get_next = iterator.get_next()
+ previous_result = 0
+ with self.cached_session() as sess:
+ for _ in range(5):
+ result = sess.run(get_next)
+ self.assertLessEqual(1, result)
+ self.assertLessEqual(result, 10)
+ # This checks if the result is somehow random by checking if we are not
+ # generating the same values.
+ self.assertNotEqual(previous_result, result)
+ previous_result = result
+ with self.assertRaises(errors.OutOfRangeError):
+ sess.run(get_next)
+
+
+if __name__ == "__main__":
+ test.main()
diff --git a/tensorflow/contrib/data/python/kernel_tests/optimization/latency_all_edges_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/latency_all_edges_test.py
similarity index 85%
rename from tensorflow/contrib/data/python/kernel_tests/optimization/latency_all_edges_test.py
rename to tensorflow/python/data/experimental/kernel_tests/optimization/latency_all_edges_test.py
index db380c0..4562387 100644
--- a/tensorflow/contrib/data/python/kernel_tests/optimization/latency_all_edges_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/optimization/latency_all_edges_test.py
@@ -17,9 +17,9 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.contrib.data.python.kernel_tests import stats_dataset_test_base
-from tensorflow.contrib.data.python.ops import optimization
-from tensorflow.contrib.data.python.ops import stats_ops
+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_ops
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import errors
from tensorflow.python.platform import test
@@ -34,8 +34,8 @@
optimization.assert_next(
["LatencyStats", "Map", "LatencyStats", "Prefetch",
"LatencyStats"])).map(lambda x: x * x).prefetch(1).apply(
- optimization.optimize(["latency_all_edges"])).apply(
- stats_ops.set_stats_aggregator(stats_aggregator))
+ stats_ops.set_stats_aggregator(stats_aggregator)).apply(
+ optimization.optimize(["latency_all_edges"]))
iterator = dataset.make_initializable_iterator()
get_next = iterator.get_next()
summary_t = stats_aggregator.get_summary()
diff --git a/tensorflow/contrib/data/python/kernel_tests/optimization/map_and_filter_fusion_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/map_and_filter_fusion_test.py
similarity index 97%
rename from tensorflow/contrib/data/python/kernel_tests/optimization/map_and_filter_fusion_test.py
rename to tensorflow/python/data/experimental/kernel_tests/optimization/map_and_filter_fusion_test.py
index e75edf6..a439635 100644
--- a/tensorflow/contrib/data/python/kernel_tests/optimization/map_and_filter_fusion_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/optimization/map_and_filter_fusion_test.py
@@ -19,7 +19,8 @@
from absl.testing import parameterized
-from tensorflow.contrib.data.python.ops import optimization
+from tensorflow.python.data.experimental.ops import optimization
+from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
@@ -28,7 +29,7 @@
from tensorflow.python.platform import test
-class MapAndFilterFusionTest(test.TestCase, parameterized.TestCase):
+class MapAndFilterFusionTest(test_base.DatasetTestBase, parameterized.TestCase):
@staticmethod
def map_functions():
diff --git a/tensorflow/contrib/data/python/kernel_tests/optimization/map_parallelization_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/map_parallelization_test.py
similarity index 93%
rename from tensorflow/contrib/data/python/kernel_tests/optimization/map_parallelization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/optimization/map_parallelization_test.py
index dd547db..334d8e3 100644
--- a/tensorflow/contrib/data/python/kernel_tests/optimization/map_parallelization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/optimization/map_parallelization_test.py
@@ -19,7 +19,8 @@
from absl.testing import parameterized
-from tensorflow.contrib.data.python.ops import optimization
+from tensorflow.python.data.experimental.ops import optimization
+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
@@ -30,7 +31,7 @@
from tensorflow.python.platform import test
-class MapParallelizationTest(test.TestCase, parameterized.TestCase):
+class MapParallelizationTest(test_base.DatasetTestBase, parameterized.TestCase):
@staticmethod
def map_functions():
diff --git a/tensorflow/contrib/data/python/kernel_tests/optimization/map_vectorization_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/map_vectorization_test.py
similarity index 94%
rename from tensorflow/contrib/data/python/kernel_tests/optimization/map_vectorization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/optimization/map_vectorization_test.py
index 5b493f4..d474927 100644
--- a/tensorflow/contrib/data/python/kernel_tests/optimization/map_vectorization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/optimization/map_vectorization_test.py
@@ -22,9 +22,9 @@
from absl.testing import parameterized
import numpy as np
-from tensorflow.contrib.data.python.kernel_tests import test_utils
-from tensorflow.contrib.data.python.ops import optimization
from tensorflow.python.client import session
+from tensorflow.python.data.experimental.ops import optimization
+from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
@@ -36,7 +36,7 @@
from tensorflow.python.platform import test
-class MapVectorizationTest(test_utils.DatasetTestBase, parameterized.TestCase):
+class MapVectorizationTest(test_base.DatasetTestBase, parameterized.TestCase):
def _get_test_datasets(self,
base_dataset,
@@ -85,7 +85,7 @@
[3, 4]]).repeat(5)
unoptimized, optimized = self._get_test_datasets(base_dataset, map_fn,
num_parallel_calls)
- self._assert_datasets_equal(unoptimized, optimized)
+ self.assertDatasetsEqual(unoptimized, optimized)
def testOptimizationBadMapFn(self):
# Test map functions that give an error
@@ -112,7 +112,7 @@
# TODO(rachelim): when this optimization works, turn on expect_optimized
unoptimized, optimized = self._get_test_datasets(
base_dataset, map_fn, expect_optimized=False)
- self._assert_datasets_equal(optimized, unoptimized)
+ self.assertDatasetsEqual(optimized, unoptimized)
def testOptimizationIgnoreStateful(self):
@@ -124,7 +124,7 @@
[3, 4]]).repeat(5)
unoptimized, optimized = self._get_test_datasets(
base_dataset, map_fn, expect_optimized=False)
- self._assert_datasets_raise_same_error(
+ self.assertDatasetsRaiseSameError(
unoptimized, optimized, errors.InvalidArgumentError,
[("OneShotIterator", "OneShotIterator_1", 1),
("IteratorGetNext", "IteratorGetNext_1", 1)])
@@ -138,7 +138,7 @@
base_dataset = dataset_ops.Dataset.range(20).batch(3, drop_remainder=False)
unoptimized, optimized = self._get_test_datasets(
base_dataset, map_fn, expect_optimized=False)
- self._assert_datasets_equal(unoptimized, optimized)
+ self.assertDatasetsEqual(unoptimized, optimized)
def testOptimizationIgnoreRaggedMap(self):
# Don't optimize when the output of the map fn shapes are unknown.
@@ -148,7 +148,7 @@
base_dataset = dataset_ops.Dataset.range(20).batch(1, drop_remainder=True)
unoptimized, optimized = self._get_test_datasets(
base_dataset, map_fn, expect_optimized=False)
- self._assert_datasets_raise_same_error(
+ self.assertDatasetsRaiseSameError(
unoptimized, optimized, errors.InvalidArgumentError,
[("OneShotIterator", "OneShotIterator_1", 1),
("IteratorGetNext", "IteratorGetNext_1", 1)])
diff --git a/tensorflow/contrib/data/python/kernel_tests/optimization/model_dataset_op_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/model_dataset_op_test.py
similarity index 96%
rename from tensorflow/contrib/data/python/kernel_tests/optimization/model_dataset_op_test.py
rename to tensorflow/python/data/experimental/kernel_tests/optimization/model_dataset_op_test.py
index 3b62a7e..a9f2ce8 100644
--- a/tensorflow/contrib/data/python/kernel_tests/optimization/model_dataset_op_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/optimization/model_dataset_op_test.py
@@ -21,14 +21,15 @@
import numpy as np
-from tensorflow.contrib.data.python.ops import batching
-from tensorflow.contrib.data.python.ops import optimization
+from tensorflow.python.data.experimental.ops import batching
+from tensorflow.python.data.experimental.ops import optimization
+from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test
-class ModelDatasetTest(test.TestCase):
+class ModelDatasetTest(test_base.DatasetTestBase):
def testModelMap(self):
k = 1024 * 1024
diff --git a/tensorflow/contrib/data/python/kernel_tests/optimization/noop_elimination_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/noop_elimination_test.py
similarity index 91%
rename from tensorflow/contrib/data/python/kernel_tests/optimization/noop_elimination_test.py
rename to tensorflow/python/data/experimental/kernel_tests/optimization/noop_elimination_test.py
index 507feda..092e0ff 100644
--- a/tensorflow/contrib/data/python/kernel_tests/optimization/noop_elimination_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/optimization/noop_elimination_test.py
@@ -17,7 +17,8 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.contrib.data.python.ops import optimization
+from tensorflow.python.data.experimental.ops import optimization
+from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
@@ -26,7 +27,7 @@
from tensorflow.python.platform import test
-class NoopEliminationTest(test.TestCase):
+class NoopEliminationTest(test_base.DatasetTestBase):
def testNoopElimination(self):
a = constant_op.constant(1, dtype=dtypes.int64)
diff --git a/tensorflow/contrib/data/python/kernel_tests/optimization/optimize_dataset_op_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/optimize_dataset_op_test.py
similarity index 95%
rename from tensorflow/contrib/data/python/kernel_tests/optimization/optimize_dataset_op_test.py
rename to tensorflow/python/data/experimental/kernel_tests/optimization/optimize_dataset_op_test.py
index a3fb824..eb66179 100644
--- a/tensorflow/contrib/data/python/kernel_tests/optimization/optimize_dataset_op_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/optimization/optimize_dataset_op_test.py
@@ -19,7 +19,8 @@
import numpy as np
-from tensorflow.contrib.data.python.ops import optimization
+from tensorflow.python.data.experimental.ops import optimization
+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
@@ -28,7 +29,7 @@
from tensorflow.python.platform import test
-class OptimizeDatasetTest(test.TestCase):
+class OptimizeDatasetTest(test_base.DatasetTestBase):
def testOptimizationDefault(self):
dataset = dataset_ops.Dataset.range(10).apply(
diff --git a/tensorflow/contrib/data/python/kernel_tests/parsing_ops_test.py b/tensorflow/python/data/experimental/kernel_tests/parsing_ops_test.py
similarity index 98%
rename from tensorflow/contrib/data/python/kernel_tests/parsing_ops_test.py
rename to tensorflow/python/data/experimental/kernel_tests/parsing_ops_test.py
index c4623bc..13f924b 100644
--- a/tensorflow/contrib/data/python/kernel_tests/parsing_ops_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/parsing_ops_test.py
@@ -22,9 +22,10 @@
import numpy as np
-from tensorflow.contrib.data.python.ops import parsing_ops as contrib_parsing_ops
from tensorflow.core.example import example_pb2
from tensorflow.core.example import feature_pb2
+from tensorflow.python.data.experimental.ops import parsing_ops as contrib_parsing_ops
+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 dtypes
@@ -72,7 +73,7 @@
i += 1
-class ParseExampleTest(test.TestCase):
+class ParseExampleTest(test_base.DatasetTestBase):
def _test(self,
input_tensor,
@@ -845,6 +846,5 @@
"allow_missing to be True."))
-
if __name__ == "__main__":
test.main()
diff --git a/tensorflow/contrib/data/python/kernel_tests/prefetching_ops_test.py b/tensorflow/python/data/experimental/kernel_tests/prefetching_ops_test.py
similarity index 98%
rename from tensorflow/contrib/data/python/kernel_tests/prefetching_ops_test.py
rename to tensorflow/python/data/experimental/kernel_tests/prefetching_ops_test.py
index 33a64ea..7d7b842 100644
--- a/tensorflow/contrib/data/python/kernel_tests/prefetching_ops_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/prefetching_ops_test.py
@@ -19,9 +19,10 @@
import threading
-from tensorflow.contrib.data.python.ops import prefetching_ops
from tensorflow.core.protobuf import config_pb2
from tensorflow.python.compat import compat
+from tensorflow.python.data.experimental.ops import prefetching_ops
+from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import iterator_ops
from tensorflow.python.framework import constant_op
@@ -35,7 +36,7 @@
from tensorflow.python.platform import test
-class PrefetchingKernelsOpsTest(test.TestCase):
+class PrefetchingKernelsOpsTest(test_base.DatasetTestBase):
def setUp(self):
self._event = threading.Event()
@@ -244,7 +245,7 @@
sess.run(destroy_op)
-class PrefetchToDeviceTest(test.TestCase):
+class PrefetchToDeviceTest(test_base.DatasetTestBase):
def testPrefetchToDevice(self):
host_dataset = dataset_ops.Dataset.range(10)
@@ -445,7 +446,7 @@
sess.run(next_element)
-class CopyToDeviceTest(test.TestCase):
+class CopyToDeviceTest(test_base.DatasetTestBase):
def testCopyToDevice(self):
host_dataset = dataset_ops.Dataset.range(10)
diff --git a/tensorflow/contrib/data/python/kernel_tests/range_dataset_op_test.py b/tensorflow/python/data/experimental/kernel_tests/range_dataset_op_test.py
similarity index 92%
rename from tensorflow/contrib/data/python/kernel_tests/range_dataset_op_test.py
rename to tensorflow/python/data/experimental/kernel_tests/range_dataset_op_test.py
index db8fe6a..22412c3 100644
--- a/tensorflow/contrib/data/python/kernel_tests/range_dataset_op_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/range_dataset_op_test.py
@@ -17,8 +17,9 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.contrib.data.python.ops import counter
-from tensorflow.contrib.data.python.ops import enumerate_ops
+from tensorflow.python.data.experimental.ops import counter
+from tensorflow.python.data.experimental.ops import enumerate_ops
+from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
@@ -27,7 +28,7 @@
from tensorflow.python.platform import test
-class RangeDatasetTest(test.TestCase):
+class RangeDatasetTest(test_base.DatasetTestBase):
def testEnumerateDataset(self):
components = (["a", "b"], [1, 2], [37.0, 38])
diff --git a/tensorflow/contrib/data/python/kernel_tests/reader_dataset_ops_test.py b/tensorflow/python/data/experimental/kernel_tests/reader_dataset_ops_test.py
similarity index 98%
rename from tensorflow/contrib/data/python/kernel_tests/reader_dataset_ops_test.py
rename to tensorflow/python/data/experimental/kernel_tests/reader_dataset_ops_test.py
index ed75b27..a02f4bd 100644
--- a/tensorflow/contrib/data/python/kernel_tests/reader_dataset_ops_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/reader_dataset_ops_test.py
@@ -23,8 +23,9 @@
import numpy as np
-from tensorflow.contrib.data.python.kernel_tests import reader_dataset_ops_test_base
-from tensorflow.contrib.data.python.ops import readers
+from tensorflow.python.data.experimental.kernel_tests import reader_dataset_ops_test_base
+from tensorflow.python.data.experimental.ops import readers
+from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import readers as core_readers
from tensorflow.python.data.util import nest
from tensorflow.python.framework import constant_op
@@ -242,7 +243,7 @@
self.assertEqual(32, shape[0])
-class MakeCsvDatasetTest(test.TestCase):
+class MakeCsvDatasetTest(test_base.DatasetTestBase):
def _make_csv_dataset(self, filenames, batch_size, num_epochs=1, **kwargs):
return readers.make_csv_dataset(
diff --git a/tensorflow/contrib/data/python/kernel_tests/reader_dataset_ops_test_base.py b/tensorflow/python/data/experimental/kernel_tests/reader_dataset_ops_test_base.py
similarity index 96%
rename from tensorflow/contrib/data/python/kernel_tests/reader_dataset_ops_test_base.py
rename to tensorflow/python/data/experimental/kernel_tests/reader_dataset_ops_test_base.py
index 08b9f03..b6ab80d 100644
--- a/tensorflow/contrib/data/python/kernel_tests/reader_dataset_ops_test_base.py
+++ b/tensorflow/python/data/experimental/kernel_tests/reader_dataset_ops_test_base.py
@@ -22,9 +22,10 @@
import os
import zlib
-from tensorflow.contrib.data.python.ops import readers
from tensorflow.core.example import example_pb2
from tensorflow.core.example import feature_pb2
+from tensorflow.python.data.experimental.ops import readers
+from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import iterator_ops
from tensorflow.python.data.ops import readers as core_readers
from tensorflow.python.framework import constant_op
@@ -32,11 +33,10 @@
from tensorflow.python.lib.io import python_io
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import parsing_ops
-from tensorflow.python.platform import test
from tensorflow.python.util import compat
-class FixedLengthRecordDatasetTestBase(test.TestCase):
+class FixedLengthRecordDatasetTestBase(test_base.DatasetTestBase):
"""Base class for setting up and testing FixedLengthRecordDataset."""
def setUp(self):
@@ -63,7 +63,7 @@
return filenames
-class ReadBatchFeaturesTestBase(test.TestCase):
+class ReadBatchFeaturesTestBase(test_base.DatasetTestBase):
"""Base class for setting up and testing `make_batched_feature_dataset`."""
def setUp(self):
@@ -273,7 +273,7 @@
self.assertAllEqual(expected_batch[i], actual_batch[i])
-class TextLineDatasetTestBase(test.TestCase):
+class TextLineDatasetTestBase(test_base.DatasetTestBase):
"""Base class for setting up and testing TextLineDataset."""
def _lineText(self, f, l):
@@ -313,7 +313,7 @@
return filenames
-class TFRecordDatasetTestBase(test.TestCase):
+class TFRecordDatasetTestBase(test_base.DatasetTestBase):
"""Base class for setting up and testing TFRecordDataset."""
def setUp(self):
diff --git a/tensorflow/contrib/data/python/kernel_tests/resample_test.py b/tensorflow/python/data/experimental/kernel_tests/resample_test.py
similarity index 96%
rename from tensorflow/contrib/data/python/kernel_tests/resample_test.py
rename to tensorflow/python/data/experimental/kernel_tests/resample_test.py
index 16b1441..775648c 100644
--- a/tensorflow/contrib/data/python/kernel_tests/resample_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/resample_test.py
@@ -23,7 +23,8 @@
import numpy as np
from six.moves import xrange # pylint: disable=redefined-builtin
-from tensorflow.contrib.data.python.ops import resampling
+from tensorflow.python.data.experimental.ops import resampling
+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
@@ -57,7 +58,7 @@
return end_time - start_time
-class ResampleTest(test.TestCase, parameterized.TestCase):
+class ResampleTest(test_base.DatasetTestBase, parameterized.TestCase):
@parameterized.named_parameters(
("InitialDistributionKnown", True),
diff --git a/tensorflow/contrib/data/python/kernel_tests/scan_dataset_op_test.py b/tensorflow/python/data/experimental/kernel_tests/scan_dataset_op_test.py
similarity index 97%
rename from tensorflow/contrib/data/python/kernel_tests/scan_dataset_op_test.py
rename to tensorflow/python/data/experimental/kernel_tests/scan_dataset_op_test.py
index dde678b..78ec80d 100644
--- a/tensorflow/contrib/data/python/kernel_tests/scan_dataset_op_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/scan_dataset_op_test.py
@@ -21,7 +21,8 @@
import numpy as np
-from tensorflow.contrib.data.python.ops import scan_ops
+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
@@ -33,7 +34,7 @@
from tensorflow.python.platform import test
-class ScanDatasetTest(test.TestCase):
+class ScanDatasetTest(test_base.DatasetTestBase):
def _counting_dataset(self, start, scan_fn):
return dataset_ops.Dataset.from_tensors(0).repeat().apply(
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/BUILD b/tensorflow/python/data/experimental/kernel_tests/serialization/BUILD
similarity index 90%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/BUILD
rename to tensorflow/python/data/experimental/kernel_tests/serialization/BUILD
index aa89674..20c02a5 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/BUILD
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/BUILD
@@ -13,7 +13,6 @@
],
srcs_version = "PY2AND3",
deps = [
- "//tensorflow/contrib/data/python/ops:iterator_ops",
"//tensorflow/python:client_testlib",
"//tensorflow/python:dtypes",
"//tensorflow/python:errors",
@@ -24,6 +23,7 @@
"//tensorflow/python:training",
"//tensorflow/python:util",
"//tensorflow/python:variables",
+ "//tensorflow/python/data/experimental/ops:iterator_ops",
"//tensorflow/python/data/ops:iterator_ops",
"//third_party/py/numpy",
],
@@ -37,10 +37,10 @@
tags = ["no_pip"],
deps = [
":dataset_serialization_test_base",
- "//tensorflow/contrib/data/python/ops:batching",
"//tensorflow/python:array_ops",
"//tensorflow/python:client_testlib",
"//tensorflow/python:sparse_tensor",
+ "//tensorflow/python/data/experimental/ops:batching",
"//tensorflow/python/data/ops:dataset_ops",
"//third_party/py/numpy",
],
@@ -81,9 +81,9 @@
tags = ["no_pip"],
deps = [
":dataset_serialization_test_base",
- "//tensorflow/contrib/data/python/ops:readers",
"//tensorflow/python:client_testlib",
"//tensorflow/python:framework_ops",
+ "//tensorflow/python/data/experimental/ops:readers",
],
)
@@ -126,8 +126,8 @@
tags = ["no_pip"],
deps = [
":dataset_serialization_test_base",
- "//tensorflow/contrib/data/python/kernel_tests:reader_dataset_ops_test_base",
"//tensorflow/python:client_testlib",
+ "//tensorflow/python/data/experimental/kernel_tests:reader_dataset_ops_test_base",
"//tensorflow/python/data/ops:readers",
],
)
@@ -160,8 +160,8 @@
srcs_version = "PY2AND3",
deps = [
":dataset_serialization_test_base",
- "//tensorflow/contrib/data/python/ops:grouping",
"//tensorflow/python:client_testlib",
+ "//tensorflow/python/data/experimental/ops:grouping",
"//tensorflow/python/data/ops:dataset_ops",
"//third_party/py/numpy",
],
@@ -174,8 +174,8 @@
srcs_version = "PY2AND3",
deps = [
":dataset_serialization_test_base",
- "//tensorflow/contrib/data/python/ops:grouping",
"//tensorflow/python:client_testlib",
+ "//tensorflow/python/data/experimental/ops:grouping",
"//tensorflow/python/data/ops:dataset_ops",
"//third_party/py/numpy",
],
@@ -189,9 +189,9 @@
tags = ["no_pip"],
deps = [
":dataset_serialization_test_base",
- "//tensorflow/contrib/data/python/ops:error_ops",
"//tensorflow/python:array_ops",
"//tensorflow/python:client_testlib",
+ "//tensorflow/python/data/experimental/ops:error_ops",
"//tensorflow/python/data/ops:dataset_ops",
"//third_party/py/numpy",
],
@@ -222,9 +222,9 @@
tags = ["no_pip"],
deps = [
":dataset_serialization_test_base",
- "//tensorflow/contrib/data/python/ops:batching",
"//tensorflow/python:client_testlib",
"//tensorflow/python:math_ops",
+ "//tensorflow/python/data/experimental/ops:batching",
"//tensorflow/python/data/ops:dataset_ops",
],
)
@@ -258,8 +258,8 @@
srcs_version = "PY2AND3",
deps = [
":dataset_serialization_test_base",
- "//tensorflow/contrib/data/python/ops:optimization",
"//tensorflow/python:client_testlib",
+ "//tensorflow/python/data/experimental/ops:optimization",
"//tensorflow/python/data/ops:dataset_ops",
],
)
@@ -288,10 +288,10 @@
tags = ["no_pip"],
deps = [
":dataset_serialization_test_base",
- "//tensorflow/contrib/data/python/ops:interleave_ops",
"//tensorflow/python:client_testlib",
"//tensorflow/python:sparse_ops",
"//tensorflow/python:sparse_tensor",
+ "//tensorflow/python/data/experimental/ops:interleave_ops",
"//tensorflow/python/data/ops:dataset_ops",
"//third_party/py/numpy",
],
@@ -326,8 +326,8 @@
tags = ["no_pip"],
deps = [
":dataset_serialization_test_base",
- "//tensorflow/contrib/data/python/kernel_tests:reader_dataset_ops_test_base",
"//tensorflow/python:client_testlib",
+ "//tensorflow/python/data/experimental/kernel_tests:reader_dataset_ops_test_base",
],
)
@@ -370,8 +370,8 @@
srcs_version = "PY2AND3",
deps = [
":dataset_serialization_test_base",
- "//tensorflow/contrib/data/python/ops:interleave_ops",
"//tensorflow/python:client_testlib",
+ "//tensorflow/python/data/experimental/ops:interleave_ops",
"//tensorflow/python/data/ops:dataset_ops",
],
)
@@ -384,8 +384,8 @@
tags = ["no_pip"],
deps = [
":dataset_serialization_test_base",
- "//tensorflow/contrib/data/python/ops:scan_ops",
"//tensorflow/python:client_testlib",
+ "//tensorflow/python/data/experimental/ops:scan_ops",
"//tensorflow/python/data/ops:dataset_ops",
],
)
@@ -411,10 +411,10 @@
srcs_version = "PY2AND3",
tags = ["no_pip"],
deps = [
- "//tensorflow/contrib/data/python/ops:iterator_ops",
"//tensorflow/python:client_testlib",
"//tensorflow/python:framework_ops",
"//tensorflow/python:training",
+ "//tensorflow/python/data/experimental/ops:iterator_ops",
"//tensorflow/python/data/ops:dataset_ops",
],
)
@@ -427,8 +427,8 @@
tags = ["no_pip"],
deps = [
":dataset_serialization_test_base",
- "//tensorflow/contrib/data/python/ops:shuffle_ops",
"//tensorflow/python:client_testlib",
+ "//tensorflow/python/data/experimental/ops:shuffle_ops",
"//tensorflow/python/data/ops:dataset_ops",
],
)
@@ -441,10 +441,10 @@
tags = ["no_pip"],
deps = [
":dataset_serialization_test_base",
- "//tensorflow/contrib/data/python/ops:iterator_ops",
"//tensorflow/python:client_testlib",
"//tensorflow/python:framework_ops",
"//tensorflow/python:training",
+ "//tensorflow/python/data/experimental/ops:iterator_ops",
"//tensorflow/python/data/ops:dataset_ops",
],
)
@@ -457,11 +457,11 @@
tags = ["no_pip"],
deps = [
":dataset_serialization_test_base",
- "//tensorflow/contrib/data/python/kernel_tests:sql_dataset_op_test_base",
- "//tensorflow/contrib/data/python/ops:readers",
"//tensorflow/python:array_ops",
"//tensorflow/python:client_testlib",
"//tensorflow/python:dtypes",
+ "//tensorflow/python/data/experimental/kernel_tests:sql_dataset_op_test_base",
+ "//tensorflow/python/data/experimental/ops:readers",
],
)
@@ -473,10 +473,10 @@
tags = ["no_pip"],
deps = [
":dataset_serialization_test_base",
- "//tensorflow/contrib/data/python/ops:stats_ops",
"//tensorflow/python:array_ops",
"//tensorflow/python:client_testlib",
"//tensorflow/python:framework_ops",
+ "//tensorflow/python/data/experimental/ops:stats_ops",
"//tensorflow/python/data/ops:dataset_ops",
],
)
@@ -490,8 +490,8 @@
tags = ["no_pip"],
deps = [
":dataset_serialization_test_base",
- "//tensorflow/contrib/data/python/kernel_tests:reader_dataset_ops_test_base",
"//tensorflow/python:client_testlib",
+ "//tensorflow/python/data/experimental/kernel_tests:reader_dataset_ops_test_base",
"//tensorflow/python/data/ops:readers",
],
)
@@ -505,8 +505,8 @@
tags = ["no_pip"],
deps = [
":dataset_serialization_test_base",
- "//tensorflow/contrib/data/python/kernel_tests:reader_dataset_ops_test_base",
"//tensorflow/python:client_testlib",
+ "//tensorflow/python/data/experimental/kernel_tests:reader_dataset_ops_test_base",
"//tensorflow/python/data/ops:readers",
],
)
@@ -519,8 +519,8 @@
tags = ["no_pip"],
deps = [
":dataset_serialization_test_base",
- "//tensorflow/contrib/data/python/ops:batching",
"//tensorflow/python:client_testlib",
+ "//tensorflow/python/data/experimental/ops:batching",
"//tensorflow/python/data/ops:dataset_ops",
"//third_party/py/numpy",
],
@@ -534,8 +534,8 @@
tags = ["no_pip"],
deps = [
":dataset_serialization_test_base",
- "//tensorflow/contrib/data/python/ops:unique",
"//tensorflow/python:client_testlib",
+ "//tensorflow/python/data/experimental/ops:unique",
"//tensorflow/python/data/ops:dataset_ops",
],
)
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/batch_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/batch_dataset_serialization_test.py
similarity index 94%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/batch_dataset_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/batch_dataset_serialization_test.py
index af87d8b..d72a6df 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/batch_dataset_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/batch_dataset_serialization_test.py
@@ -19,8 +19,8 @@
import numpy as np
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
-from tensorflow.contrib.data.python.ops import batching
+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 array_ops
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/cache_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/cache_dataset_serialization_test.py
similarity index 98%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/cache_dataset_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/cache_dataset_serialization_test.py
index 1b6059c..2bcf77f 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/cache_dataset_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/cache_dataset_serialization_test.py
@@ -21,7 +21,7 @@
from absl.testing import parameterized
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import errors
from tensorflow.python.platform import test
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/concatenate_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/concatenate_dataset_serialization_test.py
similarity index 94%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/concatenate_dataset_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/concatenate_dataset_serialization_test.py
index 96f13d7..c075dff 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/concatenate_dataset_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/concatenate_dataset_serialization_test.py
@@ -19,7 +19,7 @@
import numpy as np
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.platform import test
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/csv_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/csv_dataset_serialization_test.py
similarity index 93%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/csv_dataset_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/csv_dataset_serialization_test.py
index 247f204..d498349 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/csv_dataset_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/csv_dataset_serialization_test.py
@@ -20,8 +20,8 @@
import gzip
import os
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
-from tensorflow.contrib.data.python.ops import readers
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
+from tensorflow.python.data.experimental.ops import readers
from tensorflow.python.platform import test
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/dataset_constructor_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/dataset_constructor_serialization_test.py
similarity index 96%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/dataset_constructor_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/dataset_constructor_serialization_test.py
index 2139b5c..41a095f 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/dataset_constructor_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/dataset_constructor_serialization_test.py
@@ -19,7 +19,7 @@
import numpy as np
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.platform import test
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/dataset_serialization_test_base.py b/tensorflow/python/data/experimental/kernel_tests/serialization/dataset_serialization_test_base.py
similarity index 99%
copy from tensorflow/contrib/data/python/kernel_tests/serialization/dataset_serialization_test_base.py
copy to tensorflow/python/data/experimental/kernel_tests/serialization/dataset_serialization_test_base.py
index 595cece..7f435b8 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/dataset_serialization_test_base.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/dataset_serialization_test_base.py
@@ -22,7 +22,7 @@
import numpy as np
-from tensorflow.contrib.data.python.ops import iterator_ops as contrib_iterator_ops
+from tensorflow.python.data.experimental.ops import iterator_ops as contrib_iterator_ops
from tensorflow.python.data.ops import iterator_ops
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/filter_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/filter_dataset_serialization_test.py
similarity index 95%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/filter_dataset_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/filter_dataset_serialization_test.py
index 7c17007..225f6cb 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/filter_dataset_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/filter_dataset_serialization_test.py
@@ -19,7 +19,7 @@
import numpy as np
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.ops import math_ops
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/fixed_length_record_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/fixed_length_record_dataset_serialization_test.py
similarity index 89%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/fixed_length_record_dataset_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/fixed_length_record_dataset_serialization_test.py
index 34392d8..70caf3e 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/fixed_length_record_dataset_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/fixed_length_record_dataset_serialization_test.py
@@ -17,8 +17,8 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.contrib.data.python.kernel_tests import reader_dataset_ops_test_base
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
+from tensorflow.python.data.experimental.kernel_tests import reader_dataset_ops_test_base
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
from tensorflow.python.data.ops import readers as core_readers
from tensorflow.python.platform import test
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/flat_map_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/flat_map_dataset_serialization_test.py
similarity index 96%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/flat_map_dataset_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/flat_map_dataset_serialization_test.py
index 16051ff..c30534a 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/flat_map_dataset_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/flat_map_dataset_serialization_test.py
@@ -17,7 +17,7 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/group_by_reducer_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/group_by_reducer_serialization_test.py
similarity index 93%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/group_by_reducer_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/group_by_reducer_serialization_test.py
index 571e089..169c884 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/group_by_reducer_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/group_by_reducer_serialization_test.py
@@ -19,8 +19,8 @@
import numpy as np
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
-from tensorflow.contrib.data.python.ops import grouping
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
+from tensorflow.python.data.experimental.ops import grouping
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.platform import test
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/group_by_window_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/group_by_window_serialization_test.py
similarity index 92%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/group_by_window_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/group_by_window_serialization_test.py
index f86af40..e5bc762 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/group_by_window_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/group_by_window_serialization_test.py
@@ -19,8 +19,8 @@
import numpy as np
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
-from tensorflow.contrib.data.python.ops import grouping
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
+from tensorflow.python.data.experimental.ops import grouping
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.platform import test
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/ignore_errors_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/ignore_errors_serialization_test.py
similarity index 90%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/ignore_errors_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/ignore_errors_serialization_test.py
index 65ae992..df1f431 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/ignore_errors_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/ignore_errors_serialization_test.py
@@ -19,8 +19,8 @@
import numpy as np
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
-from tensorflow.contrib.data.python.ops import error_ops
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
+from tensorflow.python.data.experimental.ops import error_ops
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.ops import array_ops
from tensorflow.python.platform import test
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/interleave_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/interleave_dataset_serialization_test.py
similarity index 96%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/interleave_dataset_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/interleave_dataset_serialization_test.py
index 243f640..0c1d40c 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/interleave_dataset_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/interleave_dataset_serialization_test.py
@@ -20,7 +20,7 @@
from absl.testing import parameterized
import numpy as np
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.ops import sparse_ops
diff --git a/tensorflow/contrib/data/python/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
similarity index 94%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/map_and_batch_dataset_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/map_and_batch_dataset_serialization_test.py
index c9cd211..166ffa9 100644
--- a/tensorflow/contrib/data/python/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
@@ -19,8 +19,8 @@
import math
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
-from tensorflow.contrib.data.python.ops import batching
+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.ops import math_ops
from tensorflow.python.platform import test
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/map_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/map_dataset_serialization_test.py
similarity index 97%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/map_dataset_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/map_dataset_serialization_test.py
index ab783e5..b93156a 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/map_dataset_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/map_dataset_serialization_test.py
@@ -19,7 +19,7 @@
import numpy as np
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/optimize_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/optimize_dataset_serialization_test.py
similarity index 88%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/optimize_dataset_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/optimize_dataset_serialization_test.py
index d5c0349..ed4a1da 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/optimize_dataset_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/optimize_dataset_serialization_test.py
@@ -17,8 +17,8 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
-from tensorflow.contrib.data.python.ops import optimization
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
+from tensorflow.python.data.experimental.ops import optimization
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.platform import test
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/padded_batch_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/padded_batch_dataset_serialization_test.py
similarity index 95%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/padded_batch_dataset_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/padded_batch_dataset_serialization_test.py
index 9ac42a4..6f72b24 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/padded_batch_dataset_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/padded_batch_dataset_serialization_test.py
@@ -19,7 +19,7 @@
import numpy as np
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import string_ops
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/parallel_interleave_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/parallel_interleave_dataset_serialization_test.py
similarity index 95%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/parallel_interleave_dataset_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/parallel_interleave_dataset_serialization_test.py
index 1f8a584..b8f38e8 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/parallel_interleave_dataset_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/parallel_interleave_dataset_serialization_test.py
@@ -19,8 +19,8 @@
import numpy as np
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
-from tensorflow.contrib.data.python.ops import interleave_ops
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
+from tensorflow.python.data.experimental.ops import interleave_ops
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.ops import sparse_ops
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/parallel_map_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/parallel_map_dataset_serialization_test.py
similarity index 95%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/parallel_map_dataset_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/parallel_map_dataset_serialization_test.py
index 3fb7605..a0bdd4f 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/parallel_map_dataset_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/parallel_map_dataset_serialization_test.py
@@ -19,7 +19,7 @@
import numpy as np
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
@@ -65,7 +65,7 @@
for ds_fn in [self._build_ds, self._build_ds_with_prefetch]:
self.run_core_tests(
ds_fn,
- lambda: ds_fn(multiplier=15.0),
+ lambda: ds_fn(multiplier=15.0), # pylint: disable=cell-var-from-loop
self._num_outputs)
def testSaveStatefulFunction(self):
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/parse_example_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/parse_example_dataset_serialization_test.py
similarity index 89%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/parse_example_dataset_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/parse_example_dataset_serialization_test.py
index d3fa84e..a0dd696 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/parse_example_dataset_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/parse_example_dataset_serialization_test.py
@@ -17,8 +17,8 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.contrib.data.python.kernel_tests import reader_dataset_ops_test_base
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
+from tensorflow.python.data.experimental.kernel_tests import reader_dataset_ops_test_base
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
from tensorflow.python.platform import test
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/prefetch_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/prefetch_dataset_serialization_test.py
similarity index 92%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/prefetch_dataset_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/prefetch_dataset_serialization_test.py
index c802402..00d74c0 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/prefetch_dataset_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/prefetch_dataset_serialization_test.py
@@ -17,7 +17,7 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.platform import test
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/range_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/range_dataset_serialization_test.py
similarity index 97%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/range_dataset_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/range_dataset_serialization_test.py
index 6341190..ef99d01 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/range_dataset_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/range_dataset_serialization_test.py
@@ -19,7 +19,7 @@
import os
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/sample_from_datasets_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/sample_from_datasets_serialization_test.py
similarity index 89%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/sample_from_datasets_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/sample_from_datasets_serialization_test.py
index fdb35ea..c23c1ec 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/sample_from_datasets_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/sample_from_datasets_serialization_test.py
@@ -17,8 +17,8 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
-from tensorflow.contrib.data.python.ops import interleave_ops
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
+from tensorflow.python.data.experimental.ops import interleave_ops
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.platform import test
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/scan_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/scan_dataset_serialization_test.py
similarity index 89%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/scan_dataset_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/scan_dataset_serialization_test.py
index af9ef48..5f50160 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/scan_dataset_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/scan_dataset_serialization_test.py
@@ -17,8 +17,8 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
-from tensorflow.contrib.data.python.ops import scan_ops
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
+from tensorflow.python.data.experimental.ops import scan_ops
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.platform import test
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/sequence_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/sequence_dataset_serialization_test.py
similarity index 97%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/sequence_dataset_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/sequence_dataset_serialization_test.py
index 2afebca..fe99a3d 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/sequence_dataset_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/sequence_dataset_serialization_test.py
@@ -19,7 +19,7 @@
import numpy as np
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.platform import test
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/serialization_integration_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/serialization_integration_test.py
similarity index 96%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/serialization_integration_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/serialization_integration_test.py
index 6aac50e..88d5c89 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/serialization_integration_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/serialization_integration_test.py
@@ -19,7 +19,7 @@
import os
-from tensorflow.contrib.data.python.ops import iterator_ops as contrib_iterator_ops
+from tensorflow.python.data.experimental.ops import iterator_ops as contrib_iterator_ops
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import ops
from tensorflow.python.platform import test
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/shuffle_and_repeat_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/shuffle_and_repeat_dataset_serialization_test.py
similarity index 88%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/shuffle_and_repeat_dataset_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/shuffle_and_repeat_dataset_serialization_test.py
index f199ec8..f847ac1 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/shuffle_and_repeat_dataset_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/shuffle_and_repeat_dataset_serialization_test.py
@@ -17,8 +17,8 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
-from tensorflow.contrib.data.python.ops import shuffle_ops
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
+from tensorflow.python.data.experimental.ops import shuffle_ops
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.platform import test
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/shuffle_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/shuffle_dataset_serialization_test.py
similarity index 96%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/shuffle_dataset_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/shuffle_dataset_serialization_test.py
index a59fa94..a04f1dd 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/shuffle_dataset_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/shuffle_dataset_serialization_test.py
@@ -17,8 +17,8 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
-from tensorflow.contrib.data.python.ops import iterator_ops as contrib_iterator_ops
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
+from tensorflow.python.data.experimental.ops import iterator_ops as contrib_iterator_ops
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import ops
from tensorflow.python.platform import test
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/sql_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/sql_dataset_serialization_test.py
similarity index 88%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/sql_dataset_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/sql_dataset_serialization_test.py
index 93b26ed..b179770 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/sql_dataset_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/sql_dataset_serialization_test.py
@@ -19,9 +19,9 @@
import os
-from tensorflow.contrib.data.python.kernel_tests import sql_dataset_op_test_base
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
-from tensorflow.contrib.data.python.ops import readers
+from tensorflow.python.data.experimental.kernel_tests import sql_dataset_op_test_base
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
+from tensorflow.python.data.experimental.ops import readers
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.platform import test
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/stats_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/stats_dataset_serialization_test.py
similarity index 84%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/stats_dataset_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/stats_dataset_serialization_test.py
index 14cd3e9..ef7061b 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/stats_dataset_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/stats_dataset_serialization_test.py
@@ -17,9 +17,10 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
-from tensorflow.contrib.data.python.ops import stats_ops
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
+from tensorflow.python.data.experimental.ops import stats_ops
from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.platform import test
@@ -90,6 +91,16 @@
lambda: self._build_dataset_multiple_tags(num_outputs, tag1, tag2),
None, num_outputs)
+ def _build_dataset_stats_aggregator(self):
+ stats_aggregator = stats_ops.StatsAggregator()
+ return dataset_ops.Dataset.range(10).apply(
+ stats_ops.set_stats_aggregator(stats_aggregator))
+
+ def test_set_stats_aggregator_not_support_checkpointing(self):
+ with self.assertRaisesRegexp(errors.UnimplementedError,
+ "does not support checkpointing"):
+ self.run_core_tests(self._build_dataset_stats_aggregator, None, 10)
+
if __name__ == "__main__":
test.main()
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/textline_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/textline_dataset_serialization_test.py
similarity index 90%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/textline_dataset_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/textline_dataset_serialization_test.py
index 2483787..c87a744 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/textline_dataset_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/textline_dataset_serialization_test.py
@@ -17,8 +17,8 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.contrib.data.python.kernel_tests import reader_dataset_ops_test_base
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
+from tensorflow.python.data.experimental.kernel_tests import reader_dataset_ops_test_base
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
from tensorflow.python.data.ops import readers as core_readers
from tensorflow.python.platform import test
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/tf_record_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/tf_record_dataset_serialization_test.py
similarity index 94%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/tf_record_dataset_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/tf_record_dataset_serialization_test.py
index 55a6257..f0dcc13 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/tf_record_dataset_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/tf_record_dataset_serialization_test.py
@@ -21,8 +21,8 @@
import os
import zlib
-from tensorflow.contrib.data.python.kernel_tests import reader_dataset_ops_test_base
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
+from tensorflow.python.data.experimental.kernel_tests import reader_dataset_ops_test_base
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
from tensorflow.python.data.ops import readers as core_readers
from tensorflow.python.platform import test
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/unbatch_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/unbatch_dataset_serialization_test.py
similarity index 91%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/unbatch_dataset_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/unbatch_dataset_serialization_test.py
index b2a5a8a..528598d 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/unbatch_dataset_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/unbatch_dataset_serialization_test.py
@@ -19,8 +19,8 @@
import numpy as np
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
-from tensorflow.contrib.data.python.ops import batching
+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.platform import test
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/unique_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/unique_dataset_serialization_test.py
similarity index 89%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/unique_dataset_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/unique_dataset_serialization_test.py
index 22f15b8..e2862af 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/unique_dataset_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/unique_dataset_serialization_test.py
@@ -17,8 +17,8 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
-from tensorflow.contrib.data.python.ops import unique
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
+from tensorflow.python.data.experimental.ops import unique
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.platform import test
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/zip_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/zip_dataset_serialization_test.py
similarity index 94%
rename from tensorflow/contrib/data/python/kernel_tests/serialization/zip_dataset_serialization_test.py
rename to tensorflow/python/data/experimental/kernel_tests/serialization/zip_dataset_serialization_test.py
index 340a6ff..4ea6131 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/zip_dataset_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/zip_dataset_serialization_test.py
@@ -19,7 +19,7 @@
import numpy as np
-from tensorflow.contrib.data.python.kernel_tests.serialization import dataset_serialization_test_base
+from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.platform import test
diff --git a/tensorflow/contrib/data/python/kernel_tests/serialization/serialization_integration_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization_integration_test.py
similarity index 96%
copy from tensorflow/contrib/data/python/kernel_tests/serialization/serialization_integration_test.py
copy to tensorflow/python/data/experimental/kernel_tests/serialization_integration_test.py
index 6aac50e..88d5c89 100644
--- a/tensorflow/contrib/data/python/kernel_tests/serialization/serialization_integration_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization_integration_test.py
@@ -19,7 +19,7 @@
import os
-from tensorflow.contrib.data.python.ops import iterator_ops as contrib_iterator_ops
+from tensorflow.python.data.experimental.ops import iterator_ops as contrib_iterator_ops
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import ops
from tensorflow.python.platform import test
diff --git a/tensorflow/contrib/data/python/kernel_tests/shuffle_dataset_op_test.py b/tensorflow/python/data/experimental/kernel_tests/shuffle_dataset_op_test.py
similarity index 95%
rename from tensorflow/contrib/data/python/kernel_tests/shuffle_dataset_op_test.py
rename to tensorflow/python/data/experimental/kernel_tests/shuffle_dataset_op_test.py
index 440e48d..50895b5 100644
--- a/tensorflow/contrib/data/python/kernel_tests/shuffle_dataset_op_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/shuffle_dataset_op_test.py
@@ -19,14 +19,15 @@
import numpy as np
-from tensorflow.contrib.data.python.ops import shuffle_ops
+from tensorflow.python.data.experimental.ops import shuffle_ops
+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.platform import test
-class ShuffleAndRepeatTest(test.TestCase):
+class ShuffleAndRepeatTest(test_base.DatasetTestBase):
def _build_ds(self, seed, count=5, num_elements=20):
return dataset_ops.Dataset.range(num_elements).apply(
diff --git a/tensorflow/contrib/data/python/kernel_tests/sql_dataset_op_test.py b/tensorflow/python/data/experimental/kernel_tests/sql_dataset_op_test.py
similarity index 99%
rename from tensorflow/contrib/data/python/kernel_tests/sql_dataset_op_test.py
rename to tensorflow/python/data/experimental/kernel_tests/sql_dataset_op_test.py
index 52823d3..301f754 100644
--- a/tensorflow/contrib/data/python/kernel_tests/sql_dataset_op_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/sql_dataset_op_test.py
@@ -18,7 +18,7 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.contrib.data.python.kernel_tests import sql_dataset_op_test_base
+from tensorflow.python.data.experimental.kernel_tests import sql_dataset_op_test_base
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.platform import test
diff --git a/tensorflow/contrib/data/python/kernel_tests/sql_dataset_op_test_base.py b/tensorflow/python/data/experimental/kernel_tests/sql_dataset_op_test_base.py
similarity index 95%
rename from tensorflow/contrib/data/python/kernel_tests/sql_dataset_op_test_base.py
rename to tensorflow/python/data/experimental/kernel_tests/sql_dataset_op_test_base.py
index 1f5c725..a135c35 100644
--- a/tensorflow/contrib/data/python/kernel_tests/sql_dataset_op_test_base.py
+++ b/tensorflow/python/data/experimental/kernel_tests/sql_dataset_op_test_base.py
@@ -23,13 +23,14 @@
import sqlite3
-from tensorflow.contrib.data.python.ops import readers
+from tensorflow.python.data.experimental.ops import readers
+from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import array_ops
from tensorflow.python.platform import test
-class SqlDatasetTestBase(test.TestCase):
+class SqlDatasetTestBase(test_base.DatasetTestBase):
"""Base class for setting up and testing SqlDataset."""
def _createSqlDataset(self, output_types, num_repeats=1):
@@ -92,5 +93,3 @@
9007199254740992.0)])
conn.commit()
conn.close()
-
-
diff --git a/tensorflow/contrib/data/python/kernel_tests/stats_dataset_ops_test.py b/tensorflow/python/data/experimental/kernel_tests/stats_dataset_ops_test.py
similarity index 98%
rename from tensorflow/contrib/data/python/kernel_tests/stats_dataset_ops_test.py
rename to tensorflow/python/data/experimental/kernel_tests/stats_dataset_ops_test.py
index be8ae5e..6761fbd 100644
--- a/tensorflow/contrib/data/python/kernel_tests/stats_dataset_ops_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/stats_dataset_ops_test.py
@@ -19,8 +19,8 @@
import numpy as np
-from tensorflow.contrib.data.python.kernel_tests import stats_dataset_test_base
-from tensorflow.contrib.data.python.ops import stats_ops
+from tensorflow.python.data.experimental.kernel_tests import stats_dataset_test_base
+from tensorflow.python.data.experimental.ops import stats_ops
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
diff --git a/tensorflow/contrib/data/python/kernel_tests/stats_dataset_test_base.py b/tensorflow/python/data/experimental/kernel_tests/stats_dataset_test_base.py
similarity index 95%
rename from tensorflow/contrib/data/python/kernel_tests/stats_dataset_test_base.py
rename to tensorflow/python/data/experimental/kernel_tests/stats_dataset_test_base.py
index b1b4c23..80f2625 100644
--- a/tensorflow/contrib/data/python/kernel_tests/stats_dataset_test_base.py
+++ b/tensorflow/python/data/experimental/kernel_tests/stats_dataset_test_base.py
@@ -19,10 +19,10 @@
from tensorflow.core.framework import summary_pb2
-from tensorflow.python.platform import test
+from tensorflow.python.data.kernel_tests import test_base
-class StatsDatasetTestBase(test.TestCase):
+class StatsDatasetTestBase(test_base.DatasetTestBase):
"""Base class for testing statistics gathered in `StatsAggregator`."""
def _assertSummaryContains(self, summary_str, tag):
diff --git a/tensorflow/contrib/data/python/kernel_tests/threadpool_dataset_ops_test.py b/tensorflow/python/data/experimental/kernel_tests/threadpool_dataset_ops_test.py
similarity index 90%
rename from tensorflow/contrib/data/python/kernel_tests/threadpool_dataset_ops_test.py
rename to tensorflow/python/data/experimental/kernel_tests/threadpool_dataset_ops_test.py
index 8d335e8..4432dcb 100644
--- a/tensorflow/contrib/data/python/kernel_tests/threadpool_dataset_ops_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/threadpool_dataset_ops_test.py
@@ -22,8 +22,9 @@
from absl.testing import parameterized
import numpy as np
-from tensorflow.contrib.data.python.ops import threadpool
-from tensorflow.contrib.data.python.ops import unique
+from tensorflow.python.data.experimental.ops import threadpool
+from tensorflow.python.data.experimental.ops import unique
+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
@@ -31,7 +32,8 @@
from tensorflow.python.platform import test
-class OverrideThreadpoolDatasetTest(test.TestCase, parameterized.TestCase):
+class OverrideThreadpoolDatasetTest(test_base.DatasetTestBase,
+ parameterized.TestCase):
@parameterized.named_parameters(
("1", 1, None),
diff --git a/tensorflow/contrib/data/python/kernel_tests/unique_dataset_op_test.py b/tensorflow/python/data/experimental/kernel_tests/unique_dataset_op_test.py
similarity index 94%
rename from tensorflow/contrib/data/python/kernel_tests/unique_dataset_op_test.py
rename to tensorflow/python/data/experimental/kernel_tests/unique_dataset_op_test.py
index f994c85..b5a0b20 100644
--- a/tensorflow/contrib/data/python/kernel_tests/unique_dataset_op_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/unique_dataset_op_test.py
@@ -17,7 +17,8 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.contrib.data.python.ops import unique
+from tensorflow.python.data.experimental.ops import unique
+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
@@ -25,7 +26,7 @@
from tensorflow.python.util import compat
-class UniqueDatasetTest(test.TestCase):
+class UniqueDatasetTest(test_base.DatasetTestBase):
def _testSimpleHelper(self, dtype, test_cases):
"""Test the `unique()` transformation on a list of test cases.
diff --git a/tensorflow/contrib/data/python/kernel_tests/writer_ops_test.py b/tensorflow/python/data/experimental/kernel_tests/writer_ops_test.py
similarity index 95%
rename from tensorflow/contrib/data/python/kernel_tests/writer_ops_test.py
rename to tensorflow/python/data/experimental/kernel_tests/writer_ops_test.py
index 867ee2b..25a2e63 100644
--- a/tensorflow/contrib/data/python/kernel_tests/writer_ops_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/writer_ops_test.py
@@ -19,7 +19,8 @@
import os
-from tensorflow.contrib.data.python.ops import writers
+from tensorflow.python.data.experimental.ops import writers
+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
@@ -30,7 +31,7 @@
from tensorflow.python.util import compat
-class TFRecordWriterTest(test.TestCase):
+class TFRecordWriterTest(test_base.DatasetTestBase):
def setUp(self):
super(TFRecordWriterTest, self).setUp()
diff --git a/tensorflow/python/data/experimental/ops/BUILD b/tensorflow/python/data/experimental/ops/BUILD
new file mode 100644
index 0000000..915d399
--- /dev/null
+++ b/tensorflow/python/data/experimental/ops/BUILD
@@ -0,0 +1,377 @@
+package(default_visibility = ["//tensorflow:internal"])
+
+licenses(["notice"]) # Apache 2.0
+
+exports_files(["LICENSE"])
+
+load(
+ "//tensorflow:tensorflow.bzl",
+ "tf_gen_op_wrapper_py",
+ "tf_kernel_library",
+)
+load("//tensorflow:tensorflow.bzl", "tf_custom_op_py_library")
+
+py_library(
+ name = "counter",
+ srcs = ["counter.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ ":scan_ops",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python/data/ops:dataset_ops",
+ ],
+)
+
+py_library(
+ name = "get_single_element",
+ srcs = ["get_single_element.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ "//tensorflow/python:dataset_ops_gen",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/util:nest",
+ "//tensorflow/python/data/util:sparse",
+ ],
+)
+
+py_library(
+ name = "iterator_ops",
+ srcs = [
+ "iterator_ops.py",
+ ],
+ srcs_version = "PY2AND3",
+ deps = [
+ "//tensorflow/python:basic_session_run_hooks",
+ "//tensorflow/python:checkpoint_management",
+ "//tensorflow/python:dataset_ops_gen",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python:saver",
+ "//tensorflow/python:session_run_hook",
+ "//tensorflow/python/data/ops:iterator_ops",
+ "//tensorflow/python/data/ops:optional_ops",
+ ],
+)
+
+py_library(
+ name = "random_ops",
+ srcs = [
+ "random_ops.py",
+ ],
+ srcs_version = "PY2AND3",
+ deps = [
+ "//tensorflow/python:constant_op",
+ "//tensorflow/python:dataset_ops_gen",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python:random_seed",
+ "//tensorflow/python:tensor_shape",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/util:nest",
+ "//tensorflow/python/data/util:sparse",
+ ],
+)
+
+py_library(
+ name = "readers",
+ srcs = [
+ "readers.py",
+ ],
+ srcs_version = "PY2AND3",
+ deps = [
+ ":batching",
+ ":interleave_ops",
+ ":optimization",
+ ":parsing_ops",
+ ":shuffle_ops",
+ "//tensorflow/python:constant_op",
+ "//tensorflow/python:dataset_ops_gen",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:experimental_dataset_ops_gen",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python:lib",
+ "//tensorflow/python:platform",
+ "//tensorflow/python:tensor_shape",
+ "//tensorflow/python:util",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/ops:readers",
+ "//tensorflow/python/data/util:convert",
+ "//tensorflow/python/data/util:nest",
+ "//third_party/py/numpy",
+ ],
+)
+
+py_library(
+ name = "shuffle_ops",
+ srcs = [
+ "shuffle_ops.py",
+ ],
+ srcs_version = "PY2AND3",
+ deps = [
+ "//tensorflow/python/data/ops:dataset_ops",
+ ],
+)
+
+py_library(
+ name = "batching",
+ srcs = ["batching.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ ":get_single_element",
+ ":grouping",
+ "//tensorflow/python:array_ops",
+ "//tensorflow/python:dataset_ops_gen",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python:math_ops",
+ "//tensorflow/python:tensor_shape",
+ "//tensorflow/python:tensor_util",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/util:convert",
+ "//tensorflow/python/data/util:nest",
+ "//tensorflow/python/data/util:sparse",
+ "//third_party/py/numpy",
+ ],
+)
+
+py_library(
+ name = "enumerate_ops",
+ srcs = ["enumerate_ops.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python/data/ops:dataset_ops",
+ ],
+)
+
+py_library(
+ name = "error_ops",
+ srcs = ["error_ops.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ "//tensorflow/python:experimental_dataset_ops_gen",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/util:nest",
+ "//tensorflow/python/data/util:sparse",
+ ],
+)
+
+py_library(
+ name = "grouping",
+ srcs = ["grouping.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ "//tensorflow/python:array_ops",
+ "//tensorflow/python:check_ops",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python:function",
+ "//tensorflow/python:math_ops",
+ "//tensorflow/python:tensor_shape",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/util:nest",
+ "//tensorflow/python/data/util:sparse",
+ ],
+)
+
+py_library(
+ name = "interleave_ops",
+ srcs = ["interleave_ops.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ ":random_ops",
+ "//tensorflow/python:array_ops",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:experimental_dataset_ops_gen",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python:math_ops",
+ "//tensorflow/python:stateless_random_ops_gen",
+ "//tensorflow/python:util",
+ "//tensorflow/python/data/ops:readers",
+ "//tensorflow/python/data/util:nest",
+ "//tensorflow/python/data/util:sparse",
+ ],
+)
+
+py_library(
+ name = "optimization",
+ srcs = ["optimization.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:experimental_dataset_ops_gen",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python/data/util:nest",
+ "//tensorflow/python/data/util:sparse",
+ ],
+)
+
+py_library(
+ name = "parsing_ops",
+ srcs = ["parsing_ops.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ "//tensorflow/python:dataset_ops_gen",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python:parsing_ops",
+ "//tensorflow/python:sparse_tensor",
+ "//tensorflow/python:tensor_shape",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/util:nest",
+ ],
+)
+
+py_library(
+ name = "map_defun",
+ srcs = ["map_defun.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ "//tensorflow/python:dataset_ops_gen",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python:tensor_shape",
+ ],
+)
+
+py_library(
+ name = "resampling",
+ srcs = ["resampling.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ ":batching",
+ ":interleave_ops",
+ ":scan_ops",
+ "//tensorflow/python:array_ops",
+ "//tensorflow/python:control_flow_ops",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python:logging_ops",
+ "//tensorflow/python:math_ops",
+ "//tensorflow/python:random_ops",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//third_party/py/numpy",
+ ],
+)
+
+py_library(
+ name = "scan_ops",
+ srcs = ["scan_ops.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ "//tensorflow/python:dataset_ops_gen",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python:function",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/util:nest",
+ "//tensorflow/python/data/util:sparse",
+ ],
+)
+
+py_library(
+ name = "stats_ops",
+ srcs = ["stats_ops.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ "//tensorflow/python:dataset_ops_gen",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/ops:iterator_ops",
+ "//tensorflow/python/data/util:nest",
+ "//tensorflow/python/data/util:sparse",
+ ],
+)
+
+py_library(
+ name = "threadpool",
+ srcs = ["threadpool.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ "//tensorflow/python:experimental_dataset_ops_gen",
+ "//tensorflow/python:resource_variable_ops",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/util:nest",
+ "//tensorflow/python/data/util:sparse",
+ "//tensorflow/python/eager:context",
+ ],
+)
+
+py_library(
+ name = "unique",
+ srcs = [
+ "unique.py",
+ ],
+ srcs_version = "PY2AND3",
+ deps = [
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:experimental_dataset_ops_gen",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/util:nest",
+ "//tensorflow/python/data/util:sparse",
+ ],
+)
+
+py_library(
+ name = "writers",
+ srcs = [
+ "writers.py",
+ ],
+ srcs_version = "PY2AND3",
+ deps = [
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python/data/ops:dataset_ops",
+ ],
+)
+
+py_library(
+ name = "indexed_dataset_ops",
+ srcs = ["indexed_dataset_ops.py"],
+ deps = [
+ "//tensorflow/python:experimental_dataset_ops_gen",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/util:nest",
+ "//tensorflow/python/data/util:sparse",
+ ],
+)
+
+py_library(
+ name = "prefetching_ops",
+ srcs = ["prefetching_ops.py"],
+ deps = [
+ "//tensorflow/python:experimental_dataset_ops_gen",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/util:nest",
+ "//tensorflow/python/data/util:sparse",
+ ],
+)
+
+py_library(
+ name = "dataset_ops",
+ deps = [
+ ":batching",
+ ":counter",
+ ":enumerate_ops",
+ ":error_ops",
+ ":get_single_element",
+ ":grouping",
+ ":indexed_dataset_ops",
+ ":interleave_ops",
+ ":map_defun",
+ ":optimization",
+ ":prefetching_ops",
+ ":readers",
+ ":resampling",
+ ":scan_ops",
+ ":shuffle_ops",
+ ":stats_ops",
+ ":threadpool",
+ ":unique",
+ ":writers",
+ "//tensorflow/python:dataset_ops_gen",
+ "//tensorflow/python:util",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/util:nest",
+ ],
+)
diff --git a/tensorflow/python/data/experimental/ops/batching.py b/tensorflow/python/data/experimental/ops/batching.py
new file mode 100644
index 0000000..d42af9e
--- /dev/null
+++ b/tensorflow/python/data/experimental/ops/batching.py
@@ -0,0 +1,669 @@
+# 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.
+# ==============================================================================
+"""Batching dataset transformations."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from tensorflow.python.data.experimental.ops import get_single_element
+from tensorflow.python.data.experimental.ops import grouping
+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.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 check_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import gen_array_ops
+from tensorflow.python.ops import gen_dataset_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import sparse_ops
+from tensorflow.python.util.tf_export import tf_export
+
+
+def batch_window(dataset):
+ """Batches a window of tensors.
+
+ Args:
+ dataset: the input dataset.
+
+ Returns:
+ A `Tensor` representing the batch of the entire input dataset.
+ """
+ if isinstance(dataset.output_classes, tuple):
+ raise TypeError("Input dataset expected to have a single component")
+ if dataset.output_classes is ops.Tensor:
+ return _batch_dense_window(dataset)
+ elif dataset.output_classes is sparse_tensor.SparseTensor:
+ return _batch_sparse_window(dataset)
+ else:
+ raise TypeError("Unsupported dataset type: %s" % dataset.output_classes)
+
+
+def _batch_dense_window(dataset):
+ """Batches a window of dense tensors."""
+
+ def key_fn(_):
+ return np.int64(0)
+
+ def shape_init_fn(_):
+ return array_ops.shape(first_element)
+
+ def shape_reduce_fn(state, value):
+ check_ops.assert_equal(state, array_ops.shape(value))
+ return state
+
+ def finalize_fn(state):
+ return state
+
+ if dataset.output_shapes.is_fully_defined():
+ shape = dataset.output_shapes
+ else:
+ first_element = get_single_element.get_single_element(dataset.take(1))
+ shape_reducer = grouping.Reducer(shape_init_fn, shape_reduce_fn,
+ finalize_fn)
+ shape = get_single_element.get_single_element(
+ dataset.apply(grouping.group_by_reducer(key_fn, shape_reducer)))
+
+ def batch_init_fn(_):
+ batch_shape = array_ops.concat([[0], shape], 0)
+ return gen_array_ops.empty(batch_shape, dtype=dataset.output_types)
+
+ def batch_reduce_fn(state, value):
+ return array_ops.concat([state, [value]], 0)
+
+ batch_reducer = grouping.Reducer(batch_init_fn, batch_reduce_fn, finalize_fn)
+ return get_single_element.get_single_element(
+ dataset.apply(grouping.group_by_reducer(key_fn, batch_reducer)))
+
+
+def _batch_sparse_window(dataset):
+ """Batches a window of sparse tensors."""
+
+ def key_fn(_):
+ return np.int64(0)
+
+ def shape_init_fn(_):
+ return first_element.dense_shape
+
+ def shape_reduce_fn(state, value):
+ check_ops.assert_equal(state, value.dense_shape)
+ return state
+
+ def finalize_fn(state):
+ return state
+
+ if dataset.output_shapes.is_fully_defined():
+ shape = dataset.output_shapes
+ else:
+ first_element = get_single_element.get_single_element(dataset.take(1))
+ shape_reducer = grouping.Reducer(shape_init_fn, shape_reduce_fn,
+ finalize_fn)
+ shape = get_single_element.get_single_element(
+ dataset.apply(grouping.group_by_reducer(key_fn, shape_reducer)))
+
+ def batch_init_fn(_):
+ indices_shape = array_ops.concat([[0], [array_ops.size(shape) + 1]], 0)
+ return sparse_tensor.SparseTensor(
+ indices=gen_array_ops.empty(indices_shape, dtype=dtypes.int64),
+ values=constant_op.constant([], shape=[0], dtype=dataset.output_types),
+ dense_shape=array_ops.concat(
+ [np.array([0], dtype=np.int64),
+ math_ops.cast(shape, dtypes.int64)], 0))
+
+ def batch_reduce_fn(state, value):
+ return sparse_ops.sparse_concat(0, [state, value])
+
+ def reshape_fn(value):
+ return sparse_ops.sparse_reshape(
+ value,
+ array_ops.concat([np.array([1], dtype=np.int64), value.dense_shape], 0))
+
+ batch_reducer = grouping.Reducer(batch_init_fn, batch_reduce_fn, finalize_fn)
+ return get_single_element.get_single_element(
+ dataset.map(reshape_fn).apply(
+ grouping.group_by_reducer(key_fn, batch_reducer)))
+
+
+@tf_export("data.experimental.dense_to_sparse_batch")
+def dense_to_sparse_batch(batch_size, row_shape):
+ """A transformation that batches ragged elements into `tf.SparseTensor`s.
+
+ Like `Dataset.padded_batch()`, this transformation combines multiple
+ consecutive elements of the dataset, which might have different
+ shapes, into a single element. The resulting element has three
+ components (`indices`, `values`, and `dense_shape`), which
+ comprise a `tf.SparseTensor` that represents the same data. The
+ `row_shape` represents the dense shape of each row in the
+ resulting `tf.SparseTensor`, to which the effective batch size is
+ prepended. For example:
+
+ ```python
+ # NOTE: The following examples use `{ ... }` to represent the
+ # contents of a dataset.
+ a = { ['a', 'b', 'c'], ['a', 'b'], ['a', 'b', 'c', 'd'] }
+
+ a.apply(tf.data.experimental.dense_to_sparse_batch(
+ batch_size=2, row_shape=[6])) ==
+ {
+ ([[0, 0], [0, 1], [0, 2], [1, 0], [1, 1]], # indices
+ ['a', 'b', 'c', 'a', 'b'], # values
+ [2, 6]), # dense_shape
+ ([[0, 0], [0, 1], [0, 2], [0, 3]],
+ ['a', 'b', 'c', 'd'],
+ [1, 6])
+ }
+ ```
+
+ Args:
+ batch_size: A `tf.int64` scalar `tf.Tensor`, representing the
+ number of consecutive elements of this dataset to combine in a
+ single batch.
+ row_shape: A `tf.TensorShape` or `tf.int64` vector tensor-like
+ object representing the equivalent dense shape of a row in the
+ resulting `tf.SparseTensor`. Each element of this dataset must
+ have the same rank as `row_shape`, and must have size less
+ than or equal to `row_shape` in each dimension.
+
+ Returns:
+ A `Dataset` transformation function, which can be passed to
+ `tf.data.Dataset.apply`.
+ """
+
+ def _apply_fn(dataset):
+ return _DenseToSparseBatchDataset(dataset, batch_size, row_shape)
+
+ return _apply_fn
+
+
+def padded_batch_window(dataset, padded_shape, padding_value=None):
+ """Batches a window of tensors with padding.
+
+ Args:
+ dataset: the input dataset.
+ padded_shape: (Optional.) `tf.TensorShape` or `tf.int64` vector tensor-like
+ object representing the shape to which the input elements should be padded
+ prior to batching. Any unknown dimensions (e.g. `tf.Dimension(None)` in a
+ `tf.TensorShape` or `-1` in a tensor-like object) will be padded to the
+ maximum size of that dimension in each batch.
+ padding_value: (Optional.) A scalar-shaped `tf.Tensor`, representing the
+ padding value to use. Defaults are `0` for numeric types and the empty
+ string for string types. If `dataset` contains `tf.SparseTensor`, this
+ value is ignored.
+
+ Returns:
+ A `Tensor` representing the batch of the entire input dataset.
+
+ Raises:
+ ValueError: if invalid arguments are provided.
+ """
+ if not issubclass(dataset.output_classes,
+ (ops.Tensor, sparse_tensor.SparseTensor)):
+ raise TypeError("Input dataset expected to have a single tensor component")
+ if issubclass(dataset.output_classes, (ops.Tensor)):
+ return _padded_batch_dense_window(dataset, padded_shape, padding_value)
+ elif issubclass(dataset.output_classes, (sparse_tensor.SparseTensor)):
+ if padding_value is not None:
+ raise ValueError("Padding value not allowed for sparse tensors")
+ return _padded_batch_sparse_window(dataset, padded_shape)
+ else:
+ raise TypeError("Unsupported dataset type: %s" % dataset.output_classes)
+
+
+def _padded_batch_dense_window(dataset, padded_shape, padding_value=None):
+ """Batches a window of dense tensors with padding."""
+
+ padded_shape = math_ops.cast(
+ convert.partial_shape_to_tensor(padded_shape), dtypes.int32)
+
+ def key_fn(_):
+ return np.int64(0)
+
+ def max_init_fn(_):
+ return padded_shape
+
+ def max_reduce_fn(state, value):
+ """Computes the maximum shape to pad to."""
+ condition = math_ops.reduce_all(
+ math_ops.logical_or(
+ math_ops.less_equal(array_ops.shape(value), padded_shape),
+ math_ops.equal(padded_shape, -1)))
+ assert_op = control_flow_ops.Assert(condition, [
+ "Actual shape greater than padded shape: ",
+ array_ops.shape(value), padded_shape
+ ])
+ with ops.control_dependencies([assert_op]):
+ return math_ops.maximum(state, array_ops.shape(value))
+
+ def finalize_fn(state):
+ return state
+
+ # Compute the padded shape.
+ max_reducer = grouping.Reducer(max_init_fn, max_reduce_fn, finalize_fn)
+ padded_shape = get_single_element.get_single_element(
+ dataset.apply(grouping.group_by_reducer(key_fn, max_reducer)))
+
+ if padding_value is None:
+ if dataset.output_types == dtypes.string:
+ padding_value = ""
+ elif dataset.output_types == dtypes.bool:
+ padding_value = False
+ elif dataset.output_types == dtypes.variant:
+ raise TypeError("Unable to create padding for field of type 'variant'")
+ else:
+ padding_value = 0
+
+ def batch_init_fn(_):
+ batch_shape = array_ops.concat(
+ [np.array([0], dtype=np.int32), padded_shape], 0)
+ return gen_array_ops.empty(batch_shape, dtype=dataset.output_types)
+
+ def batch_reduce_fn(state, value):
+ return array_ops.concat([state, [value]], 0)
+
+ def pad_fn(value):
+ shape = array_ops.shape(value)
+ left = array_ops.zeros_like(shape)
+ right = padded_shape - shape
+ return array_ops.pad(
+ value, array_ops.stack([left, right], 1), constant_values=padding_value)
+
+ batch_reducer = grouping.Reducer(batch_init_fn, batch_reduce_fn, finalize_fn)
+ return get_single_element.get_single_element(
+ dataset.map(pad_fn).apply(
+ grouping.group_by_reducer(key_fn, batch_reducer)))
+
+
+def _padded_batch_sparse_window(dataset, padded_shape):
+ """Batches a window of sparse tensors with padding."""
+
+ def key_fn(_):
+ return np.int64(0)
+
+ def max_init_fn(_):
+ return convert.partial_shape_to_tensor(padded_shape)
+
+ def max_reduce_fn(state, value):
+ """Computes the maximum shape to pad to."""
+ condition = math_ops.reduce_all(
+ math_ops.logical_or(
+ math_ops.less_equal(value.dense_shape, padded_shape),
+ math_ops.equal(padded_shape, -1)))
+ assert_op = control_flow_ops.Assert(condition, [
+ "Actual shape greater than padded shape: ", value.dense_shape,
+ padded_shape
+ ])
+ with ops.control_dependencies([assert_op]):
+ return math_ops.maximum(state, value.dense_shape)
+
+ def finalize_fn(state):
+ return state
+
+ # Compute the padded shape.
+ max_reducer = grouping.Reducer(max_init_fn, max_reduce_fn, finalize_fn)
+ padded_shape = get_single_element.get_single_element(
+ dataset.apply(grouping.group_by_reducer(key_fn, max_reducer)))
+
+ def batch_init_fn(_):
+ indices_shape = array_ops.concat([[0], [array_ops.size(padded_shape) + 1]],
+ 0)
+ return sparse_tensor.SparseTensor(
+ indices=gen_array_ops.empty(indices_shape, dtype=dtypes.int64),
+ values=constant_op.constant([], shape=[0], dtype=dataset.output_types),
+ dense_shape=array_ops.concat(
+ [np.array([0], dtype=np.int64), padded_shape], 0))
+
+ def batch_reduce_fn(state, value):
+ padded_value = sparse_tensor.SparseTensor(
+ indices=value.indices, values=value.values, dense_shape=padded_shape)
+ reshaped_value = sparse_ops.sparse_reshape(
+ padded_value,
+ array_ops.concat(
+ [np.array([1], dtype=np.int64), padded_value.dense_shape], 0))
+ return sparse_ops.sparse_concat(0, [state, reshaped_value])
+
+ reducer = grouping.Reducer(batch_init_fn, batch_reduce_fn, finalize_fn)
+ return get_single_element.get_single_element(
+ dataset.apply(grouping.group_by_reducer(key_fn, reducer)))
+
+
+class _UnbatchDataset(dataset_ops.UnaryDataset):
+ """A dataset that splits the elements of its input into multiple elements."""
+
+ def __init__(self, input_dataset):
+ """See `unbatch()` for more details."""
+ super(_UnbatchDataset, self).__init__(input_dataset)
+ flat_shapes = nest.flatten(input_dataset.output_shapes)
+ if any(s.ndims == 0 for s in flat_shapes):
+ raise ValueError("Cannot unbatch an input with scalar components.")
+ known_batch_dim = tensor_shape.Dimension(None)
+ for s in flat_shapes:
+ try:
+ known_batch_dim = known_batch_dim.merge_with(s[0])
+ except ValueError:
+ raise ValueError("Cannot unbatch an input whose components have "
+ "different batch sizes.")
+ self._input_dataset = input_dataset
+
+ def _as_variant_tensor(self):
+ return gen_dataset_ops.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
+
+
+@tf_export("data.experimental.unbatch")
+def unbatch():
+ """Splits elements of a dataset into multiple elements on the batch dimension.
+
+ For example, if elements of the dataset are shaped `[B, a0, a1, ...]`,
+ where `B` may vary for each input element, then for each element in the
+ dataset, the unbatched dataset will contain `B` consecutive elements
+ of shape `[a0, a1, ...]`.
+
+ ```python
+ # NOTE: The following example uses `{ ... }` to represent the contents
+ # of a dataset.
+ a = { ['a', 'b', 'c'], ['a', 'b'], ['a', 'b', 'c', 'd'] }
+
+ a.apply(tf.data.experimental.unbatch()) == {
+ 'a', 'b', 'c', 'a', 'b', 'a', 'b', 'c', 'd'}
+ ```
+
+ Returns:
+ A `Dataset` transformation function, which can be passed to
+ `tf.data.Dataset.apply`.
+ """
+
+ 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.
+ def normalize(arg, *rest):
+ if rest:
+ return sparse.serialize_many_sparse_tensors((arg,) + rest)
+ else:
+ return sparse.serialize_many_sparse_tensors(arg)
+
+ normalized_dataset = dataset.map(normalize)
+
+ # NOTE(mrry): Our `map()` has lost information about the sparseness
+ # of any SparseTensor components, so re-apply the structure of the
+ # original dataset.
+ restructured_dataset = _RestructuredDataset(
+ normalized_dataset,
+ dataset.output_types,
+ dataset.output_shapes,
+ dataset.output_classes,
+ allow_unsafe_cast=True)
+ return _UnbatchDataset(restructured_dataset)
+
+ return _apply_fn
+
+
+class _DenseToSparseBatchDataset(dataset_ops.UnaryDataset):
+ """A `Dataset` that batches ragged dense elements into `tf.SparseTensor`s."""
+
+ def __init__(self, input_dataset, batch_size, row_shape):
+ """See `Dataset.dense_to_sparse_batch()` for more details."""
+ super(_DenseToSparseBatchDataset, self).__init__(input_dataset)
+ if not isinstance(input_dataset.output_types, dtypes.DType):
+ raise TypeError("DenseToSparseDataset requires an input whose elements "
+ "have a single component, whereas the input has %r." %
+ input_dataset.output_types)
+ self._input_dataset = input_dataset
+ self._batch_size = batch_size
+ self._row_shape = row_shape
+
+ def _as_variant_tensor(self):
+ return gen_dataset_ops.dense_to_sparse_batch_dataset(
+ self._input_dataset._as_variant_tensor(), # pylint: disable=protected-access
+ self._batch_size,
+ row_shape=convert.partial_shape_to_tensor(self._row_shape),
+ **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
+
+
+class _RestructuredDataset(dataset_ops.UnaryDataset):
+ """An internal helper for changing the structure and shape of a dataset."""
+
+ def __init__(self,
+ dataset,
+ output_types,
+ output_shapes=None,
+ output_classes=None,
+ allow_unsafe_cast=False):
+ """Creates a new dataset with the given output types and shapes.
+
+ The given `dataset` must have a structure that is convertible:
+ * `dataset.output_types` must be the same as `output_types` module nesting.
+ * Each shape in `dataset.output_shapes` must be compatible with each shape
+ in `output_shapes` (if given).
+
+ Note: This helper permits "unsafe casts" for shapes, equivalent to using
+ `tf.Tensor.set_shape()` where domain-specific knowledge is available.
+
+ Args:
+ dataset: A `Dataset` object.
+ output_types: A nested structure of `tf.DType` objects.
+ output_shapes: (Optional.) A nested structure of `tf.TensorShape` objects.
+ If omitted, the shapes will be inherited from `dataset`.
+ output_classes: (Optional.) A nested structure of class types.
+ If omitted, the class types will be inherited from `dataset`.
+ allow_unsafe_cast: (Optional.) If `True`, the caller may switch the
+ reported output types and shapes of the restructured dataset, e.g. to
+ switch a sparse tensor represented as `tf.variant` to its user-visible
+ type and shape.
+
+ Raises:
+ ValueError: If either `output_types` or `output_shapes` is not compatible
+ with the structure of `dataset`.
+ """
+ super(_RestructuredDataset, self).__init__(dataset)
+ self._input_dataset = dataset
+
+ if not allow_unsafe_cast:
+ # Validate that the types are compatible.
+ output_types = nest.map_structure(dtypes.as_dtype, output_types)
+ flat_original_types = nest.flatten(dataset.output_types)
+ flat_new_types = nest.flatten(output_types)
+ if flat_original_types != flat_new_types:
+ raise ValueError(
+ "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))
+ else:
+ if not allow_unsafe_cast:
+ # Validate that the shapes are compatible.
+ nest.assert_same_structure(output_types, output_shapes)
+ flat_original_shapes = nest.flatten(dataset.output_shapes)
+ flat_new_shapes = nest.flatten_up_to(output_types, output_shapes)
+
+ for original_shape, new_shape in zip(flat_original_shapes,
+ flat_new_shapes):
+ if not original_shape.is_compatible_with(new_shape):
+ raise ValueError(
+ "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_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
+
+ 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
+
+
+class _MapAndBatchDataset(dataset_ops.MapDataset):
+ """A `Dataset` that maps a function over a batch of elements."""
+
+ def __init__(self, input_dataset, map_func, batch_size, num_parallel_calls,
+ drop_remainder):
+ """See `Dataset.map()` for details."""
+ super(_MapAndBatchDataset, self).__init__(input_dataset, map_func)
+ self._batch_size_t = ops.convert_to_tensor(
+ batch_size, dtype=dtypes.int64, name="batch_size")
+ self._num_parallel_calls_t = ops.convert_to_tensor(
+ num_parallel_calls, dtype=dtypes.int64, name="num_parallel_calls")
+ self._drop_remainder_t = ops.convert_to_tensor(
+ drop_remainder, dtype=dtypes.bool, name="drop_remainder")
+
+ self._batch_size = batch_size
+ self._drop_remainder = drop_remainder
+
+ def _as_variant_tensor(self):
+ # pylint: disable=protected-access
+ input_resource = self._input_dataset._as_variant_tensor()
+ return gen_dataset_ops.map_and_batch_dataset_v2(
+ input_resource,
+ self._map_func.captured_inputs,
+ f=self._map_func,
+ batch_size=self._batch_size_t,
+ num_parallel_calls=self._num_parallel_calls_t,
+ drop_remainder=self._drop_remainder_t,
+ **dataset_ops.flat_structure(self))
+ # pylint: enable=protected-access
+
+ @property
+ def output_shapes(self):
+ dim = self._batch_size if self._drop_remainder else None
+ return nest.pack_sequence_as(self._output_shapes, [
+ tensor_shape.vector(dim).concatenate(s)
+ for s in nest.flatten(self._output_shapes)
+ ])
+
+ @property
+ def output_types(self):
+ return self._output_types
+
+
+@tf_export("data.experimental.map_and_batch")
+def map_and_batch(map_func,
+ batch_size,
+ num_parallel_batches=None,
+ drop_remainder=False,
+ num_parallel_calls=None):
+ """Fused implementation of `map` and `batch`.
+
+ Maps `map_func` across `batch_size` consecutive elements of this dataset
+ and then combines them into a batch. Functionally, it is equivalent to `map`
+ followed by `batch`. However, by fusing the two transformations together, the
+ implementation can be more efficient. Surfacing this transformation in the API
+ is temporary. Once automatic input pipeline optimization is implemented,
+ the fusing of `map` and `batch` will happen automatically and this API will be
+ deprecated.
+
+ Args:
+ map_func: A function mapping a nested structure of tensors to another
+ nested structure of tensors.
+ batch_size: A `tf.int64` scalar `tf.Tensor`, representing the number of
+ consecutive elements of this dataset to combine in a single batch.
+ num_parallel_batches: (Optional.) A `tf.int64` scalar `tf.Tensor`,
+ representing the number of batches to create in parallel. On one hand,
+ higher values can help mitigate the effect of stragglers. On the other
+ hand, higher values can increase contention if CPU is scarce.
+ drop_remainder: (Optional.) A `tf.bool` scalar `tf.Tensor`, representing
+ 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.
+
+ Returns:
+ A `Dataset` transformation function, which can be passed to
+ `tf.data.Dataset.apply`.
+
+ Raises:
+ ValueError: If both `num_parallel_batches` and `num_parallel_calls` are
+ specified.
+ """
+
+ if num_parallel_batches is None and num_parallel_calls is None:
+ num_parallel_calls = batch_size
+ elif num_parallel_batches is not None and num_parallel_calls is None:
+ num_parallel_calls = batch_size * num_parallel_batches
+ elif num_parallel_batches is not None and num_parallel_calls is not None:
+ raise ValueError("The `num_parallel_batches` and `num_parallel_calls` "
+ "arguments are mutually exclusive.")
+
+ def _apply_fn(dataset):
+ return _MapAndBatchDataset(dataset, map_func, batch_size,
+ num_parallel_calls, drop_remainder)
+
+ return _apply_fn
diff --git a/tensorflow/python/data/experimental/ops/counter.py b/tensorflow/python/data/experimental/ops/counter.py
new file mode 100644
index 0000000..42200ea
--- /dev/null
+++ b/tensorflow/python/data/experimental/ops/counter.py
@@ -0,0 +1,55 @@
+# 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.
+# ==============================================================================
+"""The Counter Dataset."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python.data.experimental.ops import scan_ops
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("data.experimental.Counter")
+def Counter(start=0, step=1, dtype=dtypes.int64):
+ """Creates a `Dataset` that counts from `start` in steps of size `step`.
+
+ For example:
+
+ ```python
+ Dataset.count() == [0, 1, 2, ...)
+ Dataset.count(2) == [2, 3, ...)
+ Dataset.count(2, 5) == [2, 7, 12, ...)
+ Dataset.count(0, -1) == [0, -1, -2, ...)
+ Dataset.count(10, -1) == [10, 9, ...)
+ ```
+
+ Args:
+ start: (Optional.) The starting value for the counter. Defaults to 0.
+ step: (Optional.) The step size for the counter. Defaults to 1.
+ dtype: (Optional.) The data type for counter elements. Defaults to
+ `tf.int64`.
+
+ Returns:
+ A `Dataset` of scalar `dtype` elements.
+ """
+ with ops.name_scope("counter"):
+ start = ops.convert_to_tensor(start, dtype=dtype, name="start")
+ step = ops.convert_to_tensor(step, dtype=dtype, name="step")
+ return dataset_ops.Dataset.from_tensors(0).repeat(None).apply(
+ scan_ops.scan(start, lambda state, _: (state + step, state)))
diff --git a/tensorflow/python/data/experimental/ops/enumerate_ops.py b/tensorflow/python/data/experimental/ops/enumerate_ops.py
new file mode 100644
index 0000000..a1af98f
--- /dev/null
+++ b/tensorflow/python/data/experimental/ops/enumerate_ops.py
@@ -0,0 +1,60 @@
+# 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.
+# ==============================================================================
+"""Enumerate dataset transformations."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.framework import dtypes
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("data.experimental.enumerate_dataset")
+def enumerate_dataset(start=0):
+ """A transformation that enumerate the elements of a dataset.
+
+ It is Similar to python's `enumerate`.
+ For example:
+
+ ```python
+ # NOTE: The following examples use `{ ... }` to represent the
+ # contents of a dataset.
+ a = { 1, 2, 3 }
+ b = { (7, 8), (9, 10) }
+
+ # The nested structure of the `datasets` argument determines the
+ # structure of elements in the resulting dataset.
+ a.apply(tf.data.experimental.enumerate(start=5)) == { (5, 1), (6, 2), (7, 3) }
+ b.apply(tf.data.experimental.enumerate()) == { (0, (7, 8)), (1, (9, 10)) }
+ ```
+
+ Args:
+ start: A `tf.int64` scalar `tf.Tensor`, representing the start
+ value for enumeration.
+
+ Returns:
+ A `Dataset` transformation function, which can be passed to
+ `tf.data.Dataset.apply`.
+ """
+
+ def _apply_fn(dataset):
+ max_value = np.iinfo(dtypes.int64.as_numpy_dtype).max
+ return dataset_ops.Dataset.zip((dataset_ops.Dataset.range(start, max_value),
+ dataset))
+
+ return _apply_fn
diff --git a/tensorflow/python/data/experimental/ops/error_ops.py b/tensorflow/python/data/experimental/ops/error_ops.py
new file mode 100644
index 0000000..82e274b
--- /dev/null
+++ b/tensorflow/python/data/experimental/ops/error_ops.py
@@ -0,0 +1,78 @@
+# 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.
+# ==============================================================================
+"""Ignore_errors dataset transformations."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.ops import gen_experimental_dataset_ops
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("data.experimental.ignore_errors")
+def ignore_errors():
+ """Creates a `Dataset` from another `Dataset` and silently ignores any errors.
+
+ Use this transformation to produce a dataset that contains the same elements
+ as the input, but silently drops any elements that caused an error. For
+ example:
+
+ ```python
+ dataset = tf.data.Dataset.from_tensor_slices([1., 2., 0., 4.])
+
+ # Computing `tf.check_numerics(1. / 0.)` will raise an InvalidArgumentError.
+ dataset = dataset.map(lambda x: tf.check_numerics(1. / x, "error"))
+
+ # Using `ignore_errors()` will drop the element that causes an error.
+ dataset =
+ dataset.apply(tf.data.experimental.ignore_errors()) # ==> {1., 0.5, 0.2}
+ ```
+
+ Returns:
+ A `Dataset` transformation function, which can be passed to
+ `tf.data.Dataset.apply`.
+ """
+
+ def _apply_fn(dataset):
+ return _IgnoreErrorsDataset(dataset)
+
+ return _apply_fn
+
+
+class _IgnoreErrorsDataset(dataset_ops.UnaryDataset):
+ """A `Dataset` that silently ignores errors when computing its input."""
+
+ def __init__(self, input_dataset):
+ """See `Dataset.ignore_errors()` for details."""
+ super(_IgnoreErrorsDataset, self).__init__(input_dataset)
+ self._input_dataset = input_dataset
+
+ def _as_variant_tensor(self):
+ return gen_experimental_dataset_ops.experimental_ignore_errors_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 self._input_dataset.output_shapes
+
+ @property
+ def output_types(self):
+ return self._input_dataset.output_types
diff --git a/tensorflow/python/data/experimental/ops/get_single_element.py b/tensorflow/python/data/experimental/ops/get_single_element.py
new file mode 100644
index 0000000..1325261
--- /dev/null
+++ b/tensorflow/python/data/experimental/ops/get_single_element.py
@@ -0,0 +1,72 @@
+# 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.
+# ==============================================================================
+"""Python wrappers for Datasets and Iterators."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.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
+
+
+@tf_export("data.experimental.get_single_element")
+def get_single_element(dataset):
+ """Returns the single element in `dataset` as a nested structure of tensors.
+
+ This function enables you to use a `tf.data.Dataset` in a stateless
+ "tensor-in tensor-out" expression, without creating a `tf.data.Iterator`.
+ This can be useful when your preprocessing transformations are expressed
+ as a `Dataset`, and you want to use the transformation at serving time.
+ For example:
+
+ ```python
+ input_batch = tf.placeholder(tf.string, shape=[BATCH_SIZE])
+
+ def preprocessing_fn(input_str):
+ # ...
+ return image, label
+
+ dataset = (tf.data.Dataset.from_tensor_slices(input_batch)
+ .map(preprocessing_fn, num_parallel_calls=BATCH_SIZE)
+ .batch(BATCH_SIZE))
+
+ image_batch, label_batch = tf.data.experimental.get_single_element(dataset)
+ ```
+
+ Args:
+ dataset: A `tf.data.Dataset` object containing a single element.
+
+ Returns:
+ A nested structure of `tf.Tensor` objects, corresponding to the single
+ element of `dataset`.
+
+ Raises:
+ TypeError: if `dataset` is not a `tf.data.Dataset` object.
+ InvalidArgumentError (at runtime): if `dataset` does not contain exactly
+ one element.
+ """
+ if not isinstance(dataset, dataset_ops.Dataset):
+ 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
+ **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
new file mode 100644
index 0000000..18ba583
--- /dev/null
+++ b/tensorflow/python/data/experimental/ops/grouping.py
@@ -0,0 +1,551 @@
+# 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.
+# ==============================================================================
+"""Grouping dataset transformations."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+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 ops
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import gen_dataset_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("data.experimental.group_by_reducer")
+def group_by_reducer(key_func, reducer):
+ """A transformation that groups elements and performs a reduction.
+
+ This transformation maps element of a dataset to a key using `key_func` and
+ groups the elements by key. The `reducer` is used to process each group; its
+ `init_func` is used to initialize state for each group when it is created, the
+ `reduce_func` is used to update the state every time an element is mapped to
+ the matching group, and the `finalize_func` is used to map the final state to
+ an output value.
+
+ Args:
+ key_func: A function mapping a nested structure of tensors
+ (having shapes and types defined by `self.output_shapes` and
+ `self.output_types`) to a scalar `tf.int64` tensor.
+ reducer: An instance of `Reducer`, which captures the reduction logic using
+ the `init_func`, `reduce_func`, and `finalize_func` functions.
+
+ Returns:
+ A `Dataset` transformation function, which can be passed to
+ `tf.data.Dataset.apply`.
+ """
+
+ def _apply_fn(dataset):
+ """Function from `Dataset` to `Dataset` that applies the transformation."""
+ return _GroupByReducerDataset(dataset, key_func, reducer)
+
+ return _apply_fn
+
+
+@tf_export("data.experimental.group_by_window")
+def group_by_window(key_func,
+ reduce_func,
+ window_size=None,
+ window_size_func=None):
+ """A transformation that groups windows of elements by key and reduces them.
+
+ This transformation maps each consecutive element in a dataset to a key
+ using `key_func` and groups the elements by key. It then applies
+ `reduce_func` to at most `window_size_func(key)` elements matching the same
+ key. All except the final window for each key will contain
+ `window_size_func(key)` elements; the final window may be smaller.
+
+ You may provide either a constant `window_size` or a window size determined by
+ the key through `window_size_func`.
+
+ Args:
+ key_func: A function mapping a nested structure of tensors
+ (having shapes and types defined by `self.output_shapes` and
+ `self.output_types`) to a scalar `tf.int64` tensor.
+ reduce_func: A function mapping a key and a dataset of up to `window_size`
+ consecutive elements matching that key to another dataset.
+ window_size: A `tf.int64` scalar `tf.Tensor`, representing the number of
+ consecutive elements matching the same key to combine in a single
+ batch, which will be passed to `reduce_func`. Mutually exclusive with
+ `window_size_func`.
+ window_size_func: A function mapping a key to a `tf.int64` scalar
+ `tf.Tensor`, representing the number of consecutive elements matching
+ the same key to combine in a single batch, which will be passed to
+ `reduce_func`. Mutually exclusive with `window_size`.
+
+ Returns:
+ A `Dataset` transformation function, which can be passed to
+ `tf.data.Dataset.apply`.
+
+ Raises:
+ ValueError: if neither or both of {`window_size`, `window_size_func`} are
+ passed.
+ """
+ if (window_size is not None and window_size_func or
+ not (window_size is not None or window_size_func)):
+ raise ValueError("Must pass either window_size or window_size_func.")
+
+ if window_size is not None:
+
+ def constant_window_func(unused_key):
+ return ops.convert_to_tensor(window_size, dtype=dtypes.int64)
+
+ window_size_func = constant_window_func
+
+ assert window_size_func is not None
+
+ def _apply_fn(dataset):
+ """Function from `Dataset` to `Dataset` that applies the transformation."""
+ return _GroupByWindowDataset(dataset, key_func, reduce_func,
+ window_size_func)
+
+ return _apply_fn
+
+
+@tf_export("data.experimental.bucket_by_sequence_length")
+def bucket_by_sequence_length(element_length_func,
+ bucket_boundaries,
+ bucket_batch_sizes,
+ padded_shapes=None,
+ padding_values=None,
+ pad_to_bucket_boundary=False,
+ no_padding=False):
+ """A transformation that buckets elements in a `Dataset` by length.
+
+ Elements of the `Dataset` are grouped together by length and then are padded
+ and batched.
+
+ This is useful for sequence tasks in which the elements have variable length.
+ Grouping together elements that have similar lengths reduces the total
+ fraction of padding in a batch which increases training step efficiency.
+
+ Args:
+ element_length_func: function from element in `Dataset` to `tf.int32`,
+ determines the length of the element, which will determine the bucket it
+ goes into.
+ bucket_boundaries: `list<int>`, upper length boundaries of the buckets.
+ bucket_batch_sizes: `list<int>`, batch size per bucket. Length should be
+ `len(bucket_boundaries) + 1`.
+ padded_shapes: Nested structure of `tf.TensorShape` to pass to
+ `tf.data.Dataset.padded_batch`. If not provided, will use
+ `dataset.output_shapes`, which will result in variable length dimensions
+ being padded out to the maximum length in each batch.
+ padding_values: Values to pad with, passed to
+ `tf.data.Dataset.padded_batch`. Defaults to padding with 0.
+ pad_to_bucket_boundary: bool, if `False`, will pad dimensions with unknown
+ size to maximum length in batch. If `True`, will pad dimensions with
+ unknown size to bucket boundary minus 1 (i.e., the maximum length in each
+ bucket), and caller must ensure that the source `Dataset` does not contain
+ any elements with length longer than `max(bucket_boundaries)`.
+ no_padding: `bool`, indicates whether to pad the batch features (features
+ need to be either of type `tf.SparseTensor` or of same shape).
+
+ Returns:
+ A `Dataset` transformation function, which can be passed to
+ `tf.data.Dataset.apply`.
+
+ Raises:
+ ValueError: if `len(bucket_batch_sizes) != len(bucket_boundaries) + 1`.
+ """
+ with ops.name_scope("bucket_by_seq_length"):
+ if len(bucket_batch_sizes) != (len(bucket_boundaries) + 1):
+ raise ValueError(
+ "len(bucket_batch_sizes) must equal len(bucket_boundaries) + 1")
+
+ batch_sizes = constant_op.constant(bucket_batch_sizes, dtype=dtypes.int64)
+
+ def element_to_bucket_id(*args):
+ """Return int64 id of the length bucket for this element."""
+ seq_length = element_length_func(*args)
+
+ boundaries = list(bucket_boundaries)
+ buckets_min = [np.iinfo(np.int32).min] + boundaries
+ buckets_max = boundaries + [np.iinfo(np.int32).max]
+ conditions_c = math_ops.logical_and(
+ math_ops.less_equal(buckets_min, seq_length),
+ math_ops.less(seq_length, buckets_max))
+ bucket_id = math_ops.reduce_min(array_ops.where(conditions_c))
+
+ return bucket_id
+
+ def window_size_fn(bucket_id):
+ # The window size is set to the batch size for this bucket
+ window_size = batch_sizes[bucket_id]
+ return window_size
+
+ def make_padded_shapes(shapes, none_filler=None):
+ padded = []
+ for shape in nest.flatten(shapes):
+ shape = tensor_shape.TensorShape(shape)
+ shape = [
+ none_filler if d.value is None else d
+ for d in shape
+ ]
+ padded.append(shape)
+ return nest.pack_sequence_as(shapes, padded)
+
+ def batching_fn(bucket_id, grouped_dataset):
+ """Batch elements in dataset."""
+ batch_size = window_size_fn(bucket_id)
+ if no_padding:
+ return grouped_dataset.batch(batch_size)
+ none_filler = None
+ if pad_to_bucket_boundary:
+ err_msg = ("When pad_to_bucket_boundary=True, elements must have "
+ "length < max(bucket_boundaries).")
+ check = check_ops.assert_less(
+ bucket_id,
+ constant_op.constant(len(bucket_batch_sizes) - 1,
+ dtype=dtypes.int64),
+ message=err_msg)
+ with ops.control_dependencies([check]):
+ boundaries = constant_op.constant(bucket_boundaries,
+ dtype=dtypes.int64)
+ bucket_boundary = boundaries[bucket_id]
+ none_filler = bucket_boundary - 1
+ shapes = make_padded_shapes(
+ padded_shapes or grouped_dataset.output_shapes,
+ none_filler=none_filler)
+ return grouped_dataset.padded_batch(batch_size, shapes, padding_values)
+
+ def _apply_fn(dataset):
+ return dataset.apply(
+ group_by_window(element_to_bucket_id, batching_fn,
+ window_size_func=window_size_fn))
+
+ return _apply_fn
+
+
+def _map_x_dataset(map_func):
+ """A transformation that maps `map_func` across its input.
+
+ This transformation is similar to `tf.data.Dataset.map`, but in addition to
+ supporting dense and sparse tensor inputs, it also supports dataset inputs.
+
+ Args:
+ map_func: A function mapping a nested structure of tensors and/or datasets
+ (having shapes and types defined by `self.output_shapes` and
+ `self.output_types`) to another nested structure of tensors and/or
+ datasets.
+
+ Returns:
+ Dataset: A `Dataset`.
+ """
+
+ def _apply_fn(dataset):
+ """Function from `Dataset` to `Dataset` that applies the transformation."""
+ return _MapXDataset(dataset, map_func)
+
+ return _apply_fn
+
+
+class _GroupByReducerDataset(dataset_ops.UnaryDataset):
+ """A `Dataset` that groups its input and performs a reduction."""
+
+ def __init__(self, input_dataset, key_func, reducer):
+ """See `group_by_reducer()` for details."""
+ super(_GroupByReducerDataset, self).__init__(input_dataset)
+
+ self._input_dataset = input_dataset
+
+ self._make_key_func(key_func, input_dataset)
+ self._make_init_func(reducer.init_func)
+ self._make_reduce_func(reducer.reduce_func, input_dataset)
+ self._make_finalize_func(reducer.finalize_func)
+
+ def _make_key_func(self, key_func, input_dataset):
+ """Make wrapping Defun for key_func."""
+ wrapped_func = dataset_ops.StructuredFunctionWrapper(
+ key_func, "tf.data.experimental.group_by_reducer()", input_dataset)
+ if not (
+ wrapped_func.output_types == dtypes.int64 and
+ wrapped_func.output_shapes.is_compatible_with(tensor_shape.scalar())):
+ 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.function
+
+ def _make_init_func(self, init_func):
+ """Make wrapping Defun for init_func."""
+ wrapped_func = dataset_ops.StructuredFunctionWrapper(
+ init_func,
+ "tf.data.experimental.group_by_reducer()",
+ input_classes=ops.Tensor,
+ input_shapes=tensor_shape.scalar(),
+ input_types=dtypes.int64)
+ self._init_func = wrapped_func.function
+ self._state_classes = wrapped_func.output_classes
+ self._state_shapes = wrapped_func.output_shapes
+ self._state_types = wrapped_func.output_types
+
+ 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`.
+ need_to_rerun = True
+ while need_to_rerun:
+
+ wrapped_func = dataset_ops.StructuredFunctionWrapper(
+ reduce_func,
+ "tf.data.experimental.group_by_reducer()",
+ 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),
+ 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)):
+ if not issubclass(new_state_class, state_class):
+ raise TypeError(
+ "The element classes for the new state must match the initial "
+ "state. Expected %s; got %s." %
+ (self._state_classes, wrapped_func.output_classes))
+
+ # 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)):
+ 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))
+
+ # Extract shape information from the returned values.
+ flat_state_shapes = nest.flatten(self._state_shapes)
+ flat_new_state_shapes = nest.flatten(wrapped_func.output_shapes)
+ weakened_state_shapes = [
+ original.most_specific_compatible_shape(new)
+ for original, new in zip(flat_state_shapes, flat_new_state_shapes)
+ ]
+
+ need_to_rerun = False
+ for original_shape, weakened_shape in zip(flat_state_shapes,
+ weakened_state_shapes):
+ if original_shape.ndims is not None and (
+ weakened_shape.ndims is None or
+ original_shape.as_list() != weakened_shape.as_list()):
+ need_to_rerun = True
+ break
+
+ if need_to_rerun:
+ self._state_shapes = nest.pack_sequence_as(self._state_shapes,
+ weakened_state_shapes)
+
+ self._reduce_func = wrapped_func.function
+ self._reduce_func.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,
+ "tf.data.experimental.group_by_reducer()",
+ input_classes=self._state_classes,
+ input_shapes=self._state_shapes,
+ input_types=self._state_types)
+ self._finalize_func = wrapped_func.function
+ self._output_classes = wrapped_func.output_classes
+ self._output_shapes = wrapped_func.output_shapes
+ self._output_types = wrapped_func.output_types
+
+ @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 _as_variant_tensor(self):
+ return gen_dataset_ops.group_by_reducer_dataset(
+ self._input_dataset._as_variant_tensor(), # pylint: disable=protected-access
+ self._key_func.captured_inputs,
+ self._init_func.captured_inputs,
+ self._reduce_func.captured_inputs,
+ self._finalize_func.captured_inputs,
+ key_func=self._key_func,
+ init_func=self._init_func,
+ reduce_func=self._reduce_func,
+ finalize_func=self._finalize_func,
+ **dataset_ops.flat_structure(self))
+
+
+class _GroupByWindowDataset(dataset_ops.UnaryDataset):
+ """A `Dataset` that groups its input and performs a windowed reduction."""
+
+ def __init__(self, input_dataset, key_func, reduce_func, window_size_func):
+ """See `group_by_window()` for details."""
+ super(_GroupByWindowDataset, self).__init__(input_dataset)
+
+ self._input_dataset = input_dataset
+
+ self._make_key_func(key_func, input_dataset)
+ self._make_reduce_func(reduce_func, input_dataset)
+ self._make_window_size_func(window_size_func)
+
+ def _make_window_size_func(self, window_size_func):
+ """Make wrapping Defun for window_size_func."""
+ def window_size_func_wrapper(key):
+ return ops.convert_to_tensor(window_size_func(key), dtype=dtypes.int64)
+ wrapped_func = dataset_ops.StructuredFunctionWrapper(
+ window_size_func_wrapper,
+ "tf.data.experimental.group_by_window()",
+ 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())):
+ raise ValueError(
+ "`window_size_func` must return a single tf.int64 scalar tensor.")
+ self._window_size_func = wrapped_func.function
+
+ 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(
+ key_func_wrapper, "tf.data.experimental.group_by_window()",
+ input_dataset)
+ if not (
+ wrapped_func.output_types == dtypes.int64 and
+ wrapped_func.output_shapes.is_compatible_with(tensor_shape.scalar())):
+ raise ValueError(
+ "`key_func` must return a single tf.int64 scalar tensor.")
+ self._key_func = wrapped_func.function
+
+ def _make_reduce_func(self, reduce_func, input_dataset):
+ """Make wrapping Defun for reduce_func."""
+ nested_dataset = dataset_ops._NestedDatasetComponent(input_dataset) # pylint: disable=protected-access
+ wrapped_func = dataset_ops.StructuredFunctionWrapper(
+ reduce_func,
+ "tf.data.experimental.reduce_by_window()",
+ input_classes=(ops.Tensor, nested_dataset),
+ input_shapes=(tensor_shape.scalar(), nested_dataset),
+ input_types=(dtypes.int64, nested_dataset),
+ experimental_nested_dataset_support=True)
+ if not isinstance(
+ wrapped_func.output_classes, dataset_ops._NestedDatasetComponent): # pylint: disable=protected-access
+ raise TypeError("`reduce_func` must return a `Dataset` object.")
+ self._output_classes = wrapped_func.output_classes.output_classes
+ self._output_types = wrapped_func.output_types.output_types
+ self._output_shapes = wrapped_func.output_shapes.output_shapes
+ self._reduce_func = wrapped_func.function
+
+ @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 _as_variant_tensor(self):
+ return gen_dataset_ops.group_by_window_dataset(
+ self._input_dataset._as_variant_tensor(), # pylint: disable=protected-access
+ self._key_func.captured_inputs,
+ self._reduce_func.captured_inputs,
+ self._window_size_func.captured_inputs,
+ key_func=self._key_func,
+ reduce_func=self._reduce_func,
+ window_size_func=self._window_size_func,
+ **dataset_ops.flat_structure(self))
+
+
+@tf_export("data.experimental.Reducer")
+class Reducer(object):
+ """A reducer is used for reducing a set of elements.
+
+ A reducer is represented as a tuple of the three functions:
+ 1) initialization function: key => initial state
+ 2) reduce function: (old state, input) => new state
+ 3) finalization function: state => result
+ """
+
+ def __init__(self, init_func, reduce_func, finalize_func):
+ self._init_func = init_func
+ self._reduce_func = reduce_func
+ self._finalize_func = finalize_func
+
+ @property
+ def init_func(self):
+ return self._init_func
+
+ @property
+ def reduce_func(self):
+ return self._reduce_func
+
+ @property
+ def finalize_func(self):
+ return self._finalize_func
+
+
+class _MapXDataset(dataset_ops.UnaryDataset):
+ """A `Dataset` that maps a function over elements in its input."""
+
+ def __init__(self, input_dataset, map_func):
+ """See `map_x_dataset()` for details."""
+ super(_MapXDataset, self).__init__(input_dataset)
+ self._input_dataset = input_dataset
+
+ wrapped_func = dataset_ops.StructuredFunctionWrapper(
+ map_func,
+ "tf.data.experimental.map_x_dataset()",
+ input_dataset,
+ experimental_nested_dataset_support=True)
+ self._output_classes = wrapped_func.output_classes
+ self._output_shapes = wrapped_func.output_shapes
+ self._output_types = wrapped_func.output_types
+ self._map_func = wrapped_func.function
+
+ def _as_variant_tensor(self):
+ input_t = self._input_dataset._as_variant_tensor() # pylint: disable=protected-access
+ return gen_dataset_ops.map_dataset(
+ input_t,
+ self._map_func.captured_inputs,
+ f=self._map_func,
+ **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
diff --git a/tensorflow/contrib/data/python/ops/indexed_dataset_ops.py b/tensorflow/python/data/experimental/ops/indexed_dataset_ops.py
similarity index 88%
rename from tensorflow/contrib/data/python/ops/indexed_dataset_ops.py
rename to tensorflow/python/data/experimental/ops/indexed_dataset_ops.py
index cc76ab0..9c06474 100644
--- a/tensorflow/contrib/data/python/ops/indexed_dataset_ops.py
+++ b/tensorflow/python/data/experimental/ops/indexed_dataset_ops.py
@@ -19,14 +19,13 @@
import abc
-from tensorflow.contrib.data.python.ops import contrib_op_loader # pylint: disable=unused-import
-from tensorflow.contrib.data.python.ops import gen_dataset_ops
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.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
class MaterializedIndexedDataset(object):
@@ -57,7 +56,7 @@
A tensor containing the values corresponding to `index`.
"""
# TODO(saeta): nest.pack_sequence_as(...)
- return gen_dataset_ops.indexed_dataset_get(
+ return ged_ops.experimental_indexed_dataset_get(
self._materialized_resource,
index,
output_types=nest.flatten(
@@ -90,16 +89,18 @@
container = ""
if shared_name is None:
shared_name = ""
- materialized_resource = gen_dataset_ops.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)))
+ materialized_resource = (
+ 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))))
with ops.colocate_with(materialized_resource):
- materializer = gen_dataset_ops.indexed_dataset_materialize(
+ materializer = ged_ops.experimental_indexed_dataset_materialize(
self._as_variant_tensor(), materialized_resource)
return MaterializedIndexedDataset(materialized_resource, materializer,
self.output_classes, self.output_types,
@@ -170,7 +171,7 @@
return tensor_shape.scalar()
def _as_variant_tensor(self):
- return gen_dataset_ops.identity_indexed_dataset(self._size)
+ return ged_ops.experimental_identity_indexed_dataset(self._size)
def _inputs(self):
return []
diff --git a/tensorflow/python/data/experimental/ops/interleave_ops.py b/tensorflow/python/data/experimental/ops/interleave_ops.py
new file mode 100644
index 0000000..a3c0948
--- /dev/null
+++ b/tensorflow/python/data/experimental/ops/interleave_ops.py
@@ -0,0 +1,262 @@
+# 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.
+# ==============================================================================
+"""Non-deterministic dataset transformations."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python.data.experimental.ops import random_ops
+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.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_experimental_dataset_ops
+from tensorflow.python.ops import gen_stateless_random_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("data.experimental.parallel_interleave")
+def parallel_interleave(map_func,
+ cycle_length,
+ block_length=1,
+ sloppy=False,
+ buffer_output_elements=None,
+ prefetch_input_elements=None):
+ """A parallel version of the `Dataset.interleave()` transformation.
+
+ `parallel_interleave()` maps `map_func` across its input to produce nested
+ datasets, and outputs their elements interleaved. Unlike
+ `tf.data.Dataset.interleave`, it gets elements from `cycle_length` nested
+ datasets in parallel, which increases the throughput, especially in the
+ presence of stragglers. Furthermore, the `sloppy` argument can be used to
+ improve performance, by relaxing the requirement that the outputs are produced
+ in a deterministic order, and allowing the implementation to skip over nested
+ datasets whose elements are not readily available when requested.
+
+ Example usage:
+
+ ```python
+ # Preprocess 4 files concurrently.
+ filenames = tf.data.Dataset.list_files("/path/to/data/train*.tfrecords")
+ dataset = filenames.apply(
+ tf.data.experimental.parallel_interleave(
+ lambda filename: tf.data.TFRecordDataset(filename),
+ cycle_length=4))
+ ```
+
+ WARNING: If `sloppy` is `True`, the order of produced elements is not
+ deterministic.
+
+ Args:
+ map_func: A function mapping a nested structure of tensors to a `Dataset`.
+ cycle_length: The number of input `Dataset`s to interleave from in parallel.
+ block_length: The number of consecutive elements to pull from an input
+ `Dataset` before advancing to the next input `Dataset`.
+ sloppy: If false, elements are produced in deterministic order. Otherwise,
+ the implementation is allowed, for the sake of expediency, to produce
+ elements in a non-deterministic order.
+ buffer_output_elements: The number of elements each iterator being
+ interleaved should buffer (similar to the `.prefetch()` transformation for
+ each interleaved iterator).
+ prefetch_input_elements: The number of input elements to transform to
+ iterators before they are needed for interleaving.
+
+ Returns:
+ A `Dataset` transformation function, which can be passed to
+ `tf.data.Dataset.apply`.
+ """
+ def _apply_fn(dataset):
+ return readers.ParallelInterleaveDataset(
+ dataset, map_func, cycle_length, block_length, sloppy,
+ buffer_output_elements, prefetch_input_elements)
+
+ return _apply_fn
+
+
+class _DirectedInterleaveDataset(dataset_ops.Dataset):
+ """A substitute for `Dataset.interleave()` on a fixed list of datasets."""
+
+ def __init__(self, selector_input, data_inputs):
+ self._selector_input = selector_input
+ self._data_inputs = list(data_inputs)
+
+ for data_input in data_inputs[1:]:
+ if (data_input.output_types != data_inputs[0].output_types or
+ data_input.output_classes != data_inputs[0].output_classes):
+ raise TypeError("All datasets must have the same type and class.")
+
+ def _as_variant_tensor(self):
+ # pylint: disable=protected-access
+ return (
+ gen_experimental_dataset_ops.experimental_directed_interleave_dataset(
+ self._selector_input._as_variant_tensor(), [
+ data_input._as_variant_tensor()
+ for data_input in self._data_inputs
+ ], **dataset_ops.flat_structure(self)))
+ # pylint: enable=protected-access
+
+ def _inputs(self):
+ 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
+
+
+@tf_export("data.experimental.sample_from_datasets")
+def sample_from_datasets(datasets, weights=None, seed=None):
+ """Samples elements at random from the datasets in `datasets`.
+
+ Args:
+ datasets: A list of `tf.data.Dataset` objects with compatible structure.
+ weights: (Optional.) A list of `len(datasets)` floating-point values where
+ `weights[i]` represents the probability with which an element should be
+ sampled from `datasets[i]`, or a `tf.data.Dataset` object where each
+ element is such a list. Defaults to a uniform distribution across
+ `datasets`.
+ seed: (Optional.) A `tf.int64` scalar `tf.Tensor`, representing the
+ random seed that will be used to create the distribution. See
+ `tf.set_random_seed` for behavior.
+
+ Returns:
+ A dataset that interleaves elements from `datasets` at random, according to
+ `weights` if provided, otherwise with uniform probability.
+
+ Raises:
+ TypeError: If the `datasets` or `weights` arguments have the wrong type.
+ ValueError: If the `weights` argument is specified and does not match the
+ length of the `datasets` element.
+ """
+ num_datasets = len(datasets)
+ if not isinstance(weights, dataset_ops.Dataset):
+ if weights is None:
+ # Select inputs with uniform probability.
+ logits = [[1.0] * num_datasets]
+
+ else:
+ # Use the given `weights` as the probability of choosing the respective
+ # input.
+ weights = ops.convert_to_tensor(weights, name="weights")
+ if weights.dtype not in (dtypes.float32, dtypes.float64):
+ raise TypeError("`weights` must be convertible to a tensor of "
+ "`tf.float32` or `tf.float64` elements.")
+ if not weights.shape.is_compatible_with([num_datasets]):
+ raise ValueError(
+ "`weights` must be a vector of length `len(datasets)`.")
+
+ # The `stateless_multinomial()` op expects log-probabilities, as opposed
+ # to weights.
+ logits = array_ops.expand_dims(math_ops.log(weights, name="logits"), 0)
+
+ # NOTE(mrry): We only specialize when `weights` is not a `Dataset`. When it
+ # is a `Dataset`, it is possible that evaluating it has a side effect the
+ # user depends on.
+ if len(datasets) == 1:
+ return datasets[0]
+
+ def select_dataset_constant_logits(seed):
+ return array_ops.squeeze(
+ gen_stateless_random_ops.stateless_multinomial(logits, 1, seed=seed),
+ axis=[0, 1])
+
+ selector_input = dataset_ops.MapDataset(
+ random_ops.RandomDataset(seed).batch(2),
+ select_dataset_constant_logits,
+ use_inter_op_parallelism=False)
+
+ else:
+ # Use each element of the given `weights` dataset as the probability of
+ # choosing the respective input.
+
+ # The `stateless_multinomial()` op expects log-probabilities, as opposed to
+ # weights.
+ logits_ds = weights.map(lambda *p: math_ops.log(p, name="logits"))
+
+ def select_dataset_varying_logits(logits, seed):
+ return array_ops.squeeze(
+ gen_stateless_random_ops.stateless_multinomial(logits, 1, seed=seed),
+ axis=[0, 1])
+
+ logits_and_seeds = dataset_ops.Dataset.zip(
+ (logits_ds, random_ops.RandomDataset(seed).batch(2)))
+ selector_input = dataset_ops.MapDataset(
+ logits_and_seeds,
+ select_dataset_varying_logits,
+ use_inter_op_parallelism=False)
+
+ return _DirectedInterleaveDataset(selector_input, datasets)
+
+
+@tf_export("data.experimental.choose_from_datasets")
+def choose_from_datasets(datasets, choice_dataset):
+ """Creates a dataset that deterministically chooses elements from `datasets`.
+
+ For example, given the following datasets:
+
+ ```python
+ datasets = [tf.data.Dataset.from_tensors("foo").repeat(),
+ tf.data.Dataset.from_tensors("bar").repeat(),
+ tf.data.Dataset.from_tensors("baz").repeat()]
+
+ # Define a dataset containing `[0, 1, 2, 0, 1, 2, 0, 1, 2]`.
+ choice_dataset = tf.data.Dataset.range(3).repeat(3)
+
+ result = tf.data.experimental.choose_from_datasets(datasets, choice_dataset)
+ ```
+
+ The elements of `result` will be:
+
+ ```
+ "foo", "bar", "baz", "foo", "bar", "baz", "foo", "bar", "baz"
+ ```
+
+ Args:
+ datasets: A list of `tf.data.Dataset` objects with compatible structure.
+ choice_dataset: A `tf.data.Dataset` of scalar `tf.int64` tensors between
+ `0` and `len(datasets) - 1`.
+
+ Returns:
+ A dataset that interleaves elements from `datasets` according to the values
+ of `choice_dataset`.
+
+ Raises:
+ TypeError: If the `datasets` or `choice_dataset` arguments have the wrong
+ type.
+ """
+ if not (choice_dataset.output_types == dtypes.int64
+ and choice_dataset.output_shapes.is_compatible_with(
+ tensor_shape.scalar())
+ and choice_dataset.output_classes == ops.Tensor):
+ raise TypeError("`choice_dataset` must be a dataset of scalar "
+ "`tf.int64` tensors.")
+ return _DirectedInterleaveDataset(choice_dataset, datasets)
diff --git a/tensorflow/python/data/experimental/ops/iterator_ops.py b/tensorflow/python/data/experimental/ops/iterator_ops.py
new file mode 100644
index 0000000..72d7d58
--- /dev/null
+++ b/tensorflow/python/data/experimental/ops/iterator_ops.py
@@ -0,0 +1,268 @@
+# 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.
+# ==============================================================================
+"""Iterator ops."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from tensorflow.python.data.ops import iterator_ops
+from tensorflow.python.data.ops import optional_ops
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import gen_dataset_ops
+from tensorflow.python.training import basic_session_run_hooks
+from tensorflow.python.training import checkpoint_management
+from tensorflow.python.training import saver as saver_lib
+from tensorflow.python.training import session_run_hook
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("data.experimental.make_saveable_from_iterator")
+def make_saveable_from_iterator(iterator):
+ """Returns a SaveableObject for saving/restore iterator state using Saver.
+
+ Args:
+ iterator: Iterator.
+
+ For example:
+
+ ```python
+ with tf.Graph().as_default():
+ ds = tf.data.Dataset.range(10)
+ iterator = ds.make_initializable_iterator()
+ # Build the iterator SaveableObject.
+ saveable_obj = tf.data.experimental.make_saveable_from_iterator(iterator)
+ # Add the SaveableObject to the SAVEABLE_OBJECTS collection so
+ # it can be automatically saved using Saver.
+ tf.add_to_collection(tf.GraphKeys.SAVEABLE_OBJECTS, saveable_obj)
+ saver = tf.train.Saver()
+
+ while continue_training:
+ ... Perform training ...
+ if should_save_checkpoint:
+ saver.save()
+ ```
+
+ Note: When restoring the iterator, the existing iterator state is completely
+ discarded. This means that any changes you may have made to the Dataset
+ graph will be discarded as well! This includes the new Dataset graph
+ that you may have built during validation. So, while running validation,
+ make sure to run the initializer for the validation input pipeline after
+ restoring the checkpoint.
+
+ Note: Not all iterators support checkpointing yet. Attempting to save the
+ state of an unsupported iterator will throw an error.
+ """
+ return _Saveable(iterator._iterator_resource) # pylint: disable=protected-access
+
+
+class _Saveable(saver_lib.BaseSaverBuilder.SaveableObject):
+ """SaveableObject for saving/restoring iterator state."""
+
+ def __init__(self, iterator_resource):
+ serialized_iterator = gen_dataset_ops.serialize_iterator(iterator_resource)
+ specs = [
+ saver_lib.BaseSaverBuilder.SaveSpec(serialized_iterator, "",
+ iterator_resource.name + "-state")
+ ]
+ super(_Saveable, self).__init__(iterator_resource, specs,
+ iterator_resource.name)
+
+ def restore(self, restored_tensors, unused_restored_shapes):
+ with ops.colocate_with(self.op):
+ return gen_dataset_ops.deserialize_iterator(self.op, restored_tensors[0])
+
+
+@tf_export("data.experimental.CheckpointInputPipelineHook")
+class CheckpointInputPipelineHook(session_run_hook.SessionRunHook):
+ """Checkpoints input pipeline state every N steps or seconds.
+
+ This hook saves the state of the iterators in the `Graph` so that when
+ training is resumed the input pipeline continues from where it left off.
+ This could potentially avoid overfitting in certain pipelines where the
+ number of training steps per eval are small compared to the dataset
+ size or if the training pipeline is pre-empted.
+
+ Differences from `CheckpointSaverHook`:
+ 1. Saves only the input pipelines in the "iterators" collection and not the
+ global variables or other saveable objects.
+ 2. Does not write the `GraphDef` and `MetaGraphDef` to the summary.
+
+ Example of checkpointing the training pipeline:
+
+ ```python
+ est = tf.estimator.Estimator(model_fn)
+ while True:
+ est.train(
+ train_input_fn,
+ hooks=[tf.data.experimental.CheckpointInputPipelineHook(est)],
+ steps=train_steps_per_eval)
+ # Note: We do not pass the hook here.
+ metrics = est.evaluate(eval_input_fn)
+ if should_stop_the_training(metrics):
+ break
+ ```
+
+ This hook should be used if the input pipeline state needs to be saved
+ separate from the model checkpoint. Doing so may be useful for a few reasons:
+ 1. The input pipeline checkpoint may be large, if there are large shuffle
+ or prefetch buffers for instance, and may bloat the checkpoint size.
+ 2. If the input pipeline is shared between training and validation, restoring
+ the checkpoint during validation may override the validation input
+ pipeline.
+
+ For saving the input pipeline checkpoint alongside the model weights use
+ `tf.data.experimental.make_saveable_from_iterator` directly to create a
+ `SaveableObject` and add to the `SAVEABLE_OBJECTS` collection. Note, however,
+ that you will need to be careful not to restore the training iterator during
+ eval. You can do that by not adding the iterator to the SAVEABLE_OBJECTS
+ collector when building the eval graph.
+ """
+
+ def __init__(self, estimator):
+ """Initializes a `CheckpointInputPipelineHook`.
+
+ Args:
+ estimator: Estimator.
+
+ Raises:
+ ValueError: One of `save_steps` or `save_secs` should be set.
+ ValueError: At most one of saver or scaffold should be set.
+ """
+ # `checkpoint_basename` is "input.ckpt" for non-distributed pipelines or
+ # of the form "input_<task_type>_<task_id>.ckpt" for distributed pipelines.
+ # Note: The default `checkpoint_basename` used by `CheckpointSaverHook` is
+ # "model.ckpt". We intentionally choose the input pipeline checkpoint prefix
+ # to be different to avoid conflicts with the model checkpoint.
+
+ # pylint: disable=protected-access
+ checkpoint_prefix = "input"
+ if estimator._config.num_worker_replicas > 1:
+ # Distributed setting.
+ suffix = "_{}_{}".format(estimator._config.task_type,
+ estimator._config.task_id)
+ checkpoint_prefix += suffix
+ # pylint: enable=protected-access
+
+ # We use a composition paradigm instead of inheriting from
+ # `CheckpointSaverHook` because `Estimator` does an `isinstance` check
+ # to check whether a `CheckpointSaverHook` is already present in the list
+ # of hooks and if not, adds one. Inheriting from `CheckpointSaverHook`
+ # would thwart this behavior. This hook checkpoints *only the iterators*
+ # and not the graph variables.
+ self._checkpoint_saver_hook = basic_session_run_hooks.CheckpointSaverHook(
+ estimator.model_dir,
+ save_secs=estimator._config.save_checkpoints_secs, # pylint: disable=protected-access
+ save_steps=estimator._config.save_checkpoints_steps, # pylint: disable=protected-access
+ checkpoint_basename=checkpoint_prefix + ".ckpt")
+
+ # Name for the protocol buffer file that will contain the list of most
+ # recent checkpoints stored as a `CheckpointState` protocol buffer.
+ # This file, kept in the same directory as the checkpoint files, is
+ # automatically managed by the `Saver` to keep track of recent checkpoints.
+ # The default name used by the `Saver` for this file is "checkpoint". Here
+ # we use the name "checkpoint_<checkpoint_prefix>" so that in case the
+ # `checkpoint_dir` is the same as the model checkpoint directory, there are
+ # no conflicts during restore.
+ self._latest_filename = "checkpoint_" + checkpoint_prefix
+ self._first_run = True
+
+ def begin(self):
+ # Build a Saver that saves all iterators in the `GLOBAL_ITERATORS`
+ # collection if no `Saver` or `Scaffold` is provided.
+ # pylint: disable=protected-access
+ if (self._checkpoint_saver_hook._saver is None and
+ self._checkpoint_saver_hook._scaffold is None):
+ iterators = ops.get_collection(iterator_ops.GLOBAL_ITERATORS)
+ saveables = [_Saveable(i) for i in iterators]
+ self._checkpoint_saver_hook._saver = _CustomSaver(saveables,
+ self._latest_filename)
+ # pylint: enable=protected-access
+ self._checkpoint_saver_hook.begin()
+
+ def _restore_or_save_initial_ckpt(self, session):
+ # Ideally this should be run in after_create_session but is not for the
+ # following reason:
+ # Currently there is no way of enforcing an order of running the
+ # `SessionRunHooks`. Hence it is possible that the `_DatasetInitializerHook`
+ # is run *after* this hook. That is troublesome because
+ # 1. If a checkpoint exists and this hook restores it, the initializer hook
+ # will override it.
+ # 2. If no checkpoint exists, this hook will try to save an initialized
+ # iterator which will result in an exception.
+ #
+ # As a temporary fix we enter the following implicit contract between this
+ # hook and the _DatasetInitializerHook.
+ # 1. The _DatasetInitializerHook initializes the iterator in the call to
+ # after_create_session.
+ # 2. This hook saves the iterator on the first call to `before_run()`, which
+ # is guaranteed to happen after `after_create_session()` of all hooks
+ # have been run.
+
+ # Check if there is an existing checkpoint. If so, restore from it.
+ # pylint: disable=protected-access
+ latest_checkpoint_path = checkpoint_management.latest_checkpoint(
+ self._checkpoint_saver_hook._checkpoint_dir,
+ latest_filename=self._latest_filename)
+ if latest_checkpoint_path:
+ self._checkpoint_saver_hook._get_saver().restore(session,
+ latest_checkpoint_path)
+ else:
+ # The checkpoint saved here is the state at step "global_step".
+ # Note: We do not save the GraphDef or MetaGraphDef here.
+ global_step = session.run(self._checkpoint_saver_hook._global_step_tensor)
+ self._checkpoint_saver_hook._save(session, global_step)
+ self._checkpoint_saver_hook._timer.update_last_triggered_step(global_step)
+ # pylint: enable=protected-access
+
+ def before_run(self, run_context):
+ if self._first_run:
+ self._restore_or_save_initial_ckpt(run_context.session)
+ self._first_run = False
+ return self._checkpoint_saver_hook.before_run(run_context)
+
+ def after_run(self, run_context, run_values):
+ self._checkpoint_saver_hook.after_run(run_context, run_values)
+
+ def end(self, session):
+ self._checkpoint_saver_hook.end(session)
+
+
+class _CustomSaver(saver_lib.Saver):
+ """`Saver` with a different default `latest_filename`.
+
+ This is used in the `CheckpointInputPipelineHook` to avoid conflicts with
+ the model ckpt saved by the `CheckpointSaverHook`.
+ """
+
+ def __init__(self, var_list, latest_filename):
+ super(_CustomSaver, self).__init__(var_list)
+ self._latest_filename = latest_filename
+
+ def save(self,
+ sess,
+ save_path,
+ global_step=None,
+ latest_filename=None,
+ meta_graph_suffix="meta",
+ write_meta_graph=True,
+ write_state=True,
+ strip_default_attrs=False):
+ return super(_CustomSaver, self).save(
+ sess, save_path, global_step, latest_filename or self._latest_filename,
+ meta_graph_suffix, write_meta_graph, write_state, strip_default_attrs)
+
+
+tf_export("data.experimental.Optional")(optional_ops.Optional)
+tf_export("data.experimental.get_next_as_optional")(
+ iterator_ops.get_next_as_optional)
diff --git a/tensorflow/contrib/data/python/ops/map_defun.py b/tensorflow/python/data/experimental/ops/map_defun.py
similarity index 100%
rename from tensorflow/contrib/data/python/ops/map_defun.py
rename to tensorflow/python/data/experimental/ops/map_defun.py
diff --git a/tensorflow/contrib/data/python/ops/optimization.py b/tensorflow/python/data/experimental/ops/optimization.py
similarity index 94%
rename from tensorflow/contrib/data/python/ops/optimization.py
rename to tensorflow/python/data/experimental/ops/optimization.py
index 3eb172a..30348ed 100644
--- a/tensorflow/contrib/data/python/ops/optimization.py
+++ b/tensorflow/python/data/experimental/ops/optimization.py
@@ -17,12 +17,11 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.contrib.data.python.ops import contrib_op_loader # pylint: disable=unused-import
-from tensorflow.contrib.data.python.ops import gen_dataset_ops as contrib_gen_dataset_ops
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import gen_dataset_ops
+from tensorflow.python.ops import gen_experimental_dataset_ops
# A constant that can be used to enable auto-tuning.
AUTOTUNE = -1
@@ -54,7 +53,7 @@
Returns:
A `Dataset` transformation function, which can be passed to
- @{tf.data.Dataset.apply}.
+ `tf.data.Dataset.apply`.
"""
def _apply_fn(dataset):
@@ -97,7 +96,7 @@
transformations, dtype=dtypes.string, name="transformations")
def _as_variant_tensor(self):
- return contrib_gen_dataset_ops.assert_next_dataset(
+ return gen_experimental_dataset_ops.experimental_assert_next_dataset(
self._input_dataset._as_variant_tensor(), # pylint: disable=protected-access
self._transformations,
**dataset_ops.flat_structure(self))
diff --git a/tensorflow/python/data/experimental/ops/parsing_ops.py b/tensorflow/python/data/experimental/ops/parsing_ops.py
new file mode 100644
index 0000000..6615b90
--- /dev/null
+++ b/tensorflow/python/data/experimental/ops/parsing_ops.py
@@ -0,0 +1,152 @@
+# 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.
+# ==============================================================================
+"""Experimental `dataset` API for parsing example."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.util import nest
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.ops import gen_dataset_ops
+from tensorflow.python.ops import parsing_ops
+from tensorflow.python.util.tf_export import tf_export
+
+
+class _ParseExampleDataset(dataset_ops.UnaryDataset):
+ """A `Dataset` that parses `example` dataset into a `dict` dataset."""
+
+ 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)):
+ raise TypeError("Input dataset should be a dataset of vectors of strings")
+ self._num_parallel_calls = num_parallel_calls
+ # pylint: disable=protected-access
+ self._features = parsing_ops._prepend_none_dimension(features)
+ # sparse_keys and dense_keys come back sorted here.
+ (sparse_keys, sparse_types, dense_keys, dense_types, dense_defaults,
+ dense_shapes) = parsing_ops._features_to_raw_params(
+ self._features, [
+ parsing_ops.VarLenFeature, parsing_ops.SparseFeature,
+ parsing_ops.FixedLenFeature, parsing_ops.FixedLenSequenceFeature
+ ])
+ # TODO(b/112859642): Pass sparse_index and sparse_values for SparseFeature.
+ (_, dense_defaults_vec, sparse_keys, sparse_types, dense_keys, dense_shapes,
+ dense_shape_as_shape) = parsing_ops._process_raw_parameters(
+ None, dense_defaults, sparse_keys, sparse_types, dense_keys,
+ dense_types, dense_shapes)
+ # pylint: enable=protected-access
+ self._sparse_keys = sparse_keys
+ self._sparse_types = sparse_types
+ self._dense_keys = dense_keys
+ self._dense_defaults = dense_defaults_vec
+ self._dense_shapes = dense_shapes
+ self._dense_types = dense_types
+ dense_output_shapes = [
+ self._input_dataset.output_shapes.concatenate(shape)
+ for shape in dense_shape_as_shape
+ ]
+ sparse_output_shapes = [
+ self._input_dataset.output_shapes.concatenate([None])
+ for _ in range(len(sparse_keys))
+ ]
+
+ self._output_shapes = dict(
+ zip(self._dense_keys + self._sparse_keys,
+ dense_output_shapes + sparse_output_shapes))
+ self._output_types = dict(
+ zip(self._dense_keys + self._sparse_keys,
+ self._dense_types + self._sparse_types))
+ self._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))
+ ]))
+
+ def _as_variant_tensor(self):
+ return gen_dataset_ops.parse_example_dataset(
+ self._input_dataset._as_variant_tensor(), # pylint: disable=protected-access
+ self._num_parallel_calls,
+ self._dense_defaults,
+ self._sparse_keys,
+ self._dense_keys,
+ self._sparse_types,
+ self._dense_shapes,
+ **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
+
+
+# TODO(b/111553342): add arguments names and example names as well.
+@tf_export("data.experimental.parse_example_dataset")
+def parse_example_dataset(features, num_parallel_calls=1):
+ """A transformation that parses `Example` protos into a `dict` of tensors.
+
+ Parses a number of serialized `Example` protos given in `serialized`. We refer
+ to `serialized` as a batch with `batch_size` many entries of individual
+ `Example` protos.
+
+ This op parses serialized examples into a dictionary mapping keys to `Tensor`
+ and `SparseTensor` objects. `features` is a dict from keys to `VarLenFeature`,
+ `SparseFeature`, and `FixedLenFeature` objects. Each `VarLenFeature`
+ and `SparseFeature` is mapped to a `SparseTensor`, and each
+ `FixedLenFeature` is mapped to a `Tensor`. See `tf.parse_example` for more
+ details about feature dictionaries.
+
+ Args:
+ features: A `dict` mapping feature keys to `FixedLenFeature`,
+ `VarLenFeature`, and `SparseFeature` values.
+ num_parallel_calls: (Optional.) A `tf.int32` scalar `tf.Tensor`,
+ representing the number of parsing processes to call in parallel.
+
+ Returns:
+ A dataset transformation function, which can be passed to
+ `tf.data.Dataset.apply`.
+
+ Raises:
+ ValueError: if features argument is None.
+ """
+ if features is None:
+ raise ValueError("Missing: features was %s." % features)
+
+ def _apply_fn(dataset):
+ """Function from `Dataset` to `Dataset` that applies the transformation."""
+ out_dataset = _ParseExampleDataset(dataset, features, num_parallel_calls)
+ if any([
+ isinstance(feature, parsing_ops.SparseFeature)
+ for _, feature in features.items()
+ ]):
+ # pylint: disable=protected-access
+ # pylint: disable=g-long-lambda
+ out_dataset = out_dataset.map(
+ lambda x: parsing_ops._construct_sparse_tensors_for_sparse_features(
+ features, x), num_parallel_calls=num_parallel_calls)
+ return out_dataset
+
+ return _apply_fn
diff --git a/tensorflow/python/data/experimental/ops/prefetching_ops.py b/tensorflow/python/data/experimental/ops/prefetching_ops.py
new file mode 100644
index 0000000..48d7136
--- /dev/null
+++ b/tensorflow/python/data/experimental/ops/prefetching_ops.py
@@ -0,0 +1,531 @@
+# 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.
+# ==============================================================================
+"""Python wrapper for prefetching_ops."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import warnings
+
+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.framework import device as framework_device
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import function
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import functional_ops
+from tensorflow.python.ops import gen_dataset_ops
+from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.util.tf_export import tf_export
+
+
+def function_buffering_resource(string_arg,
+ target_device,
+ f,
+ buffer_size,
+ output_types,
+ container="",
+ shared_name=None,
+ name=None):
+ """Creates a FunctionBufferingResource.
+
+ A FunctionBufferingResource fills up a buffer by calling a function `f` on
+ `target_device`. `f` should take in only a single string argument as input.
+
+ Args:
+ string_arg: The single string argument to the function.
+ target_device: The device to run `f` on.
+ f: The function to be executed.
+ buffer_size: Size of the buffer to be populated.
+ output_types: The output types generated by the function.
+ container: (Optional) string. Defaults to "".
+ shared_name: (Optional) string.
+ name: (Optional) string to name the op.
+
+ Returns:
+ Handle to a FunctionBufferingResource.
+ """
+ if shared_name is None:
+ shared_name = ""
+ return ged_ops.experimental_function_buffering_resource(
+ string_arg=string_arg,
+ target_device=target_device,
+ shared_name=shared_name,
+ f=f,
+ buffer_size=buffer_size,
+ container=container,
+ name=name,
+ output_types=output_types)
+
+
+def function_buffering_resource_get_next(function_buffer_resource,
+ output_types,
+ name=None):
+ return ged_ops.experimental_function_buffering_resource_get_next(
+ function_buffer_resource=function_buffer_resource,
+ output_types=output_types,
+ name=name)
+
+
+def function_buffering_resource_reset(function_buffer_resource, name=None):
+ return ged_ops.experimental_function_buffering_resource_reset(
+ function_buffer_resource=function_buffer_resource, name=name)
+
+
+# pylint: disable=protected-access
+class _PrefetchToDeviceIterator(object):
+ """A replacement for `tf.data.Iterator` that prefetches to another device.
+
+ Args:
+ input_dataset: The input dataset
+ one_shot: If true, we make a one shot iterator that's already initialized.
+ device: A fully specified device string where we want to prefetch to
+ buffer_size: Size of the prefetching buffer.
+ shared_name: (Optional.) If non-empty, the returned iterator will be
+ shared under the given name across multiple sessions that share the
+ same devices (e.g. when using a remote server).
+
+ Returns:
+ An Iterator type object.
+ """
+
+ def __init__(self,
+ input_dataset,
+ one_shot,
+ device,
+ buffer_size,
+ shared_name=None):
+ self._input_dataset = input_dataset
+ self._get_next_call_count = 0
+ self._one_shot = one_shot
+ if shared_name is None:
+ shared_name = ""
+
+ if self._one_shot:
+ self._input_iterator = input_dataset.make_one_shot_iterator()
+ else:
+ self._input_iterator = iterator_ops.Iterator.from_structure(
+ self._input_dataset.output_types, self._input_dataset.output_shapes,
+ shared_name, self._input_dataset.output_classes)
+ input_iterator_handle = self._input_iterator.string_handle()
+
+ @function.Defun(dtypes.string)
+ def _prefetch_fn(handle):
+ """Prefetches one element from `input_iterator`."""
+ remote_iterator = iterator_ops.Iterator.from_string_handle(
+ handle, self._input_iterator.output_types,
+ self._input_iterator.output_shapes,
+ self._input_iterator.output_classes)
+ ret = remote_iterator.get_next()
+ return nest.flatten(sparse.serialize_sparse_tensors(ret))
+
+ iterator_device = ged_ops.experimental_iterator_get_device(
+ self._input_iterator._iterator_resource)
+
+ with ops.device(device):
+ self._buffering_resource = function_buffering_resource(
+ f=_prefetch_fn,
+ target_device=iterator_device,
+ string_arg=input_iterator_handle,
+ buffer_size=buffer_size,
+ shared_name=shared_name,
+ output_types=nest.flatten(
+ sparse.as_dense_types(self._input_dataset.output_types,
+ self._input_dataset.output_classes)))
+
+ if not self._one_shot:
+ reset_op = function_buffering_resource_reset(self._buffering_resource)
+ with ops.control_dependencies([reset_op]):
+ self._initializer = self._input_iterator.make_initializer(
+ self._input_dataset)
+
+ def get_next(self, name=None):
+ """See `tf.data.Iterator.get_next`."""
+ self._get_next_call_count += 1
+ if self._get_next_call_count > iterator_ops.GET_NEXT_CALL_WARNING_THRESHOLD:
+ warnings.warn(iterator_ops.GET_NEXT_CALL_WARNING_MESSAGE)
+
+ flat_ret = ged_ops.experimental_function_buffering_resource_get_next(
+ self._buffering_resource,
+ output_types=nest.flatten(
+ sparse.as_dense_types(self.output_types, self.output_classes)),
+ name=name)
+
+ ret = sparse.deserialize_sparse_tensors(
+ nest.pack_sequence_as(self.output_types, flat_ret),
+ self.output_types, self.output_shapes, self.output_classes)
+
+ for tensor, shape in zip(
+ nest.flatten(ret), nest.flatten(self.output_shapes)):
+ if isinstance(tensor, ops.Tensor):
+ tensor.set_shape(shape)
+
+ return ret
+
+ @property
+ def initializer(self):
+ if self._one_shot:
+ raise NotImplementedError("Can't initialize a one_shot_iterator")
+ return self._initializer
+
+ @property
+ def output_classes(self):
+ return self._input_dataset.output_classes
+
+ @property
+ def output_shapes(self):
+ return self._input_dataset.output_shapes
+
+ @property
+ def output_types(self):
+ return self._input_dataset.output_types
+
+
+class _PrefetchToDeviceEagerIterator(iterator_ops.EagerIterator):
+ """A replacement for `tf.data.Iterator` that prefetches to another device.
+
+ Args:
+ input_dataset: The input dataset
+ one_shot: If true, we make a one shot iterator that's already initialized.
+ device: A fully specified device string where we want to prefetch to
+ buffer_size: Size of the prefetching buffer.
+ shared_name: (Optional.) If non-empty, the returned iterator will be
+ shared under the given name across multiple sessions that share the
+ same devices (e.g. when using a remote server).
+
+ Returns:
+ An Iterator type object.
+ """
+
+ def __init__(self,
+ input_dataset,
+ device,
+ buffer_size):
+ with ops.device("/device:CPU:0"):
+ super(_PrefetchToDeviceEagerIterator, self).__init__(input_dataset)
+ input_iterator_handle = gen_dataset_ops.iterator_to_string_handle(
+ self._resource)
+
+ self._device = device
+
+ @function.Defun(dtypes.string)
+ def _prefetch_fn(handle):
+ """Prefetches one element from `input_iterator`."""
+ remote_iterator = iterator_ops.Iterator.from_string_handle(
+ handle, self.output_types, self.output_shapes, self.output_classes)
+ ret = remote_iterator.get_next()
+ return nest.flatten(sparse.serialize_sparse_tensors(ret))
+
+ _prefetch_fn.add_to_graph(None)
+
+ with ops.device(device):
+ self._buffering_resource = function_buffering_resource(
+ f=_prefetch_fn,
+ output_types=self._flat_output_types,
+ target_device=ged_ops.experimental_iterator_get_device(
+ self._resource),
+ string_arg=input_iterator_handle,
+ buffer_size=buffer_size,
+ shared_name=iterator_ops._generate_shared_name(
+ "function_buffer_resource"))
+
+ def _next_internal(self):
+ """Returns a nested structure of `tf.Tensor`s containing the next element.
+ """
+ # This runs in sync mode as iterators use an error status to communicate
+ # that there is no more data to iterate over.
+ # TODO(b/77291417): Fix
+ with context.execution_mode(context.SYNC):
+ with ops.device(self._device):
+ ret = ged_ops.experimental_function_buffering_resource_get_next(
+ function_buffer_resource=self._buffering_resource,
+ output_types=self._flat_output_types)
+ return sparse.deserialize_sparse_tensors(
+ nest.pack_sequence_as(self._output_types, ret), self._output_types,
+ self._output_shapes, self._output_classes)
+# pylint: enable=protected-access
+
+
+class _PrefetchToDeviceDataset(dataset_ops.UnaryDataset):
+ """A `Dataset` whose iterator prefetches elements to another device."""
+
+ def __init__(self, input_dataset, device, buffer_size):
+ super(_PrefetchToDeviceDataset, self).__init__(input_dataset)
+ self._input_dataset = input_dataset
+ self._device = device
+ self._buffer_size = buffer_size if buffer_size is not None else 1
+
+ # The static analysis cannot tell that the eager iterator's superclass has
+ # a `next()` method.
+ # pylint: disable=non-iterator-returned
+ def __iter__(self):
+ """Creates an `Iterator` for enumerating the elements of this dataset.
+
+ The returned iterator implements the Python iterator protocol and therefore
+ can only be used in eager mode.
+
+ Returns:
+ An `Iterator` over the elements of this dataset.
+
+ Raises:
+ RuntimeError: If eager execution is enabled.
+ """
+ if context.executing_eagerly():
+ return _PrefetchToDeviceEagerIterator(self._input_dataset, self._device,
+ self._buffer_size)
+ else:
+ raise RuntimeError("dataset.__iter__() is only supported when eager "
+ "execution is enabled.")
+ # pylint: enable=non-iterator-returned
+
+ def make_one_shot_iterator(self):
+ if context.executing_eagerly():
+ return _PrefetchToDeviceEagerIterator(self._input_dataset, self._device,
+ self._buffer_size)
+ else:
+ return _PrefetchToDeviceIterator(self._input_dataset, one_shot=True,
+ device=self._device,
+ buffer_size=self._buffer_size)
+
+ def make_initializable_iterator(self, shared_name=None):
+ return _PrefetchToDeviceIterator(
+ self._input_dataset,
+ one_shot=False,
+ device=self._device,
+ buffer_size=self._buffer_size,
+ shared_name=shared_name)
+
+ def _as_variant_tensor(self):
+ # TODO(mrry): Raise this error earlier (e.g. when one of the Dataset
+ # transformation methods is called.
+ # TODO(mrry): Investigate support for chaining further transformations after
+ # the prefetch, including GPU support.
+ raise NotImplementedError("`prefetch_to_device()` must be the last "
+ "transformation in a dataset pipeline.")
+
+ @property
+ def output_types(self):
+ return self._input_dataset.output_types
+
+ @property
+ def output_shapes(self):
+ return self._input_dataset.output_shapes
+
+ @property
+ def output_classes(self):
+ return self._input_dataset.output_classes
+
+
+@tf_export("data.experimental.prefetch_to_device")
+def prefetch_to_device(device, buffer_size=None):
+ """A transformation that prefetches dataset values to the given `device`.
+
+ NOTE: Although the transformation creates a `tf.data.Dataset`, the
+ transformation must be the final `Dataset` in the input pipeline.
+
+ Args:
+ device: A string. The name of a device to which elements will be prefetched.
+ buffer_size: (Optional.) The number of elements to buffer on `device`.
+ Defaults to an automatically chosen value.
+
+ Returns:
+ A `Dataset` transformation function, which can be passed to
+ `tf.data.Dataset.apply`.
+ """
+ def _apply_fn(dataset):
+ return _PrefetchToDeviceDataset(dataset, device, buffer_size)
+
+ return _apply_fn
+
+
+@tf_export("data.experimental.copy_to_device")
+def copy_to_device(target_device, source_device="/cpu:0"):
+ """A transformation that copies dataset elements to the given `target_device`.
+
+ Args:
+ target_device: The name of a device to which elements will be copied.
+ source_device: The original device on which `input_dataset` will be placed.
+
+ Returns:
+ A `Dataset` transformation function, which can be passed to
+ `tf.data.Dataset.apply`.
+ """
+
+ def _apply_fn(dataset):
+ return _CopyToDeviceDataset(
+ dataset, target_device=target_device, source_device=source_device)
+
+ return _apply_fn
+
+
+# TODO(rohanj): Use the _input_hostmem attr on the RemoteCall ops to indicate
+# all inputs to the Op are in host memory, thereby avoiding some unnecessary
+# Sends and Recvs.
+class _CopyToDeviceDataset(dataset_ops.UnaryDataset):
+ """A `Dataset` that copies elements to another device."""
+
+ def __init__(self, input_dataset, target_device, source_device="/cpu:0"):
+ """Constructs a _CopyToDeviceDataset.
+
+ Args:
+ input_dataset: `Dataset` to be copied
+ target_device: The name of the device to which elements would be copied.
+ source_device: Device where input_dataset would be placed.
+ """
+ super(_CopyToDeviceDataset, self).__init__(input_dataset)
+ self._input_dataset = input_dataset
+ self._target_device = target_device
+ spec = framework_device.DeviceSpec().from_string(self._target_device)
+ self._is_gpu_target = (spec.device_type == "GPU")
+ 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.
+
+ Returns:
+ A `string` tensor that encapsulates the iterator created.
+ """
+ # 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)
+ with ops.control_dependencies(
+ [gen_dataset_ops.make_iterator(ds_variant, resource)]):
+ return gen_dataset_ops.iterator_to_string_handle(resource)
+
+ @function.Defun()
+ def _remote_init_func():
+ return functional_ops.remote_call(
+ target=self._source_device,
+ args=_init_func.captured_inputs,
+ Tout=[dtypes.string],
+ f=_init_func)
+
+ self._init_func = _remote_init_func
+ self._init_captured_args = _remote_init_func.captured_inputs
+
+ @function.Defun(dtypes.string)
+ def _next_func(string_handle):
+ """Calls get_next for created iterator.
+
+ Args:
+ string_handle: An iterator string handle created by _init_func
+ Returns:
+ The elements generated from `input_dataset`
+ """
+ with ops.device(self._source_device_string):
+ 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))
+
+ @function.Defun(dtypes.string)
+ def _remote_next_func(string_handle):
+ return functional_ops.remote_call(
+ target=self._source_device,
+ args=[string_handle] + _next_func.captured_inputs,
+ Tout=self._flat_output_types,
+ f=_next_func)
+
+ self._next_func = _remote_next_func
+ self._next_captured_args = _remote_next_func.captured_inputs
+
+ @function.Defun(dtypes.string)
+ def _finalize_func(string_handle):
+ """Destroys the iterator resource created.
+
+ Args:
+ string_handle: An iterator string handle created by _init_func
+ Returns:
+ Tensor constant 0
+ """
+ iterator_resource = gen_dataset_ops.iterator_from_string_handle_v2(
+ string_handle,
+ output_types=self._flat_output_types,
+ output_shapes=self._flat_output_shapes)
+ with ops.control_dependencies([
+ resource_variable_ops.destroy_resource_op(
+ iterator_resource, ignore_lookup_error=True)]):
+ return array_ops.constant(0, dtypes.int64)
+
+ @function.Defun(dtypes.string)
+ def _remote_finalize_func(string_handle):
+ return functional_ops.remote_call(
+ target=self._source_device,
+ args=[string_handle] + _finalize_func.captured_inputs,
+ Tout=[dtypes.int64],
+ f=_finalize_func)
+
+ self._finalize_func = _remote_finalize_func
+ self._finalize_captured_args = _remote_finalize_func.captured_inputs
+
+ g = ops.get_default_graph()
+ _remote_init_func.add_to_graph(g)
+ _remote_next_func.add_to_graph(g)
+ _remote_finalize_func.add_to_graph(g)
+ # pylint: enable=protected-scope
+
+ # The one_shot_iterator implementation needs a 0 arg _make_dataset function
+ # that thereby captures all the inputs required to create the dataset. Since
+ # there are strings that are inputs to the GeneratorDataset which can't be
+ # placed on a GPU, this fails for the GPU case. Therefore, disabling it for
+ # GPU
+ def make_one_shot_iterator(self):
+ if self._is_gpu_target:
+ raise ValueError("Cannot create a one shot iterator when using "
+ "`tf.data.experimental.copy_to_device()` on GPU. Please "
+ "use `Dataset.make_initializable_iterator()` instead.")
+ else:
+ return super(_CopyToDeviceDataset, self).make_one_shot_iterator()
+
+ def _as_variant_tensor(self):
+ with ops.device(self._target_device):
+ return gen_dataset_ops.generator_dataset(
+ self._init_captured_args,
+ self._next_captured_args,
+ self._finalize_captured_args,
+ 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)
+
+ @property
+ def output_types(self):
+ return self._input_dataset.output_types
+
+ @property
+ def output_shapes(self):
+ return self._input_dataset.output_shapes
+
+ @property
+ def output_classes(self):
+ return self._input_dataset.output_classes
diff --git a/tensorflow/python/data/experimental/ops/random_ops.py b/tensorflow/python/data/experimental/ops/random_ops.py
new file mode 100644
index 0000000..e3a2aea
--- /dev/null
+++ b/tensorflow/python/data/experimental/ops/random_ops.py
@@ -0,0 +1,54 @@
+# 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.
+# ==============================================================================
+"""Datasets for random number generators."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.util import random_seed
+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_dataset_ops
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("data.experimental.RandomDataset")
+class RandomDataset(dataset_ops.DatasetSource):
+ """A `Dataset` of pseudorandom values."""
+
+ def __init__(self, seed=None):
+ """A `Dataset` of pseudorandom values."""
+ super(RandomDataset, self).__init__()
+ self._seed, self._seed2 = random_seed.get_seed(seed)
+
+ def _as_variant_tensor(self):
+ return gen_dataset_ops.random_dataset(
+ seed=self._seed,
+ seed2=self._seed2,
+ **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
diff --git a/tensorflow/python/data/experimental/ops/readers.py b/tensorflow/python/data/experimental/ops/readers.py
new file mode 100644
index 0000000..3b2d094
--- /dev/null
+++ b/tensorflow/python/data/experimental/ops/readers.py
@@ -0,0 +1,904 @@
+# 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.
+# ==============================================================================
+"""Python wrappers for reader Datasets."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import collections
+import csv
+
+import numpy as np
+
+from tensorflow.python.data.experimental.ops import batching
+from tensorflow.python.data.experimental.ops import interleave_ops
+from tensorflow.python.data.experimental.ops import optimization
+from tensorflow.python.data.experimental.ops import parsing_ops
+from tensorflow.python.data.experimental.ops import shuffle_ops
+from tensorflow.python.data.ops import dataset_ops
+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.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_dataset_ops
+from tensorflow.python.ops import gen_experimental_dataset_ops
+from tensorflow.python.platform import gfile
+from tensorflow.python.util.tf_export import tf_export
+
+_ACCEPTABLE_CSV_TYPES = (dtypes.float32, dtypes.float64, dtypes.int32,
+ dtypes.int64, dtypes.string)
+
+
+def _is_valid_int32(str_val):
+ try:
+ # Checks equality to prevent int32 overflow
+ return dtypes.int32.as_numpy_dtype(str_val) == dtypes.int64.as_numpy_dtype(
+ str_val)
+ except (ValueError, OverflowError):
+ return False
+
+
+def _is_valid_int64(str_val):
+ try:
+ dtypes.int64.as_numpy_dtype(str_val)
+ return True
+ except (ValueError, OverflowError):
+ return False
+
+
+def _is_valid_float(str_val, float_dtype):
+ try:
+ return float_dtype.as_numpy_dtype(str_val) < np.inf
+ except ValueError:
+ return False
+
+
+def _infer_type(str_val, na_value, prev_type):
+ """Given a string, infers its tensor type.
+
+ Infers the type of a value by picking the least 'permissive' type possible,
+ while still allowing the previous type inference for this column to be valid.
+
+ Args:
+ str_val: String value to infer the type of.
+ na_value: Additional string to recognize as a NA/NaN CSV value.
+ prev_type: Type previously inferred based on values of this column that
+ we've seen up till now.
+ Returns:
+ Inferred dtype.
+ """
+ if str_val in ("", na_value):
+ # If the field is null, it gives no extra information about its type
+ return prev_type
+
+ type_list = [
+ dtypes.int32, dtypes.int64, dtypes.float32, dtypes.float64, dtypes.string
+ ] # list of types to try, ordered from least permissive to most
+
+ type_functions = [
+ _is_valid_int32,
+ _is_valid_int64,
+ lambda str_val: _is_valid_float(str_val, dtypes.float32),
+ lambda str_val: _is_valid_float(str_val, dtypes.float64),
+ lambda str_val: True,
+ ] # Corresponding list of validation functions
+
+ for i in range(len(type_list)):
+ validation_fn = type_functions[i]
+ if validation_fn(str_val) and (prev_type is None or
+ prev_type in type_list[:i + 1]):
+ return type_list[i]
+
+
+def _next_csv_row(filenames, num_cols, field_delim, use_quote_delim, header):
+ """Generator that yields rows of CSV file(s) in order."""
+ for fn in filenames:
+ with file_io.FileIO(fn, "r") as f:
+ rdr = csv.reader(
+ f,
+ delimiter=field_delim,
+ quoting=csv.QUOTE_MINIMAL if use_quote_delim else csv.QUOTE_NONE)
+ if header:
+ next(rdr) # Skip header lines
+
+ for csv_row in rdr:
+ if len(csv_row) != num_cols:
+ raise ValueError(
+ "Problem inferring types: CSV row has different number of fields "
+ "than expected.")
+ yield csv_row
+
+
+def _infer_column_defaults(filenames, num_cols, field_delim, use_quote_delim,
+ na_value, header, num_rows_for_inference,
+ select_columns):
+ """Infers column types from the first N valid CSV records of files."""
+ if select_columns is None:
+ select_columns = range(num_cols)
+ inferred_types = [None] * len(select_columns)
+
+ for i, csv_row in enumerate(
+ _next_csv_row(filenames, num_cols, field_delim, use_quote_delim, header)):
+ if num_rows_for_inference is not None and i >= num_rows_for_inference:
+ break
+
+ for j, col_index in enumerate(select_columns):
+ inferred_types[j] = _infer_type(csv_row[col_index], na_value,
+ inferred_types[j])
+
+ # Replace None's with a default type
+ inferred_types = [t or dtypes.string for t in inferred_types]
+ # Default to 0 or '' for null values
+ return [
+ constant_op.constant([0 if t is not dtypes.string else ""], dtype=t)
+ for t in inferred_types
+ ]
+
+
+def _infer_column_names(filenames, field_delim, use_quote_delim):
+ """Infers column names from first rows of files."""
+ csv_kwargs = {
+ "delimiter": field_delim,
+ "quoting": csv.QUOTE_MINIMAL if use_quote_delim else csv.QUOTE_NONE
+ }
+ with file_io.FileIO(filenames[0], "r") as f:
+ try:
+ column_names = next(csv.reader(f, **csv_kwargs))
+ except StopIteration:
+ raise ValueError(("Received StopIteration when reading the header line "
+ "of %s. Empty file?") % filenames[0])
+
+ for name in filenames[1:]:
+ with file_io.FileIO(name, "r") as f:
+ try:
+ if next(csv.reader(f, **csv_kwargs)) != column_names:
+ raise ValueError(
+ "Files have different column names in the header row.")
+ except StopIteration:
+ raise ValueError(("Received StopIteration when reading the header line "
+ "of %s. Empty file?") % filenames[0])
+ return column_names
+
+
+def _get_sorted_col_indices(select_columns, column_names):
+ """Transforms select_columns argument into sorted column indices."""
+ names_to_indices = {n: i for i, n in enumerate(column_names)}
+ num_cols = len(column_names)
+ for i, v in enumerate(select_columns):
+ if isinstance(v, int):
+ if v < 0 or v >= num_cols:
+ raise ValueError(
+ "Column index %d specified in select_columns out of valid range." %
+ v)
+ continue
+ if v not in names_to_indices:
+ raise ValueError(
+ "Value '%s' specified in select_columns not a valid column index or "
+ "name." % v)
+ select_columns[i] = names_to_indices[v]
+
+ # Sort and ensure there are no duplicates
+ result = sorted(set(select_columns))
+ if len(result) != len(select_columns):
+ raise ValueError("select_columns contains duplicate columns")
+ return result
+
+
+def _maybe_shuffle_and_repeat(
+ dataset, num_epochs, shuffle, shuffle_buffer_size, shuffle_seed):
+ """Optionally shuffle and repeat dataset, as requested."""
+ if num_epochs != 1 and shuffle:
+ # Use shuffle_and_repeat for perf
+ return dataset.apply(
+ shuffle_ops.shuffle_and_repeat(shuffle_buffer_size, num_epochs,
+ shuffle_seed))
+ elif shuffle:
+ return dataset.shuffle(shuffle_buffer_size, shuffle_seed)
+ elif num_epochs != 1:
+ return dataset.repeat(num_epochs)
+ return dataset
+
+
+def make_tf_record_dataset(file_pattern,
+ batch_size,
+ parser_fn=None,
+ num_epochs=None,
+ shuffle=True,
+ shuffle_buffer_size=None,
+ shuffle_seed=None,
+ prefetch_buffer_size=optimization.AUTOTUNE,
+ num_parallel_reads=None,
+ num_parallel_parser_calls=None,
+ drop_final_batch=False):
+ """Reads and optionally parses TFRecord files into a dataset.
+
+ Provides common functionality such as batching, optional parsing, shuffling,
+ and performant defaults.
+
+ Args:
+ file_pattern: List of files or patterns of TFRecord file paths.
+ See `tf.gfile.Glob` for pattern rules.
+ batch_size: An int representing the number of records to combine
+ in a single batch.
+ parser_fn: (Optional.) A function accepting string input to parse
+ and process the record contents. This function must map records
+ to components of a fixed shape, so they may be batched. By
+ default, uses the record contents unmodified.
+ num_epochs: (Optional.) An int specifying the number of times this
+ dataset is repeated. If None (the default), cycles through the
+ dataset forever.
+ shuffle: (Optional.) A bool that indicates whether the input
+ should be shuffled. Defaults to `True`.
+ shuffle_buffer_size: (Optional.) Buffer size to use for
+ shuffling. A large buffer size ensures better shuffling, but
+ increases memory usage and startup time.
+ shuffle_seed: (Optional.) Randomization seed to use for shuffling.
+ prefetch_buffer_size: (Optional.) An int specifying the number of
+ feature batches to prefetch for performance improvement.
+ Defaults to auto-tune. Set to 0 to disable prefetching.
+ num_parallel_reads: (Optional.) Number of threads used to read
+ records from files. By default or if set to a value >1, the
+ results will be interleaved.
+ num_parallel_parser_calls: (Optional.) Number of parallel
+ records to parse in parallel. Defaults to an automatic selection.
+ drop_final_batch: (Optional.) Whether the last batch should be
+ dropped in case its size is smaller than `batch_size`; the
+ default behavior is not to drop the smaller batch.
+
+ Returns:
+ A dataset, where each element matches the output of `parser_fn`
+ except it will have an additional leading `batch-size` dimension,
+ or a `batch_size`-length 1-D tensor of strings if `parser_fn` is
+ unspecified.
+ """
+ files = dataset_ops.Dataset.list_files(
+ file_pattern, shuffle=shuffle, seed=shuffle_seed)
+
+ if num_parallel_reads is None:
+ # Note: We considered auto-tuning this value, but there is a concern
+ # that this affects the mixing of records from different files, which
+ # could affect training convergence/accuracy, so we are defaulting to
+ # a constant for now.
+ num_parallel_reads = 24
+ dataset = core_readers.TFRecordDataset(
+ files, num_parallel_reads=num_parallel_reads)
+
+ if shuffle_buffer_size is None:
+ # TODO(josh11b): Auto-tune this value when not specified
+ shuffle_buffer_size = 10000
+ dataset = _maybe_shuffle_and_repeat(
+ dataset, num_epochs, shuffle, shuffle_buffer_size, shuffle_seed)
+
+ # NOTE(mrry): We set `drop_final_batch=True` when `num_epochs is None` to
+ # improve the shape inference, because it makes the batch dimension static.
+ # It is safe to do this because in that case we are repeating the input
+ # indefinitely, and all batches will be full-sized.
+ drop_final_batch = drop_final_batch or num_epochs is None
+
+ if parser_fn is None:
+ dataset = dataset.batch(batch_size, drop_remainder=drop_final_batch)
+ else:
+ # TODO(josh11b): if num_parallel_parser_calls is None, use some function
+ # of num cores instead of map_and_batch's default behavior of one batch.
+ dataset = dataset.apply(batching.map_and_batch(
+ parser_fn, batch_size, num_parallel_calls=num_parallel_parser_calls,
+ drop_remainder=drop_final_batch))
+
+ if prefetch_buffer_size == 0:
+ return dataset
+ else:
+ return dataset.prefetch(buffer_size=prefetch_buffer_size)
+
+
+@tf_export("data.experimental.make_csv_dataset")
+def make_csv_dataset(
+ file_pattern,
+ batch_size,
+ column_names=None,
+ column_defaults=None,
+ label_name=None,
+ select_columns=None,
+ field_delim=",",
+ use_quote_delim=True,
+ na_value="",
+ header=True,
+ num_epochs=None,
+ shuffle=True,
+ shuffle_buffer_size=10000,
+ shuffle_seed=None,
+ prefetch_buffer_size=optimization.AUTOTUNE,
+ num_parallel_reads=1,
+ sloppy=False,
+ num_rows_for_inference=100,
+ compression_type=None,
+):
+ """Reads CSV files into a dataset.
+
+ Reads CSV files into a dataset, where each element is a (features, labels)
+ tuple that corresponds to a batch of CSV rows. The features dictionary
+ maps feature column names to `Tensor`s containing the corresponding
+ feature data, and labels is a `Tensor` containing the batch's label data.
+
+ Args:
+ file_pattern: List of files or patterns of file paths containing CSV
+ records. See `tf.gfile.Glob` for pattern rules.
+ batch_size: An int representing the number of records to combine
+ in a single batch.
+ column_names: An optional list of strings that corresponds to the CSV
+ columns, in order. One per column of the input record. If this is not
+ provided, infers the column names from the first row of the records.
+ These names will be the keys of the features dict of each dataset element.
+ column_defaults: A optional list of default values for the CSV fields. One
+ item per selected column of the input record. Each item in the list is
+ either a valid CSV dtype (float32, float64, int32, int64, or string), or a
+ `Tensor` with one of the aforementioned types. The tensor can either be
+ a scalar default value (if the column is optional), or an empty tensor (if
+ the column is required). If a dtype is provided instead of a tensor, the
+ column is also treated as required. If this list is not provided, tries
+ to infer types based on reading the first num_rows_for_inference rows of
+ files specified, and assumes all columns are optional, defaulting to `0`
+ for numeric values and `""` for string values. If both this and
+ `select_columns` are specified, these must have the same lengths, and
+ `column_defaults` is assumed to be sorted in order of increasing column
+ index.
+ label_name: A optional string corresponding to the label column. If
+ provided, the data for this column is returned as a separate `Tensor` from
+ the features dictionary, so that the dataset complies with the format
+ expected by a `tf.Estimator.train` or `tf.Estimator.evaluate` input
+ function.
+ select_columns: An optional list of integer indices or string column
+ names, that specifies a subset of columns of CSV data to select. If
+ column names are provided, these must correspond to names provided in
+ `column_names` or inferred from the file header lines. When this argument
+ is specified, only a subset of CSV columns will be parsed and returned,
+ corresponding to the columns specified. Using this results in faster
+ parsing and lower memory usage. If both this and `column_defaults` are
+ specified, these must have the same lengths, and `column_defaults` is
+ assumed to be sorted in order of increasing column index.
+ field_delim: An optional `string`. Defaults to `","`. Char delimiter to
+ separate fields in a record.
+ use_quote_delim: An optional bool. Defaults to `True`. If false, treats
+ double quotation marks as regular characters inside of the string fields.
+ na_value: Additional string to recognize as NA/NaN.
+ header: A bool that indicates whether the first rows of provided CSV files
+ correspond to header lines with column names, and should not be included
+ in the data.
+ num_epochs: An int specifying the number of times this dataset is repeated.
+ If None, cycles through the dataset forever.
+ shuffle: A bool that indicates whether the input should be shuffled.
+ shuffle_buffer_size: Buffer size to use for shuffling. A large buffer size
+ ensures better shuffling, but increases memory usage and startup time.
+ shuffle_seed: Randomization seed to use for shuffling.
+ prefetch_buffer_size: An int specifying the number of feature
+ batches to prefetch for performance improvement. Recommended value is the
+ number of batches consumed per training step. Defaults to auto-tune.
+
+ num_parallel_reads: Number of threads used to read CSV records from files.
+ If >1, the results will be interleaved.
+ sloppy: If `True`, reading performance will be improved at
+ the cost of non-deterministic ordering. If `False`, the order of elements
+ produced is deterministic prior to shuffling (elements are still
+ randomized if `shuffle=True`. Note that if the seed is set, then order
+ of elements after shuffling is deterministic). Defaults to `False`.
+ num_rows_for_inference: Number of rows of a file to use for type inference
+ if record_defaults is not provided. If None, reads all the rows of all
+ the files. Defaults to 100.
+ compression_type: (Optional.) A `tf.string` scalar evaluating to one of
+ `""` (no compression), `"ZLIB"`, or `"GZIP"`. Defaults to no compression.
+
+ Returns:
+ A dataset, where each element is a (features, labels) tuple that corresponds
+ to a batch of `batch_size` CSV rows. The features dictionary maps feature
+ column names to `Tensor`s containing the corresponding column data, and
+ labels is a `Tensor` containing the column data for the label column
+ specified by `label_name`.
+
+ Raises:
+ ValueError: If any of the arguments is malformed.
+ """
+ # Create dataset of all matching filenames
+ filenames = _get_file_names(file_pattern, False)
+ dataset = dataset_ops.Dataset.from_tensor_slices(filenames)
+ if shuffle:
+ dataset = dataset.shuffle(len(filenames), shuffle_seed)
+
+ # Clean arguments; figure out column names and defaults
+
+ if column_names is None:
+ if not header:
+ raise ValueError("Cannot infer column names without a header line.")
+ # If column names are not provided, infer from the header lines
+ column_names = _infer_column_names(filenames, field_delim, use_quote_delim)
+ if len(column_names) != len(set(column_names)):
+ raise ValueError("Cannot have duplicate column names.")
+
+ if select_columns is not None:
+ select_columns = _get_sorted_col_indices(select_columns, column_names)
+
+ if column_defaults is not None:
+ column_defaults = [
+ constant_op.constant([], dtype=x) if x in _ACCEPTABLE_CSV_TYPES else x
+ for x in column_defaults
+ ]
+ else:
+ # If column defaults are not provided, infer from records at graph
+ # construction time
+ column_defaults = _infer_column_defaults(
+ filenames, len(column_names), field_delim, use_quote_delim, na_value,
+ header, num_rows_for_inference, select_columns)
+
+ if select_columns is not None and len(column_defaults) != len(select_columns):
+ raise ValueError(
+ "If specified, column_defaults and select_columns must have same "
+ "length."
+ )
+ if select_columns is not None and len(column_names) > len(select_columns):
+ # Pick the relevant subset of column names
+ column_names = [column_names[i] for i in select_columns]
+
+ if label_name is not None and label_name not in column_names:
+ raise ValueError("`label_name` provided must be one of the columns.")
+
+ def filename_to_dataset(filename):
+ return CsvDataset(
+ filename,
+ record_defaults=column_defaults,
+ field_delim=field_delim,
+ use_quote_delim=use_quote_delim,
+ na_value=na_value,
+ select_cols=select_columns,
+ header=header,
+ compression_type=compression_type,
+ )
+
+ def map_fn(*columns):
+ """Organizes columns into a features dictionary.
+
+ Args:
+ *columns: list of `Tensor`s corresponding to one csv record.
+ Returns:
+ An OrderedDict of feature names to values for that particular record. If
+ label_name is provided, extracts the label feature to be returned as the
+ second element of the tuple.
+ """
+ features = collections.OrderedDict(zip(column_names, columns))
+ if label_name is not None:
+ label = features.pop(label_name)
+ return features, label
+ return features
+
+ # Read files sequentially (if num_parallel_reads=1) or in parallel
+ dataset = dataset.apply(
+ interleave_ops.parallel_interleave(
+ filename_to_dataset, cycle_length=num_parallel_reads, sloppy=sloppy))
+
+ dataset = _maybe_shuffle_and_repeat(
+ dataset, num_epochs, shuffle, shuffle_buffer_size, shuffle_seed)
+
+ # Apply batch before map for perf, because map has high overhead relative
+ # to the size of the computation in each map.
+ # NOTE(mrry): We set `drop_remainder=True` when `num_epochs is None` to
+ # improve the shape inference, because it makes the batch dimension static.
+ # It is safe to do this because in that case we are repeating the input
+ # indefinitely, and all batches will be full-sized.
+ dataset = dataset.batch(batch_size=batch_size,
+ drop_remainder=num_epochs is None)
+ dataset = dataset_ops.MapDataset(
+ dataset, map_fn, use_inter_op_parallelism=False)
+ dataset = dataset.prefetch(prefetch_buffer_size)
+
+ return dataset
+
+
+_DEFAULT_READER_BUFFER_SIZE_BYTES = 4 * 1024 * 1024 # 4 MB
+
+
+@tf_export("data.experimental.CsvDataset")
+class CsvDataset(dataset_ops.DatasetSource):
+ """A Dataset comprising lines from one or more CSV files."""
+
+ def __init__(self,
+ filenames,
+ record_defaults,
+ compression_type=None,
+ buffer_size=None,
+ header=False,
+ field_delim=",",
+ use_quote_delim=True,
+ na_value="",
+ select_cols=None):
+ """Creates a `CsvDataset` by reading and decoding CSV files.
+
+ The elements of this dataset correspond to records from the file(s).
+ RFC 4180 format is expected for CSV files
+ (https://tools.ietf.org/html/rfc4180)
+ Note that we allow leading and trailing spaces with int or float field.
+
+
+ For example, suppose we have a file 'my_file0.csv' with four CSV columns of
+ different data types:
+ ```
+ abcdefg,4.28E10,5.55E6,12
+ hijklmn,-5.3E14,,2
+ ```
+
+ We can construct a CsvDataset from it as follows:
+ ```python
+ dataset = tf.data.experimental.CsvDataset(
+ "my_file*.csv",
+ [tf.float32, # Required field, use dtype or empty tensor
+ tf.constant([0.0], dtype=tf.float32), # Optional field, default to 0.0
+ tf.int32, # Required field, use dtype or empty tensor
+ ],
+ select_cols=[1,2,3] # Only parse last three columns
+ )
+ ```
+
+ The expected output of its iterations is:
+ ```python
+ next_element = dataset.make_one_shot_iterator().get_next()
+ with tf.Session() as sess:
+ while True:
+ try:
+ print(sess.run(next_element))
+ except tf.errors.OutOfRangeError:
+ break
+
+ >> (4.28e10, 5.55e6, 12)
+ >> (-5.3e14, 0.0, 2)
+ ```
+
+ Args:
+ filenames: A `tf.string` tensor containing one or more filenames.
+ record_defaults: A list of default values for the CSV fields. Each item in
+ the list is either a valid CSV `DType` (float32, float64, int32, int64,
+ string), or a `Tensor` object with one of the above types. One per
+ column of CSV data, with either a scalar `Tensor` default value for the
+ column if it is optional, or `DType` or empty `Tensor` if required. If
+ both this and `select_columns` are specified, these must have the same
+ lengths, and `column_defaults` is assumed to be sorted in order of
+ increasing column index.
+ compression_type: (Optional.) A `tf.string` scalar evaluating to one of
+ `""` (no compression), `"ZLIB"`, or `"GZIP"`. Defaults to no
+ compression.
+ buffer_size: (Optional.) A `tf.int64` scalar denoting the number of bytes
+ to buffer while reading files. Defaults to 4MB.
+ header: (Optional.) A `tf.bool` scalar indicating whether the CSV file(s)
+ have header line(s) that should be skipped when parsing. Defaults to
+ `False`.
+ field_delim: (Optional.) A `tf.string` scalar containing the delimiter
+ character that separates fields in a record. Defaults to `","`.
+ use_quote_delim: (Optional.) A `tf.bool` scalar. If `False`, treats
+ double quotation marks as regular characters inside of string fields
+ (ignoring RFC 4180, Section 2, Bullet 5). Defaults to `True`.
+ na_value: (Optional.) A `tf.string` scalar indicating a value that will
+ be treated as NA/NaN.
+ select_cols: (Optional.) A sorted list of column indices to select from
+ the input data. If specified, only this subset of columns will be
+ parsed. Defaults to parsing all columns.
+ """
+ super(CsvDataset, self).__init__()
+ self._filenames = ops.convert_to_tensor(
+ filenames, dtype=dtypes.string, name="filenames")
+ self._compression_type = convert.optional_param_to_tensor(
+ "compression_type",
+ compression_type,
+ argument_default="",
+ argument_dtype=dtypes.string)
+ record_defaults = [
+ constant_op.constant([], dtype=x) if x in _ACCEPTABLE_CSV_TYPES else x
+ for x in record_defaults
+ ]
+ self._record_defaults = ops.convert_n_to_tensor(
+ record_defaults, name="record_defaults")
+ self._buffer_size = convert.optional_param_to_tensor(
+ "buffer_size", buffer_size, _DEFAULT_READER_BUFFER_SIZE_BYTES)
+ self._header = ops.convert_to_tensor(
+ header, dtype=dtypes.bool, name="header")
+ self._field_delim = ops.convert_to_tensor(
+ field_delim, dtype=dtypes.string, name="field_delim")
+ self._use_quote_delim = ops.convert_to_tensor(
+ use_quote_delim, dtype=dtypes.bool, name="use_quote_delim")
+ self._na_value = ops.convert_to_tensor(
+ na_value, dtype=dtypes.string, name="na_value")
+ self._select_cols = convert.optional_param_to_tensor(
+ "select_cols",
+ select_cols,
+ 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)))
+
+ def _as_variant_tensor(self):
+ # Constructs graph node for the dataset op.
+ return gen_experimental_dataset_ops.experimental_csv_dataset(
+ filenames=self._filenames,
+ record_defaults=self._record_defaults,
+ buffer_size=self._buffer_size,
+ header=self._header,
+ output_shapes=self._output_shapes,
+ field_delim=self._field_delim,
+ use_quote_delim=self._use_quote_delim,
+ na_value=self._na_value,
+ select_cols=self._select_cols,
+ compression_type=self._compression_type,
+ )
+
+ @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
+
+
+@tf_export("data.experimental.make_batched_features_dataset")
+def make_batched_features_dataset(file_pattern,
+ batch_size,
+ features,
+ reader=core_readers.TFRecordDataset,
+ label_key=None,
+ reader_args=None,
+ num_epochs=None,
+ shuffle=True,
+ shuffle_buffer_size=10000,
+ shuffle_seed=None,
+ prefetch_buffer_size=optimization.AUTOTUNE,
+ reader_num_threads=1,
+ parser_num_threads=2,
+ sloppy_ordering=False,
+ drop_final_batch=False):
+ """Returns a `Dataset` of feature dictionaries from `Example` protos.
+
+ If label_key argument is provided, returns a `Dataset` of tuple
+ comprising of feature dictionaries and label.
+
+ Example:
+
+ ```
+ serialized_examples = [
+ features {
+ feature { key: "age" value { int64_list { value: [ 0 ] } } }
+ feature { key: "gender" value { bytes_list { value: [ "f" ] } } }
+ feature { key: "kws" value { bytes_list { value: [ "code", "art" ] } } }
+ },
+ features {
+ feature { key: "age" value { int64_list { value: [] } } }
+ feature { key: "gender" value { bytes_list { value: [ "f" ] } } }
+ feature { key: "kws" value { bytes_list { value: [ "sports" ] } } }
+ }
+ ]
+ ```
+
+ We can use arguments:
+
+ ```
+ features: {
+ "age": FixedLenFeature([], dtype=tf.int64, default_value=-1),
+ "gender": FixedLenFeature([], dtype=tf.string),
+ "kws": VarLenFeature(dtype=tf.string),
+ }
+ ```
+
+ And the expected output is:
+
+ ```python
+ {
+ "age": [[0], [-1]],
+ "gender": [["f"], ["f"]],
+ "kws": SparseTensor(
+ indices=[[0, 0], [0, 1], [1, 0]],
+ values=["code", "art", "sports"]
+ dense_shape=[2, 2]),
+ }
+ ```
+
+ Args:
+ file_pattern: List of files or patterns of file paths containing
+ `Example` records. See `tf.gfile.Glob` for pattern rules.
+ batch_size: An int representing the number of records to combine
+ in a single batch.
+ features: A `dict` mapping feature keys to `FixedLenFeature` or
+ `VarLenFeature` values. See `tf.parse_example`.
+ reader: A function or class that can be
+ called with a `filenames` tensor and (optional) `reader_args` and returns
+ a `Dataset` of `Example` tensors. Defaults to `tf.data.TFRecordDataset`.
+ label_key: (Optional) A string corresponding to the key labels are stored in
+ `tf.Examples`. If provided, it must be one of the `features` key,
+ otherwise results in `ValueError`.
+ reader_args: Additional arguments to pass to the reader class.
+ num_epochs: Integer specifying the number of times to read through the
+ dataset. If None, cycles through the dataset forever. Defaults to `None`.
+ shuffle: A boolean, indicates whether the input should be shuffled. Defaults
+ to `True`.
+ shuffle_buffer_size: Buffer size of the ShuffleDataset. A large capacity
+ ensures better shuffling but would increase memory usage and startup time.
+ shuffle_seed: Randomization seed to use for shuffling.
+ prefetch_buffer_size: Number of feature batches to prefetch in order to
+ improve performance. Recommended value is the number of batches consumed
+ per training step. Defaults to auto-tune.
+ reader_num_threads: Number of threads used to read `Example` records. If >1,
+ the results will be interleaved.
+ parser_num_threads: Number of threads to use for parsing `Example` tensors
+ into a dictionary of `Feature` tensors.
+ sloppy_ordering: If `True`, reading performance will be improved at
+ the cost of non-deterministic ordering. If `False`, the order of elements
+ produced is deterministic prior to shuffling (elements are still
+ randomized if `shuffle=True`. Note that if the seed is set, then order
+ of elements after shuffling is deterministic). Defaults to `False`.
+ drop_final_batch: If `True`, and the batch size does not evenly divide the
+ input dataset size, the final smaller batch will be dropped. Defaults to
+ `False`.
+
+ Returns:
+ A dataset of `dict` elements, (or a tuple of `dict` elements and label).
+ Each `dict` maps feature keys to `Tensor` or `SparseTensor` objects.
+
+ Raises:
+ ValueError: If `label_key` is not one of the `features` keys.
+ """
+ # Create dataset of all matching filenames
+ filenames = _get_file_names(file_pattern, False)
+ dataset = dataset_ops.Dataset.from_tensor_slices(filenames)
+ if shuffle:
+ dataset = dataset.shuffle(len(filenames), shuffle_seed)
+
+ # Read `Example` records from files as tensor objects.
+ if reader_args is None:
+ reader_args = []
+
+ # Read files sequentially (if reader_num_threads=1) or in parallel
+ dataset = dataset.apply(
+ interleave_ops.parallel_interleave(
+ lambda filename: reader(filename, *reader_args),
+ cycle_length=reader_num_threads,
+ sloppy=sloppy_ordering))
+
+ # Extract values if the `Example` tensors are stored as key-value tuples.
+ if dataset.output_types == (dtypes.string, dtypes.string):
+ dataset = dataset_ops.MapDataset(
+ dataset, lambda _, v: v, use_inter_op_parallelism=False)
+
+ # Apply dataset repeat and shuffle transformations.
+ dataset = _maybe_shuffle_and_repeat(
+ dataset, num_epochs, shuffle, shuffle_buffer_size, shuffle_seed)
+
+ # NOTE(mrry): We set `drop_remainder=True` when `num_epochs is None` to
+ # improve the shape inference, because it makes the batch dimension static.
+ # It is safe to do this because in that case we are repeating the input
+ # indefinitely, and all batches will be full-sized.
+ dataset = dataset.batch(
+ batch_size, drop_remainder=drop_final_batch or num_epochs is None)
+
+ # Parse `Example` tensors to a dictionary of `Feature` tensors.
+ dataset = dataset.apply(
+ parsing_ops.parse_example_dataset(
+ features, num_parallel_calls=parser_num_threads))
+
+ if label_key:
+ if label_key not in features:
+ raise ValueError(
+ "The `label_key` provided (%r) must be one of the `features` keys." %
+ label_key)
+ dataset = dataset.map(lambda x: (x, x.pop(label_key)))
+
+ dataset = dataset.prefetch(prefetch_buffer_size)
+ return dataset
+
+
+def _get_file_names(file_pattern, shuffle):
+ """Parse list of file names from pattern, optionally shuffled.
+
+ Args:
+ file_pattern: File glob pattern, or list of glob patterns.
+ shuffle: Whether to shuffle the order of file names.
+
+ Returns:
+ List of file names matching `file_pattern`.
+
+ Raises:
+ ValueError: If `file_pattern` is empty, or pattern matches no files.
+ """
+ if isinstance(file_pattern, list):
+ if not file_pattern:
+ raise ValueError("File pattern is empty.")
+ file_names = []
+ for entry in file_pattern:
+ file_names.extend(gfile.Glob(entry))
+ else:
+ file_names = list(gfile.Glob(file_pattern))
+
+ if not file_names:
+ raise ValueError("No files match %s." % file_pattern)
+
+ # Sort files so it will be deterministic for unit tests.
+ if not shuffle:
+ file_names = sorted(file_names)
+ return file_names
+
+
+@tf_export("data.experimental.SqlDataset")
+class SqlDataset(dataset_ops.DatasetSource):
+ """A `Dataset` consisting of the results from a SQL query."""
+
+ def __init__(self, driver_name, data_source_name, query, output_types):
+ """Creates a `SqlDataset`.
+
+ `SqlDataset` allows a user to read data from the result set of a SQL query.
+ For example:
+
+ ```python
+ dataset = tf.data.experimental.SqlDataset("sqlite", "/foo/bar.sqlite3",
+ "SELECT name, age FROM people",
+ (tf.string, tf.int32))
+ iterator = dataset.make_one_shot_iterator()
+ next_element = iterator.get_next()
+ # Prints the rows of the result set of the above query.
+ while True:
+ try:
+ print(sess.run(next_element))
+ except tf.errors.OutOfRangeError:
+ break
+ ```
+
+ Args:
+ driver_name: A 0-D `tf.string` tensor containing the database type.
+ Currently, the only supported value is 'sqlite'.
+ data_source_name: A 0-D `tf.string` tensor containing a connection string
+ to connect to the database.
+ query: A 0-D `tf.string` tensor containing the SQL query to execute.
+ output_types: A tuple of `tf.DType` objects representing the types of the
+ columns returned by `query`.
+ """
+ super(SqlDataset, self).__init__()
+ self._driver_name = ops.convert_to_tensor(
+ driver_name, dtype=dtypes.string, name="driver_name")
+ self._data_source_name = ops.convert_to_tensor(
+ 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
+
+ def _as_variant_tensor(self):
+ return gen_dataset_ops.sql_dataset(self._driver_name,
+ self._data_source_name, self._query,
+ 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
diff --git a/tensorflow/python/data/experimental/ops/resampling.py b/tensorflow/python/data/experimental/ops/resampling.py
new file mode 100644
index 0000000..3a3040a
--- /dev/null
+++ b/tensorflow/python/data/experimental/ops/resampling.py
@@ -0,0 +1,296 @@
+# 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.
+# ==============================================================================
+"""Resampling dataset transformations."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from tensorflow.python.data.experimental.ops import batching
+from tensorflow.python.data.experimental.ops import interleave_ops
+from tensorflow.python.data.experimental.ops import scan_ops
+from tensorflow.python.data.ops import dataset_ops
+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_ops
+from tensorflow.python.ops import logging_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("data.experimental.rejection_resample")
+def rejection_resample(class_func, target_dist, initial_dist=None, seed=None):
+ """A transformation that resamples a dataset to achieve a target distribution.
+
+ **NOTE** Resampling is performed via rejection sampling; some fraction
+ of the input values will be dropped.
+
+ Args:
+ class_func: A function mapping an element of the input dataset to a scalar
+ `tf.int32` tensor. Values should be in `[0, num_classes)`.
+ target_dist: A floating point type tensor, shaped `[num_classes]`.
+ initial_dist: (Optional.) A floating point type tensor, shaped
+ `[num_classes]`. If not provided, the true class distribution is
+ estimated live in a streaming fashion.
+ seed: (Optional.) Python integer seed for the resampler.
+
+ Returns:
+ A `Dataset` transformation function, which can be passed to
+ `tf.data.Dataset.apply`.
+ """
+ def _apply_fn(dataset):
+ """Function from `Dataset` to `Dataset` that applies the transformation."""
+ target_dist_t = ops.convert_to_tensor(target_dist, name="target_dist")
+ class_values_ds = dataset.map(class_func)
+
+ # Get initial distribution.
+ if initial_dist is not None:
+ initial_dist_t = ops.convert_to_tensor(initial_dist, name="initial_dist")
+ acceptance_dist, prob_of_original = (
+ _calculate_acceptance_probs_with_mixing(initial_dist_t,
+ target_dist_t))
+ initial_dist_ds = dataset_ops.Dataset.from_tensors(
+ initial_dist_t).repeat()
+ acceptance_dist_ds = dataset_ops.Dataset.from_tensors(
+ acceptance_dist).repeat()
+ prob_of_original_ds = dataset_ops.Dataset.from_tensors(
+ prob_of_original).repeat()
+ else:
+ initial_dist_ds = _estimate_initial_dist_ds(
+ target_dist_t, class_values_ds)
+ acceptance_and_original_prob_ds = initial_dist_ds.map(
+ lambda initial: _calculate_acceptance_probs_with_mixing( # pylint: disable=g-long-lambda
+ initial, target_dist_t))
+ acceptance_dist_ds = acceptance_and_original_prob_ds.map(
+ lambda accept_prob, _: accept_prob)
+ prob_of_original_ds = acceptance_and_original_prob_ds.map(
+ lambda _, prob_original: prob_original)
+ filtered_ds = _filter_ds(dataset, acceptance_dist_ds, initial_dist_ds,
+ class_values_ds, seed)
+ # Prefetch filtered dataset for speed.
+ filtered_ds = filtered_ds.prefetch(3)
+
+ prob_original_static = _get_prob_original_static(
+ initial_dist_t, target_dist_t) if initial_dist is not None else None
+ if prob_original_static == 1:
+ return dataset_ops.Dataset.zip((class_values_ds, dataset))
+ elif prob_original_static == 0:
+ return filtered_ds
+ else:
+ return interleave_ops.sample_from_datasets(
+ [dataset_ops.Dataset.zip((class_values_ds, dataset)), filtered_ds],
+ weights=prob_of_original_ds.map(lambda prob: [(prob, 1.0 - prob)]),
+ seed=seed)
+
+ return _apply_fn
+
+
+def _get_prob_original_static(initial_dist_t, target_dist_t):
+ """Returns the static probability of sampling from the original.
+
+ `tensor_util.constant_value(prob_of_original)` returns `None` if it encounters
+ an Op that it isn't defined for. We have some custom logic to avoid this.
+
+ Args:
+ initial_dist_t: A tensor of the initial distribution.
+ target_dist_t: A tensor of the target distribution.
+
+ Returns:
+ The probability of sampling from the original distribution as a constant,
+ if it is a constant, or `None`.
+ """
+ init_static = tensor_util.constant_value(initial_dist_t)
+ target_static = tensor_util.constant_value(target_dist_t)
+
+ if init_static is None or target_static is None:
+ return None
+ else:
+ return np.min(target_static / init_static)
+
+
+def _filter_ds(dataset, acceptance_dist_ds, initial_dist_ds, class_values_ds,
+ seed):
+ """Filters a dataset based on per-class acceptance probabilities.
+
+ Args:
+ dataset: The dataset to be filtered.
+ acceptance_dist_ds: A dataset of acceptance probabilities.
+ initial_dist_ds: A dataset of the initial probability distribution, given or
+ estimated.
+ class_values_ds: A dataset of the corresponding classes.
+ seed: (Optional.) Python integer seed for the resampler.
+
+ Returns:
+ A dataset of (class value, data) after filtering.
+ """
+ def maybe_warn_on_large_rejection(accept_dist, initial_dist):
+ proportion_rejected = math_ops.reduce_sum((1 - accept_dist) * initial_dist)
+ return control_flow_ops.cond(
+ math_ops.less(proportion_rejected, .5),
+ lambda: accept_dist,
+ lambda: logging_ops.Print( # pylint: disable=g-long-lambda
+ accept_dist, [proportion_rejected, initial_dist, accept_dist],
+ message="Proportion of examples rejected by sampler is high: ",
+ summarize=100,
+ first_n=10))
+
+ acceptance_dist_ds = (dataset_ops.Dataset.zip((acceptance_dist_ds,
+ initial_dist_ds))
+ .map(maybe_warn_on_large_rejection))
+
+ def _gather_and_copy(class_val, acceptance_prob, data):
+ return class_val, array_ops.gather(acceptance_prob, class_val), data
+
+ current_probabilities_and_class_and_data_ds = dataset_ops.Dataset.zip(
+ (class_values_ds, acceptance_dist_ds, dataset)).map(_gather_and_copy)
+ filtered_ds = (
+ current_probabilities_and_class_and_data_ds
+ .filter(lambda _1, p, _2: random_ops.random_uniform([], seed=seed) < p))
+ return filtered_ds.map(lambda class_value, _, data: (class_value, data))
+
+
+def _estimate_initial_dist_ds(
+ target_dist_t, class_values_ds, dist_estimation_batch_size=32,
+ smoothing_constant=10):
+ num_classes = (target_dist_t.shape[0].value or
+ array_ops.shape(target_dist_t)[0])
+ initial_examples_per_class_seen = array_ops.fill(
+ [num_classes], np.int64(smoothing_constant))
+
+ def update_estimate_and_tile(num_examples_per_class_seen, c):
+ updated_examples_per_class_seen, dist = _estimate_data_distribution(
+ c, num_examples_per_class_seen)
+ tiled_dist = array_ops.tile(
+ array_ops.expand_dims(dist, 0), [dist_estimation_batch_size, 1])
+ return updated_examples_per_class_seen, tiled_dist
+
+ initial_dist_ds = (class_values_ds.batch(dist_estimation_batch_size)
+ .apply(scan_ops.scan(initial_examples_per_class_seen,
+ update_estimate_and_tile))
+ .apply(batching.unbatch()))
+
+ return initial_dist_ds
+
+
+def _get_target_to_initial_ratio(initial_probs, target_probs):
+ # Add tiny to initial_probs to avoid divide by zero.
+ denom = (initial_probs + np.finfo(initial_probs.dtype.as_numpy_dtype).tiny)
+ return target_probs / denom
+
+
+def _estimate_data_distribution(c, num_examples_per_class_seen):
+ """Estimate data distribution as labels are seen.
+
+ Args:
+ c: The class labels. Type `int32`, shape `[batch_size]`.
+ num_examples_per_class_seen: Type `int64`, shape `[num_classes]`,
+ containing counts.
+
+ Returns:
+ num_examples_per_lass_seen: Updated counts. Type `int64`, shape
+ `[num_classes]`.
+ dist: The updated distribution. Type `float32`, shape `[num_classes]`.
+ """
+ num_classes = num_examples_per_class_seen.get_shape()[0].value
+ # Update the class-count based on what labels are seen in batch.
+ num_examples_per_class_seen = math_ops.add(
+ num_examples_per_class_seen, math_ops.reduce_sum(
+ array_ops.one_hot(c, num_classes, dtype=dtypes.int64), 0))
+ init_prob_estimate = math_ops.truediv(
+ num_examples_per_class_seen,
+ math_ops.reduce_sum(num_examples_per_class_seen))
+ dist = math_ops.cast(init_prob_estimate, dtypes.float32)
+ return num_examples_per_class_seen, dist
+
+
+def _calculate_acceptance_probs_with_mixing(initial_probs, target_probs):
+ """Calculates the acceptance probabilities and mixing ratio.
+
+ In this case, we assume that we can *either* sample from the original data
+ distribution with probability `m`, or sample from a reshaped distribution
+ that comes from rejection sampling on the original distribution. This
+ rejection sampling is done on a per-class basis, with `a_i` representing the
+ probability of accepting data from class `i`.
+
+ This method is based on solving the following analysis for the reshaped
+ distribution:
+
+ Let F be the probability of a rejection (on any example).
+ Let p_i be the proportion of examples in the data in class i (init_probs)
+ Let a_i is the rate the rejection sampler should *accept* class i
+ Let t_i is the target proportion in the minibatches for class i (target_probs)
+
+ ```
+ F = sum_i(p_i * (1-a_i))
+ = 1 - sum_i(p_i * a_i) using sum_i(p_i) = 1
+ ```
+
+ An example with class `i` will be accepted if `k` rejections occur, then an
+ example with class `i` is seen by the rejector, and it is accepted. This can
+ be written as follows:
+
+ ```
+ t_i = sum_k=0^inf(F^k * p_i * a_i)
+ = p_i * a_j / (1 - F) using geometric series identity, since 0 <= F < 1
+ = p_i * a_i / sum_j(p_j * a_j) using F from above
+ ```
+
+ Note that the following constraints hold:
+ ```
+ 0 <= p_i <= 1, sum_i(p_i) = 1
+ 0 <= a_i <= 1
+ 0 <= t_i <= 1, sum_i(t_i) = 1
+ ```
+
+ A solution for a_i in terms of the other variables is the following:
+ ```a_i = (t_i / p_i) / max_i[t_i / p_i]```
+
+ If we try to minimize the amount of data rejected, we get the following:
+
+ M_max = max_i [ t_i / p_i ]
+ M_min = min_i [ t_i / p_i ]
+
+ The desired probability of accepting data if it comes from class `i`:
+
+ a_i = (t_i/p_i - m) / (M_max - m)
+
+ The desired probability of pulling a data element from the original dataset,
+ rather than the filtered one:
+
+ m = M_min
+
+ Args:
+ initial_probs: A Tensor of the initial probability distribution, given or
+ estimated.
+ target_probs: A Tensor of the corresponding classes.
+
+ Returns:
+ (A 1D Tensor with the per-class acceptance probabilities, the desired
+ probability of pull from the original distribution.)
+ """
+ ratio_l = _get_target_to_initial_ratio(initial_probs, target_probs)
+ max_ratio = math_ops.reduce_max(ratio_l)
+ min_ratio = math_ops.reduce_min(ratio_l)
+
+ # Target prob to sample from original distribution.
+ m = min_ratio
+
+ # TODO(joelshor): Simplify fraction, if possible.
+ a_i = (ratio_l - m) / (max_ratio - m)
+ return a_i, m
diff --git a/tensorflow/python/data/experimental/ops/scan_ops.py b/tensorflow/python/data/experimental/ops/scan_ops.py
new file mode 100644
index 0000000..e05e7c5
--- /dev/null
+++ b/tensorflow/python/data/experimental/ops/scan_ops.py
@@ -0,0 +1,177 @@
+# 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.
+# ==============================================================================
+"""Scan dataset transformation."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import collections
+
+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.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.ops import gen_dataset_ops
+from tensorflow.python.util.tf_export import tf_export
+
+
+class _ScanDataset(dataset_ops.UnaryDataset):
+ """A dataset that scans a function across its input."""
+
+ def __init__(self, input_dataset, initial_state, scan_func):
+ """See `scan()` for details."""
+ super(_ScanDataset, self).__init__(input_dataset)
+ self._input_dataset = input_dataset
+
+ with ops.name_scope("initial_state"):
+ # Convert any `SparseTensorValue`s to `SparseTensor`s and all other
+ # values to tensors.
+ self._initial_state = nest.pack_sequence_as(initial_state, [
+ sparse_tensor.SparseTensor.from_value(t)
+ if sparse_tensor.is_sparse(t) else ops.convert_to_tensor(
+ t, name="component_%d" % i)
+ for i, t in enumerate(nest.flatten(initial_state))
+ ])
+
+ # 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
+
+ # Iteratively rerun the scan function until reaching a fixed point on
+ # `self._state_shapes`.
+ need_to_rerun = True
+ while need_to_rerun:
+
+ wrapped_func = dataset_ops.StructuredFunctionWrapper(
+ scan_func,
+ "tf.data.experimental.scan()",
+ 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),
+ add_to_graph=False)
+ if not (
+ isinstance(wrapped_func.output_types, collections.Sequence) and
+ len(wrapped_func.output_types) == 2):
+ raise TypeError("The scan function must return a pair comprising the "
+ "new state and the output value.")
+
+ 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(
+ nest.flatten(new_state_classes),
+ nest.flatten(self._state_classes)):
+ if not issubclass(new_state_class, 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))
+
+ # 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:
+ raise TypeError(
+ "The element types for the new state must match the initial "
+ "state. Expected %s; got %s." %
+ (self._state_types, new_state_types))
+
+ # Extract shape information from the returned values.
+ new_state_shapes, self._output_shapes = wrapped_func.output_shapes
+
+ flat_state_shapes = nest.flatten(self._state_shapes)
+ flat_new_state_shapes = nest.flatten(new_state_shapes)
+ weakened_state_shapes = [
+ original.most_specific_compatible_shape(new)
+ for original, new in zip(flat_state_shapes, flat_new_state_shapes)
+ ]
+
+ need_to_rerun = False
+ for original_shape, weakened_shape in zip(flat_state_shapes,
+ weakened_state_shapes):
+ if original_shape.ndims is not None and (
+ weakened_shape.ndims is None or
+ original_shape.as_list() != weakened_shape.as_list()):
+ need_to_rerun = True
+ break
+
+ if need_to_rerun:
+ self._state_shapes = nest.pack_sequence_as(self._state_shapes,
+ weakened_state_shapes)
+
+ self._scan_func = wrapped_func.function
+ self._scan_func.add_to_graph(ops.get_default_graph())
+
+ def _as_variant_tensor(self):
+ input_t = self._input_dataset._as_variant_tensor() # pylint: disable=protected-access
+ return gen_dataset_ops.scan_dataset(
+ input_t,
+ nest.flatten(sparse.serialize_sparse_tensors(self._initial_state)),
+ self._scan_func.captured_inputs,
+ f=self._scan_func,
+ **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
+
+
+@tf_export("data.experimental.scan")
+def scan(initial_state, scan_func):
+ """A transformation that scans a function across an input dataset.
+
+ This transformation is a stateful relative of `tf.data.Dataset.map`.
+ In addition to mapping `scan_func` across the elements of the input dataset,
+ `scan()` accumulates one or more state tensors, whose initial values are
+ `initial_state`.
+
+ Args:
+ initial_state: A nested structure of tensors, representing the initial state
+ of the accumulator.
+ scan_func: A function that maps `(old_state, input_element)` to
+ `(new_state, output_element). It must take two arguments and return a
+ pair of nested structures of tensors. The `new_state` must match the
+ structure of `initial_state`.
+
+ Returns:
+ A `Dataset` transformation function, which can be passed to
+ `tf.data.Dataset.apply`.
+ """
+ def _apply_fn(dataset):
+ return _ScanDataset(dataset, initial_state, scan_func)
+
+ return _apply_fn
diff --git a/tensorflow/python/data/experimental/ops/shuffle_ops.py b/tensorflow/python/data/experimental/ops/shuffle_ops.py
new file mode 100644
index 0000000..a430721
--- /dev/null
+++ b/tensorflow/python/data/experimental/ops/shuffle_ops.py
@@ -0,0 +1,102 @@
+# 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.
+# ==============================================================================
+"""Experimental shuffle ops."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.util import random_seed
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import gen_dataset_ops
+from tensorflow.python.util.tf_export import tf_export
+
+
+class _ShuffleAndRepeatDataset(dataset_ops.UnaryDataset):
+ """A `Dataset` that fuses `shuffle` and `repeat`."""
+
+ def __init__(self, input_dataset, buffer_size, count=None, seed=None):
+ super(_ShuffleAndRepeatDataset, self).__init__(input_dataset)
+ self._input_dataset = input_dataset
+ self._buffer_size = ops.convert_to_tensor(
+ buffer_size, dtype=dtypes.int64, name="buffer_size")
+ if count is None:
+ self._count = constant_op.constant(-1, dtype=dtypes.int64, name="count")
+ else:
+ self._count = ops.convert_to_tensor(
+ count, dtype=dtypes.int64, name="count")
+ self._seed, self._seed2 = random_seed.get_seed(seed)
+
+ def _as_variant_tensor(self):
+ # pylint: disable=protected-access
+ input_resource = self._input_dataset._as_variant_tensor()
+ return gen_dataset_ops.shuffle_and_repeat_dataset(
+ input_resource,
+ buffer_size=self._buffer_size,
+ count=self._count,
+ seed=self._seed,
+ seed2=self._seed2,
+ **dataset_ops.flat_structure(self))
+ # pylint: enable=protected-access
+
+ @property
+ def output_classes(self):
+ return self._input_dataset.output_classes
+
+ @property
+ def output_shapes(self):
+ return self._input_dataset.output_shapes
+
+ @property
+ def output_types(self):
+ return self._input_dataset.output_types
+
+
+@tf_export("data.experimental.shuffle_and_repeat")
+def shuffle_and_repeat(buffer_size, count=None, seed=None):
+ """Shuffles and repeats a Dataset returning a new permutation for each epoch.
+
+ `dataset.apply(tf.data.experimental.shuffle_and_repeat(buffer_size, count))`
+
+ is equivalent to
+
+ `dataset.shuffle(buffer_size, reshuffle_each_iteration=True).repeat(count)`
+
+ The difference is that the latter dataset is not serializable. So,
+ if you need to checkpoint an input pipeline with reshuffling you must use
+ this implementation.
+
+ Args:
+ buffer_size: A `tf.int64` scalar `tf.Tensor`, representing the
+ maximum number elements that will be buffered when prefetching.
+ count: (Optional.) A `tf.int64` scalar `tf.Tensor`, representing the
+ number of times the dataset should be repeated. The default behavior
+ (if `count` is `None` or `-1`) is for the dataset be repeated
+ indefinitely.
+ seed: (Optional.) A `tf.int64` scalar `tf.Tensor`, representing the
+ random seed that will be used to create the distribution. See
+ `tf.set_random_seed` for behavior.
+
+ Returns:
+ A `Dataset` transformation function, which can be passed to
+ `tf.data.Dataset.apply`.
+ """
+
+ def _apply_fn(dataset): # pylint: disable=missing-docstring
+ return _ShuffleAndRepeatDataset(dataset, buffer_size, count, seed)
+
+ return _apply_fn
diff --git a/tensorflow/contrib/data/python/ops/stats_ops.py b/tensorflow/python/data/experimental/ops/stats_ops.py
similarity index 92%
rename from tensorflow/contrib/data/python/ops/stats_ops.py
rename to tensorflow/python/data/experimental/ops/stats_ops.py
index bc47c59..c918d22 100644
--- a/tensorflow/contrib/data/python/ops/stats_ops.py
+++ b/tensorflow/python/data/experimental/ops/stats_ops.py
@@ -21,8 +21,10 @@
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import gen_dataset_ops
+from tensorflow.python.util.tf_export import tf_export
+@tf_export("data.experimental.StatsAggregator")
class StatsAggregator(object):
"""A stateful resource that aggregates statistics from one or more iterators.
@@ -34,7 +36,7 @@
```python
dataset = ...
- dataset = dataset.apply(stats_ops.latency_stats("total_bytes"))
+ dataset = dataset.apply(tf.data.experimental.latency_stats("total_bytes"))
```
To associate a `StatsAggregator` with a `tf.data.Dataset` object, use
@@ -46,7 +48,7 @@
# Apply `set_stats_aggregator` to associate `dataset` with `stats_aggregator`.
dataset = dataset.apply(
- tf.contrib.data.set_stats_aggregator(stats_aggregator))
+ tf.data.experimental.set_stats_aggregator(stats_aggregator))
iterator = dataset.make_one_shot_iterator()
```
@@ -111,11 +113,12 @@
return self._input_dataset.output_classes
+@tf_export("data.experimental.set_stats_aggregator")
def set_stats_aggregator(stats_aggregator):
"""Set the given `stats_aggregator` for aggregating the input dataset stats.
Args:
- stats_aggregator: A `tf.contrib.data.StatsAggregator` object.
+ stats_aggregator: A `tf.data.experimental.StatsAggregator` object.
Returns:
A `Dataset` transformation function, which can be passed to
@@ -128,8 +131,8 @@
return _apply_fn
-# TODO(b/38416882): Properly export in the `tf.contrib.data` API when stable
-# or make private / remove.
+# TODO(b/38416882): Properly export in the `tf.data.experimental` API when
+# stable or make private / remove.
def bytes_produced_stats(tag):
"""Records the number of bytes produced by each element of the input dataset.
@@ -152,6 +155,7 @@
return _apply_fn
+@tf_export("data.experimental.latency_stats")
def latency_stats(tag):
"""Records the latency of producing each element of the input dataset.
diff --git a/tensorflow/python/data/experimental/ops/threadpool.py b/tensorflow/python/data/experimental/ops/threadpool.py
new file mode 100644
index 0000000..3ea017c
--- /dev/null
+++ b/tensorflow/python/data/experimental/ops/threadpool.py
@@ -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.
+# ==============================================================================
+"""Experimental API for controlling threading in `tf.data` pipelines."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import threading
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.eager import context
+from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
+from tensorflow.python.ops import resource_variable_ops
+
+_uid_counter = 0
+_uid_lock = threading.Lock()
+
+
+def _generate_shared_name(prefix):
+ with _uid_lock:
+ global _uid_counter
+ uid = _uid_counter
+ _uid_counter += 1
+ return "{}{}".format(prefix, uid)
+
+
+# TODO(b/73383364): Properly export in the `tf.data.experimental` API when
+# stable or make private / remove.
+class PrivateThreadPool(object):
+ """A stateful resource that represents a private thread pool."""
+
+ def __init__(self, num_threads, display_name=None,
+ max_intra_op_parallelism=1):
+ """Creates a `PrivateThreadPool` with the given number of threads."""
+ if context.executing_eagerly():
+ shared_name = _generate_shared_name("privatethreadpool")
+ self._resource = ged_ops.experimental_thread_pool_handle(
+ num_threads=num_threads,
+ max_intra_op_parallelism=max_intra_op_parallelism,
+ display_name=display_name,
+ shared_name=shared_name)
+ self._resource_deleter = resource_variable_ops.EagerResourceDeleter(
+ handle=self._resource, handle_device=context.context().device_name)
+ else:
+ self._resource = ged_ops.experimental_thread_pool_handle(
+ num_threads=num_threads,
+ max_intra_op_parallelism=max_intra_op_parallelism,
+ display_name=display_name)
+
+
+class _ThreadPoolDataset(dataset_ops.UnaryDataset):
+ """A `Dataset` that acts as an identity, and sets a custom threadpool."""
+
+ def __init__(self, input_dataset, thread_pool):
+ super(_ThreadPoolDataset, self).__init__(input_dataset)
+ self._input_dataset = input_dataset
+ self._thread_pool = thread_pool
+
+ def _as_variant_tensor(self):
+ return ged_ops.experimental_thread_pool_dataset(
+ self._input_dataset._as_variant_tensor(), # pylint: disable=protected-access
+ self._thread_pool._resource, # pylint: disable=protected-access
+ **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
+
+
+# TODO(b/73383364): Properly export in the `tf.data.experimental` API when
+# stable or make private / remove.
+def override_threadpool(dataset, thread_pool):
+ """Returns a new dataset that uses the given thread pool for its operations.
+
+ Args:
+ dataset: A `tf.data.Dataset` object.
+ thread_pool: A `PrivateThreadPool` object.
+
+ Returns:
+ A dataset containing the same values as `dataset`, but which uses
+ `thread_pool` to compute any of its parallel operations (such as
+ `tf.data.Dataset.map`).
+ """
+ return _ThreadPoolDataset(dataset, thread_pool)
diff --git a/tensorflow/python/data/experimental/ops/unique.py b/tensorflow/python/data/experimental/ops/unique.py
new file mode 100644
index 0000000..2a7775c
--- /dev/null
+++ b/tensorflow/python/data/experimental/ops/unique.py
@@ -0,0 +1,79 @@
+# 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.
+# ==============================================================================
+"""Unique element dataset transformations."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import gen_experimental_dataset_ops
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("data.experimental.unique")
+def unique():
+ """Creates a `Dataset` from another `Dataset`, discarding duplicates.
+
+ Use this transformation to produce a dataset that contains one instance of
+ each unique element in the input. For example:
+
+ ```python
+ dataset = tf.data.Dataset.from_tensor_slices([1, 37, 2, 37, 2, 1])
+
+ # Using `unique()` will drop the duplicate elements.
+ dataset = dataset.apply(tf.data.experimental.unique()) # ==> { 1, 37, 2 }
+ ```
+
+ Returns:
+ A `Dataset` transformation function, which can be passed to
+ `tf.data.Dataset.apply`.
+ """
+
+ def _apply_fn(dataset):
+ return _UniqueDataset(dataset)
+
+ return _apply_fn
+
+
+class _UniqueDataset(dataset_ops.UnaryDataset):
+ """A `Dataset` contains the unique elements from its input."""
+
+ def __init__(self, input_dataset):
+ """See `unique()` for details."""
+ super(_UniqueDataset, self).__init__(input_dataset)
+ self._input_dataset = input_dataset
+ if input_dataset.output_types not in (dtypes.int32, dtypes.int64,
+ dtypes.string):
+ raise TypeError(
+ "`tf.data.experimental.unique()` only supports inputs with a single "
+ "`tf.int32`, `tf.int64`, or `tf.string` component.")
+
+ def _as_variant_tensor(self):
+ return gen_experimental_dataset_ops.experimental_unique_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 self._input_dataset.output_shapes
+
+ @property
+ def output_types(self):
+ return self._input_dataset.output_types
diff --git a/tensorflow/python/data/experimental/ops/writers.py b/tensorflow/python/data/experimental/ops/writers.py
new file mode 100644
index 0000000..994447c
--- /dev/null
+++ b/tensorflow/python/data/experimental/ops/writers.py
@@ -0,0 +1,60 @@
+# 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.
+# ==============================================================================
+"""Python wrappers for tf.data writers."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.util import convert
+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_dataset_ops
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("data.experimental.TFRecordWriter")
+class TFRecordWriter(object):
+ """Writes data to a TFRecord file."""
+
+ def __init__(self, filename, compression_type=None):
+ self._filename = ops.convert_to_tensor(
+ filename, dtypes.string, name="filename")
+ self._compression_type = convert.optional_param_to_tensor(
+ "compression_type",
+ compression_type,
+ argument_default="",
+ argument_dtype=dtypes.string)
+
+ def write(self, dataset):
+ """Returns a `tf.Operation` to write a dataset to a file.
+
+ Args:
+ dataset: a `tf.data.Dataset` whose elements are to be written to a file
+
+ Returns:
+ A `tf.Operation` that, when run, writes contents of `dataset` to a file.
+ """
+ if not isinstance(dataset, dataset_ops.Dataset):
+ raise TypeError("`dataset` must be a `tf.data.Dataset` object.")
+ if (dataset.output_types != dtypes.string or
+ dataset.output_shapes != tensor_shape.scalar()):
+ raise TypeError(
+ "`dataset` must produce scalar `DT_STRING` tensors whereas it "
+ "produces shape {0} and types {1}".format(dataset.output_shapes,
+ dataset.output_types))
+ return gen_dataset_ops.dataset_to_tf_record(
+ dataset._as_variant_tensor(), self._filename, self._compression_type) # pylint: disable=protected-access
diff --git a/tensorflow/python/data/kernel_tests/BUILD b/tensorflow/python/data/kernel_tests/BUILD
index 7a6f03d..cadfe7f 100644
--- a/tensorflow/python/data/kernel_tests/BUILD
+++ b/tensorflow/python/data/kernel_tests/BUILD
@@ -15,6 +15,7 @@
size = "small",
srcs = ["batch_dataset_op_test.py"],
additional_deps = [
+ ":test_base",
"@absl_py//absl/testing:parameterized",
"//third_party/py/numpy",
"//tensorflow/python:array_ops",
@@ -31,10 +32,44 @@
)
tf_py_test(
+ name = "cache_dataset_op_test",
+ size = "small",
+ srcs = ["cache_dataset_op_test.py"],
+ additional_deps = [
+ ":test_base",
+ "//third_party/py/numpy",
+ "//tensorflow/python:array_ops",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:constant_op",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:errors",
+ "//tensorflow/python:variables",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/ops:iterator_ops",
+ ],
+)
+
+tf_py_test(
+ name = "concatenate_dataset_op_test",
+ size = "small",
+ srcs = ["concatenate_dataset_op_test.py"],
+ additional_deps = [
+ ":test_base",
+ "//third_party/py/numpy",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:errors",
+ "//tensorflow/python:tensor_shape",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/util:nest",
+ ],
+)
+
+tf_py_test(
name = "dataset_constructor_op_test",
size = "small",
srcs = ["dataset_constructor_op_test.py"],
additional_deps = [
+ ":test_base",
"//third_party/py/numpy",
"//tensorflow/core:protos_all_py",
"//tensorflow/python:array_ops",
@@ -63,6 +98,7 @@
size = "medium",
srcs = ["dataset_from_generator_op_test.py"],
additional_deps = [
+ ":test_base",
"//third_party/py/numpy",
"//tensorflow/python:client_testlib",
"//tensorflow/python:dtypes",
@@ -78,6 +114,7 @@
size = "small",
srcs = ["dataset_ops_test.py"],
additional_deps = [
+ ":test_base",
"//tensorflow/core:protos_all_py",
"//tensorflow/python:client_testlib",
"//tensorflow/python/data/ops:dataset_ops",
@@ -89,6 +126,7 @@
size = "small",
srcs = ["filter_dataset_op_test.py"],
additional_deps = [
+ ":test_base",
"//third_party/py/numpy",
"//tensorflow/python:array_ops",
"//tensorflow/python:client_testlib",
@@ -106,6 +144,7 @@
size = "small",
srcs = ["flat_map_dataset_op_test.py"],
additional_deps = [
+ ":test_base",
"//third_party/py/numpy",
"//tensorflow/python:client_testlib",
"//tensorflow/python:errors",
@@ -123,6 +162,7 @@
size = "small",
srcs = ["list_files_dataset_op_test.py"],
additional_deps = [
+ ":test_base",
"//tensorflow/python:array_ops",
"//tensorflow/python:client_testlib",
"//tensorflow/python:dtypes",
@@ -133,10 +173,25 @@
)
tf_py_test(
+ name = "inputs_test",
+ size = "small",
+ srcs = ["inputs_test.py"],
+ additional_deps = [
+ ":test_base",
+ "@absl_py//absl/testing:parameterized",
+ "//third_party/py/numpy",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:sparse_tensor",
+ "//tensorflow/python/data/ops:dataset_ops",
+ ],
+)
+
+tf_py_test(
name = "interleave_dataset_op_test",
size = "small",
srcs = ["interleave_dataset_op_test.py"],
additional_deps = [
+ ":test_base",
"@absl_py//absl/testing:parameterized",
"//third_party/py/numpy",
"//tensorflow/python:array_ops",
@@ -151,176 +206,6 @@
],
)
-tf_py_test(
- name = "map_dataset_op_test",
- size = "small",
- srcs = ["map_dataset_op_test.py"],
- additional_deps = [
- "@absl_py//absl/testing:parameterized",
- "//third_party/py/numpy",
- "//tensorflow/python:array_ops",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:constant_op",
- "//tensorflow/python:data_flow_ops",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:errors",
- "//tensorflow/python:functional_ops",
- "//tensorflow/python:lookup_ops",
- "//tensorflow/python:math_ops",
- "//tensorflow/python:random_ops",
- "//tensorflow/python:script_ops",
- "//tensorflow/python:sparse_ops",
- "//tensorflow/python:sparse_tensor",
- "//tensorflow/python:string_ops",
- "//tensorflow/python:variable_scope",
- "//tensorflow/python/data/ops:dataset_ops",
- ],
-)
-
-tf_py_test(
- name = "prefetch_dataset_op_test",
- size = "small",
- srcs = ["prefetch_dataset_op_test.py"],
- additional_deps = [
- "@absl_py//absl/testing:parameterized",
- "//tensorflow/python:array_ops",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:dataset_ops_gen",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:errors",
- "//tensorflow/python/data/ops:dataset_ops",
- ],
-)
-
-tf_py_test(
- name = "range_dataset_op_test",
- size = "small",
- srcs = ["range_dataset_op_test.py"],
- additional_deps = [
- "//tensorflow/python:array_ops",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:dataset_ops_gen",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:errors",
- "//tensorflow/python:io_ops",
- "//tensorflow/python:framework_ops",
- "//tensorflow/python:parsing_ops",
- "//tensorflow/python:platform",
- "//tensorflow/python:tensor_shape",
- "//tensorflow/python:variables",
- "//tensorflow/python/data/ops:dataset_ops",
- "//tensorflow/python/data/ops:iterator_ops",
- ],
-)
-
-tf_py_test(
- name = "reader_dataset_ops_test",
- size = "small",
- srcs = ["reader_dataset_ops_test.py"],
- additional_deps = [
- "//tensorflow/python:array_ops",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:constant_op",
- "//tensorflow/python:dataset_ops_gen",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:errors",
- "//tensorflow/python:framework_ops",
- "//tensorflow/python:io_ops",
- "//tensorflow/python:lib",
- "//tensorflow/python:parsing_ops",
- "//tensorflow/python:tensor_shape",
- "//tensorflow/python:util",
- "//tensorflow/python/data/ops:iterator_ops",
- "//tensorflow/python/data/ops:readers",
- ],
-)
-
-tf_py_test(
- name = "sequence_dataset_op_test",
- size = "small",
- srcs = ["sequence_dataset_op_test.py"],
- additional_deps = [
- "//third_party/py/numpy",
- "//tensorflow/python:array_ops",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:errors",
- "//tensorflow/python/data/ops:dataset_ops",
- ],
-)
-
-tf_py_test(
- name = "shuffle_dataset_op_test",
- size = "small",
- srcs = ["shuffle_dataset_op_test.py"],
- additional_deps = [
- "//third_party/py/numpy",
- "//tensorflow/python:array_ops",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:constant_op",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:errors",
- "//tensorflow/python/data/ops:dataset_ops",
- "//tensorflow/python/data/ops:iterator_ops",
- ],
-)
-
-tf_py_test(
- name = "shard_dataset_op_test",
- size = "small",
- srcs = ["shard_dataset_op_test.py"],
- additional_deps = [
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:errors",
- "//tensorflow/python/data/ops:dataset_ops",
- ],
-)
-
-tf_py_test(
- name = "cache_dataset_op_test",
- size = "small",
- srcs = ["cache_dataset_op_test.py"],
- additional_deps = [
- "//third_party/py/numpy",
- "//tensorflow/python:array_ops",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:constant_op",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:errors",
- "//tensorflow/python:variables",
- "//tensorflow/python/data/ops:dataset_ops",
- "//tensorflow/python/data/ops:iterator_ops",
- ],
-)
-
-tf_py_test(
- name = "zip_dataset_op_test",
- size = "small",
- srcs = ["zip_dataset_op_test.py"],
- additional_deps = [
- "//third_party/py/numpy",
- "//tensorflow/python:array_ops",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:errors",
- "//tensorflow/python/data/ops:dataset_ops",
- ],
-)
-
-tf_py_test(
- name = "concatenate_dataset_op_test",
- size = "small",
- srcs = ["concatenate_dataset_op_test.py"],
- additional_deps = [
- "//third_party/py/numpy",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:errors",
- "//tensorflow/python:tensor_shape",
- "//tensorflow/python/data/ops:dataset_ops",
- "//tensorflow/python/data/util:nest",
- ],
-)
-
cuda_py_test(
name = "iterator_ops_test",
size = "small",
@@ -389,23 +274,30 @@
],
)
-cuda_py_test(
- name = "optional_ops_test",
+tf_py_test(
+ name = "map_dataset_op_test",
size = "small",
- srcs = ["optional_ops_test.py"],
+ srcs = ["map_dataset_op_test.py"],
additional_deps = [
+ ":test_base",
"@absl_py//absl/testing:parameterized",
- "//tensorflow/python/data/ops:dataset_ops",
- "//tensorflow/python/data/ops:iterator_ops",
- "//tensorflow/python/data/ops:optional_ops",
- "//tensorflow/python:client_testlib",
+ "//third_party/py/numpy",
"//tensorflow/python:array_ops",
+ "//tensorflow/python:client_testlib",
"//tensorflow/python:constant_op",
+ "//tensorflow/python:data_flow_ops",
"//tensorflow/python:dtypes",
"//tensorflow/python:errors",
- "//tensorflow/python:framework_ops",
- "//tensorflow/python:framework_test_lib",
- "//tensorflow/python:tensor_shape",
+ "//tensorflow/python:functional_ops",
+ "//tensorflow/python:lookup_ops",
+ "//tensorflow/python:math_ops",
+ "//tensorflow/python:random_ops",
+ "//tensorflow/python:script_ops",
+ "//tensorflow/python:sparse_ops",
+ "//tensorflow/python:sparse_tensor",
+ "//tensorflow/python:string_ops",
+ "//tensorflow/python:variable_scope",
+ "//tensorflow/python/data/ops:dataset_ops",
],
)
@@ -414,6 +306,7 @@
size = "small",
srcs = ["multi_device_iterator_test.py"],
additional_deps = [
+ ":test_base",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/ops:multi_device_iterator_ops",
@@ -429,11 +322,94 @@
],
)
-tf_py_test(
- name = "window_dataset_op_test",
+cuda_py_test(
+ name = "optional_ops_test",
size = "small",
- srcs = ["window_dataset_op_test.py"],
+ srcs = ["optional_ops_test.py"],
additional_deps = [
+ ":test_base",
+ "@absl_py//absl/testing:parameterized",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/ops:iterator_ops",
+ "//tensorflow/python/data/ops:optional_ops",
+ "//tensorflow/python:client_testlib",
+ "//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:tensor_shape",
+ ],
+)
+
+tf_py_test(
+ name = "prefetch_dataset_op_test",
+ size = "small",
+ srcs = ["prefetch_dataset_op_test.py"],
+ additional_deps = [
+ ":test_base",
+ "@absl_py//absl/testing:parameterized",
+ "//tensorflow/python:array_ops",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:dataset_ops_gen",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:errors",
+ "//tensorflow/python/data/ops:dataset_ops",
+ ],
+)
+
+tf_py_test(
+ name = "range_dataset_op_test",
+ size = "small",
+ srcs = ["range_dataset_op_test.py"],
+ additional_deps = [
+ ":test_base",
+ "//tensorflow/python:array_ops",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:dataset_ops_gen",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:errors",
+ "//tensorflow/python:io_ops",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python:parsing_ops",
+ "//tensorflow/python:platform",
+ "//tensorflow/python:tensor_shape",
+ "//tensorflow/python:variables",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/ops:iterator_ops",
+ ],
+)
+
+tf_py_test(
+ name = "reader_dataset_ops_test",
+ size = "small",
+ srcs = ["reader_dataset_ops_test.py"],
+ additional_deps = [
+ ":test_base",
+ "//tensorflow/python:array_ops",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:constant_op",
+ "//tensorflow/python:dataset_ops_gen",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:errors",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python:io_ops",
+ "//tensorflow/python:lib",
+ "//tensorflow/python:parsing_ops",
+ "//tensorflow/python:tensor_shape",
+ "//tensorflow/python:util",
+ "//tensorflow/python/data/ops:iterator_ops",
+ "//tensorflow/python/data/ops:readers",
+ ],
+)
+
+tf_py_test(
+ name = "reduce_dataset_op_test",
+ size = "small",
+ srcs = ["reduce_dataset_op_test.py"],
+ additional_deps = [
+ ":test_base",
"@absl_py//absl/testing:parameterized",
"//third_party/py/numpy",
"//tensorflow/python:array_ops",
@@ -447,14 +423,89 @@
)
tf_py_test(
- name = "inputs_test",
+ name = "sequence_dataset_op_test",
size = "small",
- srcs = ["inputs_test.py"],
+ srcs = ["sequence_dataset_op_test.py"],
additional_deps = [
+ ":test_base",
+ "//third_party/py/numpy",
+ "//tensorflow/python:array_ops",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:errors",
+ "//tensorflow/python/data/ops:dataset_ops",
+ ],
+)
+
+tf_py_test(
+ name = "shard_dataset_op_test",
+ size = "small",
+ srcs = ["shard_dataset_op_test.py"],
+ additional_deps = [
+ ":test_base",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:errors",
+ "//tensorflow/python/data/ops:dataset_ops",
+ ],
+)
+
+tf_py_test(
+ name = "shuffle_dataset_op_test",
+ size = "small",
+ srcs = ["shuffle_dataset_op_test.py"],
+ additional_deps = [
+ ":test_base",
+ "//third_party/py/numpy",
+ "//tensorflow/python:array_ops",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:constant_op",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:errors",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/ops:iterator_ops",
+ ],
+)
+
+py_library(
+ name = "test_base",
+ srcs = ["test_base.py"],
+ deps = [
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:errors",
+ "//tensorflow/python:sparse_tensor",
+ "//tensorflow/python/data/util:nest",
+ ],
+)
+
+tf_py_test(
+ name = "window_dataset_op_test",
+ size = "small",
+ srcs = ["window_dataset_op_test.py"],
+ additional_deps = [
+ ":test_base",
"@absl_py//absl/testing:parameterized",
"//third_party/py/numpy",
+ "//tensorflow/python:array_ops",
"//tensorflow/python:client_testlib",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:errors",
+ "//tensorflow/python:math_ops",
"//tensorflow/python:sparse_tensor",
"//tensorflow/python/data/ops:dataset_ops",
],
)
+
+tf_py_test(
+ name = "zip_dataset_op_test",
+ size = "small",
+ srcs = ["zip_dataset_op_test.py"],
+ additional_deps = [
+ ":test_base",
+ "//third_party/py/numpy",
+ "//tensorflow/python:array_ops",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:errors",
+ "//tensorflow/python/data/ops:dataset_ops",
+ ],
+)
diff --git a/tensorflow/python/data/kernel_tests/batch_dataset_op_test.py b/tensorflow/python/data/kernel_tests/batch_dataset_op_test.py
index c48708a..9cb4daf 100644
--- a/tensorflow/python/data/kernel_tests/batch_dataset_op_test.py
+++ b/tensorflow/python/data/kernel_tests/batch_dataset_op_test.py
@@ -24,6 +24,7 @@
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
from tensorflow.python.framework import dtypes
@@ -37,7 +38,7 @@
from tensorflow.python.util import compat
-class BatchDatasetTest(test.TestCase, parameterized.TestCase):
+class BatchDatasetTest(test_base.DatasetTestBase, parameterized.TestCase):
@parameterized.named_parameters(
('even', 28, 14, False),
@@ -115,11 +116,6 @@
with self.assertRaises(errors.InvalidArgumentError):
sess.run(get_next)
- def assertSparseValuesEqual(self, a, b):
- self.assertAllEqual(a.indices, b.indices)
- self.assertAllEqual(a.values, b.values)
- self.assertAllEqual(a.dense_shape, b.dense_shape)
-
def testBatchSparse(self):
def _sparse(i):
@@ -227,7 +223,7 @@
return np.random.randint(20, size=(count,)).astype(np.int32)
-class PaddedBatchDatasetTest(test.TestCase, parameterized.TestCase):
+class PaddedBatchDatasetTest(test_base.DatasetTestBase, parameterized.TestCase):
@parameterized.named_parameters(
('default_padding', _random_seq_lens(32), 4, [-1], False),
diff --git a/tensorflow/python/data/kernel_tests/cache_dataset_op_test.py b/tensorflow/python/data/kernel_tests/cache_dataset_op_test.py
index d5f5b2f..63625fa 100644
--- a/tensorflow/python/data/kernel_tests/cache_dataset_op_test.py
+++ b/tensorflow/python/data/kernel_tests/cache_dataset_op_test.py
@@ -23,6 +23,7 @@
import numpy as np
+from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import iterator_ops
from tensorflow.python.framework import constant_op
@@ -34,7 +35,7 @@
from tensorflow.python.platform import test
-class FileCacheDatasetTest(test.TestCase):
+class FileCacheDatasetTest(test_base.DatasetTestBase):
def setUp(self):
self.tmp_dir = tempfile.mkdtemp()
@@ -200,7 +201,7 @@
self.assertAllEqual(elements, elements_itr2)
-class MemoryCacheDatasetTest(test.TestCase):
+class MemoryCacheDatasetTest(test_base.DatasetTestBase):
def testCacheDatasetPassthrough(self):
with ops.device("cpu:0"):
diff --git a/tensorflow/python/data/kernel_tests/concatenate_dataset_op_test.py b/tensorflow/python/data/kernel_tests/concatenate_dataset_op_test.py
index 5dfb84f..83af31f 100644
--- a/tensorflow/python/data/kernel_tests/concatenate_dataset_op_test.py
+++ b/tensorflow/python/data/kernel_tests/concatenate_dataset_op_test.py
@@ -19,6 +19,7 @@
import numpy as np
+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 errors
@@ -26,7 +27,7 @@
from tensorflow.python.platform import test
-class ConcatenateDatasetTest(test.TestCase):
+class ConcatenateDatasetTest(test_base.DatasetTestBase):
def testConcatenateDataset(self):
input_components = (
diff --git a/tensorflow/python/data/kernel_tests/dataset_constructor_op_test.py b/tensorflow/python/data/kernel_tests/dataset_constructor_op_test.py
index e43564a..bc6b362 100644
--- a/tensorflow/python/data/kernel_tests/dataset_constructor_op_test.py
+++ b/tensorflow/python/data/kernel_tests/dataset_constructor_op_test.py
@@ -23,6 +23,7 @@
from tensorflow.core.protobuf import config_pb2
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.data.util import nest
from tensorflow.python.framework import dtypes
@@ -36,7 +37,7 @@
from tensorflow.python.platform import test
-class DatasetConstructorTest(test.TestCase):
+class DatasetConstructorTest(test_base.DatasetTestBase):
def testFromTensors(self):
"""Test a dataset that represents a single tuple of tensors."""
@@ -58,11 +59,6 @@
with self.assertRaises(errors.OutOfRangeError):
sess.run(get_next)
- def assertSparseValuesEqual(self, a, b):
- self.assertAllEqual(a.indices, b.indices)
- self.assertAllEqual(a.values, b.values)
- self.assertAllEqual(a.dense_shape, b.dense_shape)
-
def testFromTensorsSparse(self):
"""Test a dataset that represents a single tuple of tensors."""
components = (sparse_tensor.SparseTensorValue(
diff --git a/tensorflow/python/data/kernel_tests/dataset_from_generator_op_test.py b/tensorflow/python/data/kernel_tests/dataset_from_generator_op_test.py
index cd0c1dd..cb8cb9a 100644
--- a/tensorflow/python/data/kernel_tests/dataset_from_generator_op_test.py
+++ b/tensorflow/python/data/kernel_tests/dataset_from_generator_op_test.py
@@ -22,6 +22,7 @@
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
from tensorflow.python.framework import dtypes
@@ -30,7 +31,7 @@
from tensorflow.python.platform import test
-class DatasetConstructorTest(test.TestCase):
+class DatasetConstructorTest(test_base.DatasetTestBase):
def _testFromGenerator(self, generator, elem_sequence, num_repeats,
output_types=None):
diff --git a/tensorflow/python/data/kernel_tests/dataset_ops_test.py b/tensorflow/python/data/kernel_tests/dataset_ops_test.py
index 239aa85..f115f9d 100644
--- a/tensorflow/python/data/kernel_tests/dataset_ops_test.py
+++ b/tensorflow/python/data/kernel_tests/dataset_ops_test.py
@@ -19,11 +19,12 @@
from __future__ import print_function
from tensorflow.core.framework import graph_pb2
+from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.platform import test
-class DatasetOpsTest(test.TestCase):
+class DatasetOpsTest(test_base.DatasetTestBase):
def testAsSerializedGraph(self):
dataset = dataset_ops.Dataset.range(10)
diff --git a/tensorflow/python/data/kernel_tests/filter_dataset_op_test.py b/tensorflow/python/data/kernel_tests/filter_dataset_op_test.py
index 19944d3..6b7afaf 100644
--- a/tensorflow/python/data/kernel_tests/filter_dataset_op_test.py
+++ b/tensorflow/python/data/kernel_tests/filter_dataset_op_test.py
@@ -22,6 +22,7 @@
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 dtypes
from tensorflow.python.framework import errors
@@ -33,7 +34,7 @@
from tensorflow.python.platform import test
-class FilterDatasetTest(test.TestCase):
+class FilterDatasetTest(test_base.DatasetTestBase):
def testFilterDataset(self):
components = (
@@ -129,11 +130,6 @@
with self.assertRaises(errors.OutOfRangeError):
sess.run(get_next)
- def assertSparseValuesEqual(self, a, b):
- self.assertAllEqual(a.indices, b.indices)
- self.assertAllEqual(a.values, b.values)
- self.assertAllEqual(a.dense_shape, b.dense_shape)
-
def testSparse(self):
def _map_fn(i):
diff --git a/tensorflow/python/data/kernel_tests/flat_map_dataset_op_test.py b/tensorflow/python/data/kernel_tests/flat_map_dataset_op_test.py
index 1123cbf..68038f9 100644
--- a/tensorflow/python/data/kernel_tests/flat_map_dataset_op_test.py
+++ b/tensorflow/python/data/kernel_tests/flat_map_dataset_op_test.py
@@ -22,6 +22,7 @@
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 errors
from tensorflow.python.framework import sparse_tensor
@@ -30,7 +31,7 @@
from tensorflow.python.training import server_lib
-class FlatMapDatasetTest(test.TestCase):
+class FlatMapDatasetTest(test_base.DatasetTestBase):
# pylint: disable=g-long-lambda
def testFlatMapDataset(self):
diff --git a/tensorflow/python/data/kernel_tests/inputs_test.py b/tensorflow/python/data/kernel_tests/inputs_test.py
index 4c9279d..d089b49 100644
--- a/tensorflow/python/data/kernel_tests/inputs_test.py
+++ b/tensorflow/python/data/kernel_tests/inputs_test.py
@@ -19,6 +19,7 @@
from absl.testing import parameterized
import numpy as np
+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.data.util import nest
@@ -27,7 +28,7 @@
from tensorflow.python.platform import test
-class InputsTest(test.TestCase, parameterized.TestCase):
+class InputsTest(test_base.DatasetTestBase, parameterized.TestCase):
@staticmethod
def make_apply_fn(dataset):
diff --git a/tensorflow/python/data/kernel_tests/interleave_dataset_op_test.py b/tensorflow/python/data/kernel_tests/interleave_dataset_op_test.py
index e7e51df..92bb67b 100644
--- a/tensorflow/python/data/kernel_tests/interleave_dataset_op_test.py
+++ b/tensorflow/python/data/kernel_tests/interleave_dataset_op_test.py
@@ -22,6 +22,7 @@
from absl.testing import parameterized
import numpy as np
+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 sparse_tensor
@@ -30,7 +31,7 @@
from tensorflow.python.platform import test
-class InterleaveDatasetTest(test.TestCase, parameterized.TestCase):
+class InterleaveDatasetTest(test_base.DatasetTestBase, parameterized.TestCase):
def _interleave(self, lists, cycle_length, block_length):
num_open = 0
diff --git a/tensorflow/python/data/kernel_tests/list_files_dataset_op_test.py b/tensorflow/python/data/kernel_tests/list_files_dataset_op_test.py
index c4b338a..8eb1381 100644
--- a/tensorflow/python/data/kernel_tests/list_files_dataset_op_test.py
+++ b/tensorflow/python/data/kernel_tests/list_files_dataset_op_test.py
@@ -22,6 +22,7 @@
import shutil
import tempfile
+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
@@ -30,7 +31,7 @@
from tensorflow.python.util import compat
-class ListFilesDatasetOpTest(test.TestCase):
+class ListFilesDatasetOpTest(test_base.DatasetTestBase):
def setUp(self):
self.tmp_dir = tempfile.mkdtemp()
diff --git a/tensorflow/python/data/kernel_tests/map_dataset_op_test.py b/tensorflow/python/data/kernel_tests/map_dataset_op_test.py
index ae04995..230ae3f 100644
--- a/tensorflow/python/data/kernel_tests/map_dataset_op_test.py
+++ b/tensorflow/python/data/kernel_tests/map_dataset_op_test.py
@@ -27,6 +27,7 @@
from tensorflow.core.framework import attr_value_pb2
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
from tensorflow.python.framework import dtypes
@@ -47,7 +48,7 @@
from tensorflow.python.platform import test
-class MapDatasetTest(test.TestCase, parameterized.TestCase):
+class MapDatasetTest(test_base.DatasetTestBase, parameterized.TestCase):
def _buildMapDataset(self, components, count):
def _map_fn(x, y, z):
@@ -574,11 +575,6 @@
with self.assertRaises(errors.OutOfRangeError):
sess.run(get_next)
- def assertSparseValuesEqual(self, a, b):
- self.assertAllEqual(a.indices, b.indices)
- self.assertAllEqual(a.values, b.values)
- self.assertAllEqual(a.dense_shape, b.dense_shape)
-
def testSparse(self):
def _sparse(i):
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 056664b..1cf6dd1 100644
--- a/tensorflow/python/data/kernel_tests/multi_device_iterator_test.py
+++ b/tensorflow/python/data/kernel_tests/multi_device_iterator_test.py
@@ -19,6 +19,7 @@
from __future__ import print_function
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.data.ops import multi_device_iterator_ops
from tensorflow.python.framework import dtypes
@@ -29,7 +30,7 @@
from tensorflow.python.platform import test
-class MultiDeviceIteratorTest(test.TestCase):
+class MultiDeviceIteratorTest(test_base.DatasetTestBase):
def testNoGetNext(self):
dataset = dataset_ops.Dataset.range(10)
diff --git a/tensorflow/python/data/kernel_tests/optional_ops_test.py b/tensorflow/python/data/kernel_tests/optional_ops_test.py
index 706a65f..604e3ad 100644
--- a/tensorflow/python/data/kernel_tests/optional_ops_test.py
+++ b/tensorflow/python/data/kernel_tests/optional_ops_test.py
@@ -20,6 +20,7 @@
from absl.testing import parameterized
import numpy as np
+from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import iterator_ops
from tensorflow.python.data.ops import optional_ops
@@ -35,7 +36,7 @@
from tensorflow.python.platform import test
-class OptionalTest(test.TestCase, parameterized.TestCase):
+class OptionalTest(test_base.DatasetTestBase, parameterized.TestCase):
@test_util.run_in_graph_and_eager_modes
def testFromValue(self):
diff --git a/tensorflow/python/data/kernel_tests/prefetch_dataset_op_test.py b/tensorflow/python/data/kernel_tests/prefetch_dataset_op_test.py
index cc97bac..76e2697 100644
--- a/tensorflow/python/data/kernel_tests/prefetch_dataset_op_test.py
+++ b/tensorflow/python/data/kernel_tests/prefetch_dataset_op_test.py
@@ -19,6 +19,7 @@
from absl.testing import parameterized
+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
@@ -26,7 +27,7 @@
from tensorflow.python.platform import test
-class PrefetchDatasetTest(test.TestCase, parameterized.TestCase):
+class PrefetchDatasetTest(test_base.DatasetTestBase, parameterized.TestCase):
@parameterized.parameters((-1), (0), (5))
def testBufferSize(self, buffer_size):
diff --git a/tensorflow/python/data/kernel_tests/range_dataset_op_test.py b/tensorflow/python/data/kernel_tests/range_dataset_op_test.py
index 51e9078..b7e2a5f 100644
--- a/tensorflow/python/data/kernel_tests/range_dataset_op_test.py
+++ b/tensorflow/python/data/kernel_tests/range_dataset_op_test.py
@@ -19,6 +19,7 @@
import os
+from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import iterator_ops
from tensorflow.python.framework import dtypes
@@ -34,7 +35,7 @@
from tensorflow.python.platform import test
-class RangeDatasetTest(test.TestCase):
+class RangeDatasetTest(test_base.DatasetTestBase):
def tearDown(self):
# Remove all checkpoint files.
diff --git a/tensorflow/python/data/kernel_tests/reader_dataset_ops_test.py b/tensorflow/python/data/kernel_tests/reader_dataset_ops_test.py
index aa36363..aef2dd1 100644
--- a/tensorflow/python/data/kernel_tests/reader_dataset_ops_test.py
+++ b/tensorflow/python/data/kernel_tests/reader_dataset_ops_test.py
@@ -21,6 +21,7 @@
import os
import zlib
+from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import iterator_ops
from tensorflow.python.data.ops import readers
@@ -46,7 +47,7 @@
psutil_import_succeeded = False
-class TextLineDatasetTest(test.TestCase):
+class TextLineDatasetTest(test_base.DatasetTestBase):
def _lineText(self, f, l):
return compat.as_bytes("%d: %d" % (f, l))
@@ -199,7 +200,7 @@
self.assertNotIn(filename, [open_file.path for open_file in open_files])
-class FixedLengthRecordReaderTest(test.TestCase):
+class FixedLengthRecordReaderTest(test_base.DatasetTestBase):
def setUp(self):
super(FixedLengthRecordReaderTest, self).setUp()
@@ -621,7 +622,7 @@
sess.run(get_next_op)
-class TFRecordDatasetTest(test.TestCase):
+class TFRecordDatasetTest(test_base.DatasetTestBase):
def setUp(self):
super(TFRecordDatasetTest, self).setUp()
diff --git a/tensorflow/python/data/kernel_tests/reduce_dataset_op_test.py b/tensorflow/python/data/kernel_tests/reduce_dataset_op_test.py
new file mode 100644
index 0000000..11e0730
--- /dev/null
+++ b/tensorflow/python/data/kernel_tests/reduce_dataset_op_test.py
@@ -0,0 +1,124 @@
+# 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 experimental input pipeline 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.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 sparse_tensor
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+class ReduceDatasetTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+ def testSum(self):
+ for i in range(10):
+ ds = dataset_ops.Dataset.range(1, i + 1)
+ result = ds.reduce(np.int64(0), lambda x, y: x + y)
+ with self.cached_session() as sess:
+ self.assertEqual(((i + 1) * i) // 2, sess.run(result))
+
+ def testSumTuple(self):
+
+ def reduce_fn(state, value):
+ v1, v2 = value
+ return state + v1 + v2
+
+ for i in range(10):
+ ds = dataset_ops.Dataset.range(1, i + 1)
+ ds = dataset_ops.Dataset.zip((ds, ds))
+ result = ds.reduce(np.int64(0), reduce_fn)
+ with self.cached_session() as sess:
+ self.assertEqual(((i + 1) * i), sess.run(result))
+
+ def testSumAndCount(self):
+
+ def reduce_fn(state, value):
+ s, c = state
+ return s + value, c + 1
+
+ for i in range(10):
+ ds = dataset_ops.Dataset.range(1, i + 1)
+ result = ds.reduce((np.int64(0), np.int64(0)), reduce_fn)
+ with self.cached_session() as sess:
+ s, c = sess.run(result)
+ self.assertEqual(((i + 1) * i) // 2, s)
+ self.assertEqual(i, c)
+
+ def testSquareUsingPlaceholder(self):
+ delta = array_ops.placeholder(dtype=dtypes.int64)
+
+ def reduce_fn(state, _):
+ return state + delta
+
+ for i in range(10):
+ ds = dataset_ops.Dataset.range(1, i + 1)
+ result = ds.reduce(np.int64(0), reduce_fn)
+ with self.cached_session() as sess:
+ square = sess.run(result, feed_dict={delta: i})
+ self.assertEqual(i * i, square)
+
+ def testSparse(self):
+
+ def reduce_fn(_, value):
+ return value
+
+ def make_sparse_fn(i):
+ return sparse_tensor.SparseTensorValue(
+ indices=np.array([[0, 0]]),
+ values=(i * np.array([1])),
+ dense_shape=np.array([1, 1]))
+
+ for i in range(10):
+ ds = dataset_ops.Dataset.from_tensors(make_sparse_fn(i+1))
+ result = ds.reduce(make_sparse_fn(0), reduce_fn)
+ with self.cached_session() as sess:
+ self.assertSparseValuesEqual(make_sparse_fn(i+1), sess.run(result))
+
+ def testNested(self):
+
+ def reduce_fn(state, value):
+ state["dense"] += value["dense"]
+ state["sparse"] = value["sparse"]
+ return state
+
+ def make_sparse_fn(i):
+ return sparse_tensor.SparseTensorValue(
+ indices=np.array([[0, 0]]),
+ values=(i * np.array([1])),
+ dense_shape=np.array([1, 1]))
+
+ def map_fn(i):
+ return {"dense": math_ops.cast(i, dtype=dtypes.int64),
+ "sparse": make_sparse_fn(math_ops.cast(i, dtype=dtypes.int64))}
+
+ for i in range(10):
+ ds = dataset_ops.Dataset.range(1, i + 1).map(map_fn)
+ result = ds.reduce(map_fn(0), reduce_fn)
+ with self.cached_session() as sess:
+ result = sess.run(result)
+ self.assertEqual(((i + 1) * i) // 2, result["dense"])
+ self.assertSparseValuesEqual(make_sparse_fn(i), result["sparse"])
+
+if __name__ == "__main__":
+ test.main()
diff --git a/tensorflow/python/data/kernel_tests/sequence_dataset_op_test.py b/tensorflow/python/data/kernel_tests/sequence_dataset_op_test.py
index 37e2333..e86356d 100644
--- a/tensorflow/python/data/kernel_tests/sequence_dataset_op_test.py
+++ b/tensorflow/python/data/kernel_tests/sequence_dataset_op_test.py
@@ -19,6 +19,7 @@
import numpy as np
+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
@@ -26,7 +27,7 @@
from tensorflow.python.platform import test
-class SequenceDatasetTest(test.TestCase):
+class SequenceDatasetTest(test_base.DatasetTestBase):
def testRepeatTensorDataset(self):
"""Test a dataset that repeats its input multiple times."""
diff --git a/tensorflow/python/data/kernel_tests/shard_dataset_op_test.py b/tensorflow/python/data/kernel_tests/shard_dataset_op_test.py
index 137f634..b9f3c79 100644
--- a/tensorflow/python/data/kernel_tests/shard_dataset_op_test.py
+++ b/tensorflow/python/data/kernel_tests/shard_dataset_op_test.py
@@ -17,12 +17,13 @@
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 errors
from tensorflow.python.platform import test
-class ShardDatasetOpTest(test.TestCase):
+class ShardDatasetOpTest(test_base.DatasetTestBase):
def testSimpleCase(self):
dataset = dataset_ops.Dataset.range(10).shard(5, 2)
diff --git a/tensorflow/python/data/kernel_tests/shuffle_dataset_op_test.py b/tensorflow/python/data/kernel_tests/shuffle_dataset_op_test.py
index f294840..347af18 100644
--- a/tensorflow/python/data/kernel_tests/shuffle_dataset_op_test.py
+++ b/tensorflow/python/data/kernel_tests/shuffle_dataset_op_test.py
@@ -21,6 +21,7 @@
import numpy as np
+from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import iterator_ops
from tensorflow.python.framework import constant_op
@@ -30,7 +31,7 @@
from tensorflow.python.platform import test
-class ShuffleDatasetTest(test.TestCase):
+class ShuffleDatasetTest(test_base.DatasetTestBase):
def testShuffleDataset(self):
components = (
diff --git a/tensorflow/python/data/kernel_tests/test_base.py b/tensorflow/python/data/kernel_tests/test_base.py
new file mode 100644
index 0000000..b730e10
--- /dev/null
+++ b/tensorflow/python/data/kernel_tests/test_base.py
@@ -0,0 +1,109 @@
+# 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.
+# ==============================================================================
+"""Test utilities for tf.data functionality."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import re
+
+from tensorflow.python.data.util import nest
+from tensorflow.python.eager import context
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.platform import test
+
+
+class DatasetTestBase(test.TestCase):
+ """Base class for dataset tests."""
+
+ def assertSparseValuesEqual(self, a, b):
+ """Asserts that two SparseTensors/SparseTensorValues are equal."""
+ self.assertAllEqual(a.indices, b.indices)
+ self.assertAllEqual(a.values, b.values)
+ self.assertAllEqual(a.dense_shape, b.dense_shape)
+
+ def getNext(self, dataset):
+ """Returns a callable that returns the next element of the dataset.
+
+ Example use:
+ ```python
+ # In both graph and eager modes
+ dataset = ...
+ nxt = self.getNext(dataset)
+ result = self.evaluate(nxt())
+ ```
+
+ Args:
+ dataset: A dataset whose next element is returned
+
+ Returns:
+ A callable that returns the next element of `dataset`
+ """
+ it = dataset.make_one_shot_iterator()
+ if context.executing_eagerly():
+ return it.get_next
+ else:
+ nxt = it.get_next()
+ return lambda: nxt
+
+ def assertDatasetsEqual(self, dataset1, dataset2):
+ """Checks that datasets are equal. Supports both graph and eager mode."""
+ self.assertEqual(dataset1.output_types, dataset2.output_types)
+ self.assertEqual(dataset1.output_classes, dataset2.output_classes)
+
+ next1 = self.getNext(dataset1)
+ next2 = self.getNext(dataset2)
+ while True:
+ try:
+ op1 = self.evaluate(next1())
+ except errors.OutOfRangeError:
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(next2())
+ break
+ op2 = self.evaluate(next2())
+
+ op1 = nest.flatten(op1)
+ op2 = nest.flatten(op2)
+ assert len(op1) == len(op2)
+ for i in range(len(op1)):
+ if isinstance(
+ op1[i],
+ (sparse_tensor.SparseTensor, sparse_tensor.SparseTensorValue)):
+ self.assertSparseValuesEqual(op1[i], op2[i])
+ else:
+ self.assertAllEqual(op1[i], op2[i])
+
+ def assertDatasetsRaiseSameError(self,
+ dataset1,
+ dataset2,
+ exception_class,
+ replacements=None):
+ """Checks that datasets raise the same error on the first get_next call."""
+ next1 = self.getNext(dataset1)
+ next2 = self.getNext(dataset2)
+ try:
+ self.evaluate(next1())
+ raise ValueError(
+ 'Expected dataset to raise an error of type %s, but it did not.' %
+ repr(exception_class))
+ except exception_class as e:
+ expected_message = e.message
+ for old, new, count in replacements:
+ expected_message = expected_message.replace(old, new, count)
+ # Check that the first segment of the error messages are the same.
+ with self.assertRaisesRegexp(exception_class,
+ re.escape(expected_message)):
+ self.evaluate(next2())
diff --git a/tensorflow/python/data/kernel_tests/window_dataset_op_test.py b/tensorflow/python/data/kernel_tests/window_dataset_op_test.py
index fd43484..9d06781 100644
--- a/tensorflow/python/data/kernel_tests/window_dataset_op_test.py
+++ b/tensorflow/python/data/kernel_tests/window_dataset_op_test.py
@@ -20,6 +20,7 @@
from absl.testing import parameterized
import numpy as np
+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
@@ -29,7 +30,7 @@
from tensorflow.python.platform import test
-class WindowDatasetTest(test.TestCase, parameterized.TestCase):
+class WindowDatasetTest(test_base.DatasetTestBase, parameterized.TestCase):
@parameterized.named_parameters(
("1", 20, 14, 7, 1),
@@ -150,11 +151,6 @@
stride_t: stride
})
- def assertSparseValuesEqual(self, a, b):
- self.assertAllEqual(a.indices, b.indices)
- self.assertAllEqual(a.values, b.values)
- self.assertAllEqual(a.dense_shape, b.dense_shape)
-
def testWindowSparse(self):
def _sparse(i):
diff --git a/tensorflow/python/data/kernel_tests/zip_dataset_op_test.py b/tensorflow/python/data/kernel_tests/zip_dataset_op_test.py
index 3106eff..9d76387 100644
--- a/tensorflow/python/data/kernel_tests/zip_dataset_op_test.py
+++ b/tensorflow/python/data/kernel_tests/zip_dataset_op_test.py
@@ -19,6 +19,7 @@
import numpy as np
+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
@@ -26,7 +27,7 @@
from tensorflow.python.platform import test
-class ZipDatasetTest(test.TestCase):
+class ZipDatasetTest(test_base.DatasetTestBase):
def testZipDataset(self):
component_placeholders = [
diff --git a/tensorflow/python/data/ops/dataset_ops.py b/tensorflow/python/data/ops/dataset_ops.py
index ac87a45..3b9d3a6 100644
--- a/tensorflow/python/data/ops/dataset_ops.py
+++ b/tensorflow/python/data/ops/dataset_ops.py
@@ -889,8 +889,8 @@
will be padded out to the maximum length of all elements in that
dimension.
- See also `tf.contrib.data.dense_to_sparse_batch`, which combines elements
- that may have different shapes into a `tf.SparseTensor`.
+ See also `tf.data.experimental.dense_to_sparse_batch`, which combines
+ elements that may have different shapes into a `tf.SparseTensor`.
Args:
batch_size: A `tf.int64` scalar `tf.Tensor`, representing the number of
@@ -1205,6 +1205,126 @@
shift = size
return WindowDataset(self, size, shift, stride, drop_remainder)
+ def reduce(self, initial_state, reduce_func):
+ """Reduces the input dataset to a single element.
+
+ The transformation calls `reduce_func` successively on every element of
+ the input dataset until the dataset is exhausted, aggregating information in
+ its internal state. The `initial_state` argument is used for the initial
+ state and the final state is returned as the result.
+
+ For example:
+ - `tf.data.Dataset.range(5).reduce(np.int64(0), lambda x, _: x + 1)`
+ produces `5`
+ - `tf.data.Dataset.range(5).reduce(np.int64(0), lambda x, y: x + y)`
+ produces `10`
+
+ Args:
+ initial_state: A nested structure of tensors, representing the initial
+ state of the transformation.
+ reduce_func: A function that maps `(old_state, input_element)` to
+ `new_state`. It must take two arguments and return a nested structure
+ of tensors. The structure of `new_state` must match the structure of
+ `initial_state`.
+
+ Returns:
+ A nested structure of `tf.Tensor` objects, corresponding to the final
+ state of the transformation.
+
+ """
+
+ with ops.name_scope("initial_state"):
+ # Convert any `SparseTensorValue`s to `SparseTensor`s and all other
+ # values to tensors.
+ initial_state = nest.pack_sequence_as(initial_state, [
+ sparse_tensor_lib.SparseTensor.from_value(t)
+ if sparse_tensor_lib.is_sparse(t) else ops.convert_to_tensor(
+ t, name="component_%d" % i)
+ for i, t in enumerate(nest.flatten(initial_state))
+ ])
+
+ # 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)])
+
+ # Iteratively rerun the reduce function until reaching a fixed point on
+ # `self._state_shapes`.
+ 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),
+ add_to_graph=False)
+
+ # Extract and validate class information from the returned values.
+ output_classes = wrapped_func.output_classes
+ for new_state_class, state_class in zip(
+ nest.flatten(output_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 "
+ "state. Expected %s; got %s." % (state_classes,
+ wrapped_func.output_classes))
+
+ # Extract and validate type information from the returned values.
+ output_types = wrapped_func.output_types
+ for new_state_type, state_type in zip(
+ nest.flatten(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." % (state_types,
+ wrapped_func.output_types))
+
+ # Extract shape information from the returned values.
+ output_shapes = wrapped_func.output_shapes
+ flat_state_shapes = nest.flatten(state_shapes)
+ flat_new_state_shapes = nest.flatten(output_shapes)
+ weakened_state_shapes = [
+ original.most_specific_compatible_shape(new)
+ for original, new in zip(flat_state_shapes, flat_new_state_shapes)
+ ]
+
+ need_to_rerun = False
+ for original_shape, weakened_shape in zip(flat_state_shapes,
+ weakened_state_shapes):
+ if original_shape.ndims is not None and (
+ weakened_shape.ndims is None or
+ original_shape.as_list() != weakened_shape.as_list()):
+ need_to_rerun = True
+ break
+
+ if need_to_rerun:
+ state_shapes = nest.pack_sequence_as(state_shapes,
+ weakened_state_shapes)
+
+ 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)
+
class DatasetSource(Dataset):
"""Abstract class representing a dataset with no inputs."""
diff --git a/tensorflow/python/data/ops/optional_ops.py b/tensorflow/python/data/ops/optional_ops.py
index 3bbebd7..aca989e 100644
--- a/tensorflow/python/data/ops/optional_ops.py
+++ b/tensorflow/python/data/ops/optional_ops.py
@@ -31,7 +31,7 @@
An `Optional` can represent the result of an operation that may fail as a
value, rather than raising an exception and halting execution. For example,
- `tf.contrib.data.get_next_as_optional` returns an `Optional` that either
+ `tf.data.experimental.get_next_as_optional` returns an `Optional` that either
contains the next value from a `tf.data.Iterator` if one exists, or a "none"
value that indicates the end of the sequence has been reached.
"""
@@ -111,7 +111,7 @@
class _OptionalImpl(Optional):
- """Concrete implementation of `tf.contrib.data.Optional`.
+ """Concrete implementation of `tf.data.experimental.Optional`.
NOTE(mrry): This implementation is kept private, to avoid defining
`Optional.__init__()` in the public API.
diff --git a/tensorflow/python/data/ops/readers.py b/tensorflow/python/data/ops/readers.py
index b0f2663..d08da67 100644
--- a/tensorflow/python/data/ops/readers.py
+++ b/tensorflow/python/data/ops/readers.py
@@ -129,7 +129,7 @@
def __init__(self, input_dataset, map_func, cycle_length, block_length,
sloppy, buffer_output_elements, prefetch_input_elements):
- """See `tf.contrib.data.parallel_interleave()` for details."""
+ """See `tf.data.experimental.parallel_interleave()` for details."""
super(ParallelInterleaveDataset, self).__init__(input_dataset, map_func,
cycle_length, block_length)
self._sloppy = ops.convert_to_tensor(
@@ -158,7 +158,7 @@
# pylint: enable=protected-access
def _transformation_name(self):
- return "tf.contrib.data.parallel_interleave()"
+ return "tf.data.experimental.parallel_interleave()"
@tf_export("data.TFRecordDataset")
diff --git a/tensorflow/python/debug/BUILD b/tensorflow/python/debug/BUILD
index 849d165..e84482d 100644
--- a/tensorflow/python/debug/BUILD
+++ b/tensorflow/python/debug/BUILD
@@ -18,6 +18,7 @@
load("//tensorflow:tensorflow.bzl", "cuda_py_test")
load("//tensorflow:tensorflow.bzl", "py_test")
+load("//tensorflow:tensorflow.bzl", "py_binary")
load("//tensorflow:tensorflow.bzl", "if_not_windows")
py_library(
diff --git a/tensorflow/python/debug/cli/analyzer_cli_test.py b/tensorflow/python/debug/cli/analyzer_cli_test.py
index 4630bda..f197a9e 100644
--- a/tensorflow/python/debug/cli/analyzer_cli_test.py
+++ b/tensorflow/python/debug/cli/analyzer_cli_test.py
@@ -599,11 +599,11 @@
v_name = "simple_mul_add/v"
u_init = constant_op.constant(u_init_val, shape=[2, 2], name="u_init")
- u = variables.Variable(u_init, name=u_name)
+ u = variables.VariableV1(u_init, name=u_name)
cls._u_line_number = line_number_above()
v_init = constant_op.constant(v_init_val, shape=[2, 1], name="v_init")
- v = variables.Variable(v_init, name=v_name)
+ v = variables.VariableV1(v_init, name=v_name)
cls._v_line_number = line_number_above()
w = math_ops.matmul(u, v, name="simple_mul_add/matmul")
@@ -612,7 +612,7 @@
x = math_ops.add(w, w, name="simple_mul_add/add")
cls._x_line_number = line_number_above()
- a = variables.Variable([1, 3, 3, 7], name="a")
+ a = variables.VariableV1([1, 3, 3, 7], name="a")
u.initializer.run()
v.initializer.run()
@@ -1371,7 +1371,7 @@
# Verify the annotation of the line that creates u.
index = self._findSourceLine(out, self._u_line_number)
self.assertEqual(
- ["L%d u = variables.Variable(u_init, name=u_name)" %
+ ["L%d u = variables.VariableV1(u_init, name=u_name)" %
self._u_line_number,
" simple_mul_add/u",
" simple_mul_add/u/Assign",
@@ -1388,7 +1388,7 @@
# Verify the annotation of the line that creates v.
index = self._findSourceLine(out, self._v_line_number)
self.assertEqual(
- ["L%d v = variables.Variable(v_init, name=v_name)" %
+ ["L%d v = variables.VariableV1(v_init, name=v_name)" %
self._v_line_number,
" simple_mul_add/v"],
out.lines[index : index + 2])
@@ -1425,7 +1425,7 @@
# Verify the annotation of the line that creates u.
index = self._findSourceLine(out, self._u_line_number)
self.assertEqual(
- ["L%d u = variables.Variable(u_init, name=u_name)" %
+ ["L%d u = variables.VariableV1(u_init, name=u_name)" %
self._u_line_number,
" simple_mul_add/u/read:0",
" simple_mul_add/u:0"],
@@ -1447,7 +1447,7 @@
index = self._findSourceLine(out, self._u_line_number)
self.assertEqual(
- ["L%d u = variables.Variable(u_init, name=u_name)" %
+ ["L%d u = variables.VariableV1(u_init, name=u_name)" %
self._u_line_number,
" simple_mul_add/u",
" simple_mul_add/u/Assign",
@@ -1470,7 +1470,7 @@
index = self._findSourceLine(out, self._u_line_number)
self.assertEqual(
- ["L%d u = variables.Variable(u_init, name=u_name)" %
+ ["L%d u = variables.VariableV1(u_init, name=u_name)" %
self._u_line_number,
" simple_mul_add/u",
" (... Omitted 2 of 3 op(s) ...) +5"],
@@ -1580,7 +1580,7 @@
"""List an input tree containing tensors from non-:0 output slot."""
with session.Session(config=no_rewrite_session_config()) as sess:
- x = variables.Variable([1, 3, 3, 7], name="x")
+ x = variables.VariableV1([1, 3, 3, 7], name="x")
_, idx = array_ops.unique(x, name="x_unique")
idx_times_two = math_ops.multiply(idx, 2, name="idx_times_two")
sess.run(x.initializer)
@@ -1684,7 +1684,7 @@
with session.Session(config=no_rewrite_session_config()) as sess:
x_init_val = np.array([5.0, 3.0])
x_init = constant_op.constant(x_init_val, shape=[2])
- x = variables.Variable(x_init, name="control_deps/x")
+ x = variables.VariableV1(x_init, name="control_deps/x")
y = math_ops.add(x, x, name="control_deps/y")
y = control_flow_ops.with_dependencies(
diff --git a/tensorflow/python/debug/cli/stepper_cli_test.py b/tensorflow/python/debug/cli/stepper_cli_test.py
index ee8cabc..7b8a42c 100644
--- a/tensorflow/python/debug/cli/stepper_cli_test.py
+++ b/tensorflow/python/debug/cli/stepper_cli_test.py
@@ -132,8 +132,8 @@
class NodeStepperSimpleGraphTest(test_util.TensorFlowTestCase):
def setUp(self):
- self.a = variables.Variable(10.0, name="a")
- self.b = variables.Variable(20.0, name="b")
+ self.a = variables.VariableV1(10.0, name="a")
+ self.b = variables.VariableV1(20.0, name="b")
self.c = math_ops.add(self.a, self.b, name="c") # Should be 30.0.
self.d = math_ops.subtract(self.a, self.c, name="d") # Should be -20.0.
diff --git a/tensorflow/python/debug/examples/debug_tflearn_iris.py b/tensorflow/python/debug/examples/debug_tflearn_iris.py
index 019f13c..f9bb314 100644
--- a/tensorflow/python/debug/examples/debug_tflearn_iris.py
+++ b/tensorflow/python/debug/examples/debug_tflearn_iris.py
@@ -94,13 +94,15 @@
"sepal_length", "sepal_width", "petal_length", "petal_width", "label"]
batch_size = 32
def training_input_fn():
- return tf.contrib.data.make_csv_dataset(
- [training_data_path], batch_size,
- column_names=column_names, label_name="label")
+ return tf.data.experimental.make_csv_dataset([training_data_path],
+ batch_size,
+ column_names=column_names,
+ label_name="label")
def test_input_fn():
- return tf.contrib.data.make_csv_dataset(
- [test_data_path], batch_size,
- column_names=column_names, label_name="label")
+ return tf.data.experimental.make_csv_dataset([test_data_path],
+ batch_size,
+ column_names=column_names,
+ label_name="label")
feature_columns = [tf.feature_column.numeric_column(feature)
for feature in column_names[:-1]]
diff --git a/tensorflow/python/debug/lib/debug_utils_test.py b/tensorflow/python/debug/lib/debug_utils_test.py
index 5b1875e..23ab984 100644
--- a/tensorflow/python/debug/lib/debug_utils_test.py
+++ b/tensorflow/python/debug/lib/debug_utils_test.py
@@ -46,8 +46,8 @@
cls._b_init = constant_op.constant(
cls._b_init_val, shape=[2, 1], name="b_init")
- cls._a = variables.Variable(cls._a_init, name="a1")
- cls._b = variables.Variable(cls._b_init, name="b")
+ cls._a = variables.VariableV1(cls._a_init, name="a1")
+ cls._b = variables.VariableV1(cls._b_init, name="b")
cls._c = constant_op.constant(cls._c_val, shape=[2, 1], name="c")
# Matrix product of a and b.
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 46a7be5..74498c8 100644
--- a/tensorflow/python/debug/lib/dist_session_debug_grpc_test.py
+++ b/tensorflow/python/debug/lib/dist_session_debug_grpc_test.py
@@ -118,8 +118,8 @@
"""
with ops.Graph().as_default() as graph:
with ops.device("/job:worker/task:0/cpu:0"):
- self.a = variables.Variable(10.0, name="a")
- self.b = variables.Variable(100.0, name="b")
+ self.a = variables.VariableV1(10.0, name="a")
+ self.b = variables.VariableV1(100.0, name="b")
self.inc_a = state_ops.assign_add(self.a, 2.0, name="inc_a")
self.dec_b = state_ops.assign_add(self.b, -5.0, name="dec_b")
self.p = math_ops.multiply(self.inc_a, self.dec_b, name="p")
diff --git a/tensorflow/python/debug/lib/grpc_large_data_test.py b/tensorflow/python/debug/lib/grpc_large_data_test.py
index 5bc477a..ccc21bc 100644
--- a/tensorflow/python/debug/lib/grpc_large_data_test.py
+++ b/tensorflow/python/debug/lib/grpc_large_data_test.py
@@ -61,7 +61,7 @@
with self.test_session(
use_gpu=True,
config=session_debug_testlib.no_rewrite_session_config()) as sess:
- u = variables.Variable(42.0, name="original_u")
+ u = variables.VariableV1(42.0, name="original_u")
for _ in xrange(50 * 1000):
u = array_ops.identity(u)
sess.run(variables.global_variables_initializer())
@@ -94,7 +94,7 @@
u_init = constant_op.constant(
u_init_val_array, dtype=dtypes.float32, name="u_init")
- u = variables.Variable(u_init, name="u")
+ u = variables.VariableV1(u_init, name="u")
def watch_fn(fetches, feeds):
del fetches, feeds # Unused by this watch_fn.
@@ -117,7 +117,7 @@
b"", b"spam", b"A" * 2500 * 1024, b"B" * 2500 * 1024, b"egg", b""]
u_init = constant_op.constant(
u_init_val, dtype=dtypes.string, name="u_init")
- u = variables.Variable(u_init, name="u")
+ u = variables.VariableV1(u_init, name="u")
def watch_fn(fetches, feeds):
del fetches, feeds
@@ -146,7 +146,7 @@
u_init = constant_op.constant(
u_init_val_array, dtype=dtypes.string, name="u_init")
- u = variables.Variable(u_init, name="u")
+ u = variables.VariableV1(u_init, name="u")
def watch_fn(fetches, feeds):
del fetches, feeds
@@ -167,7 +167,7 @@
config=session_debug_testlib.no_rewrite_session_config()) as sess:
u_init = constant_op.constant(
[], dtype=dtypes.float32, shape=[0], name="u_init")
- u = variables.Variable(u_init, name="u")
+ u = variables.VariableV1(u_init, name="u")
def watch_fn(fetches, feeds):
del fetches, feeds
@@ -189,7 +189,7 @@
config=session_debug_testlib.no_rewrite_session_config()) as sess:
u_init = constant_op.constant(
[], dtype=dtypes.string, shape=[0], name="u_init")
- u = variables.Variable(u_init, name="u")
+ u = variables.VariableV1(u_init, name="u")
def watch_fn(fetches, feeds):
del fetches, feeds
diff --git a/tensorflow/python/debug/lib/session_debug_file_test.py b/tensorflow/python/debug/lib/session_debug_file_test.py
index ba0f15b..1874160 100644
--- a/tensorflow/python/debug/lib/session_debug_file_test.py
+++ b/tensorflow/python/debug/lib/session_debug_file_test.py
@@ -58,9 +58,9 @@
v_name = "diff_Watch/v"
u_init = constant_op.constant(u_init_val, shape=[2, 2])
- u = variables.Variable(u_init, name=u_name)
+ u = variables.VariableV1(u_init, name=u_name)
v_init = constant_op.constant(v_init_val, shape=[2, 1])
- v = variables.Variable(v_init, name=v_name)
+ v = variables.VariableV1(v_init, name=v_name)
w = math_ops.matmul(u, v, name="diff_Watch/matmul")
diff --git a/tensorflow/python/debug/lib/session_debug_grpc_test.py b/tensorflow/python/debug/lib/session_debug_grpc_test.py
index 91f21cb..bfc9a3a 100644
--- a/tensorflow/python/debug/lib/session_debug_grpc_test.py
+++ b/tensorflow/python/debug/lib/session_debug_grpc_test.py
@@ -148,8 +148,8 @@
sess, "localhost:%d" % self._server_port, watch_fn="foo")
def testGrpcDebugWrapperSessionWithoutWatchFnWorks(self):
- u = variables.Variable(2.1, name="u")
- v = variables.Variable(20.0, name="v")
+ u = variables.VariableV1(2.1, name="u")
+ v = variables.VariableV1(20.0, name="v")
w = math_ops.multiply(u, v, name="w")
sess = session.Session(
@@ -175,8 +175,8 @@
del feeds, fetch_keys
return ["DebugIdentity", "DebugNumericSummary"], r".*/read", None
- u = variables.Variable(2.1, name="u")
- v = variables.Variable(20.0, name="v")
+ u = variables.VariableV1(2.1, name="u")
+ v = variables.VariableV1(20.0, name="v")
w = math_ops.multiply(u, v, name="w")
sess = session.Session(
@@ -209,8 +209,8 @@
op_type_regex_whitelist=None,
tolerate_debug_op_creation_failures=True)
- u = variables.Variable(2.1, name="u")
- v = variables.Variable(20.0, name="v")
+ u = variables.VariableV1(2.1, name="u")
+ v = variables.VariableV1(20.0, name="v")
w = math_ops.multiply(u, v, name="w")
sess = session.Session(
@@ -241,8 +241,8 @@
14, len(dump.get_tensors("v/read", 0, "DebugNumericSummary")[0]))
def testTensorBoardDebugHookWorks(self):
- u = variables.Variable(2.1, name="u")
- v = variables.Variable(20.0, name="v")
+ u = variables.VariableV1(2.1, name="u")
+ v = variables.VariableV1(20.0, name="v")
w = math_ops.multiply(u, v, name="w")
sess = session.Session(
@@ -286,8 +286,8 @@
self._server.query_source_file_line(__file__, 1)
def testTensorBoardDebugHookDisablingTracebackSourceCodeSendingWorks(self):
- u = variables.Variable(2.1, name="u")
- v = variables.Variable(20.0, name="v")
+ u = variables.VariableV1(2.1, name="u")
+ v = variables.VariableV1(20.0, name="v")
w = math_ops.multiply(u, v, name="w")
sess = session.Session(
@@ -381,8 +381,8 @@
def testToggleEnableTwoDebugWatchesNoCrosstalkBetweenDebugNodes(self):
with session.Session(
config=session_debug_testlib.no_rewrite_session_config()) as sess:
- v_1 = variables.Variable(50.0, name="v_1")
- v_2 = variables.Variable(-50.0, name="v_1")
+ v_1 = variables.VariableV1(50.0, name="v_1")
+ v_2 = variables.VariableV1(-50.0, name="v_1")
delta_1 = constant_op.constant(5.0, name="delta_1")
delta_2 = constant_op.constant(-5.0, name="delta_2")
inc_v_1 = state_ops.assign_add(v_1, delta_1, name="inc_v_1")
@@ -451,8 +451,8 @@
with session.Session(
config=session_debug_testlib.no_rewrite_session_config()) as sess:
- v_1 = variables.Variable(50.0, name="v_1")
- v_2 = variables.Variable(-50.0, name="v_1")
+ v_1 = variables.VariableV1(50.0, name="v_1")
+ v_2 = variables.VariableV1(-50.0, name="v_1")
# These two nodes have names that match those in the
# toggle_watch_on_core_metadata argument used when calling
# start_server_on_separate_thread().
@@ -491,7 +491,7 @@
def testToggleEnableTwoDebugWatchesNoCrosstalkBetweenServers(self):
with session.Session(
config=session_debug_testlib.no_rewrite_session_config()) as sess:
- v = variables.Variable(50.0, name="v")
+ v = variables.VariableV1(50.0, name="v")
delta = constant_op.constant(5.0, name="delta")
inc_v = state_ops.assign_add(v, delta, name="inc_v")
@@ -534,8 +534,8 @@
def testToggleBreakpointsWorks(self):
with session.Session(
config=session_debug_testlib.no_rewrite_session_config()) as sess:
- v_1 = variables.Variable(50.0, name="v_1")
- v_2 = variables.Variable(-50.0, name="v_2")
+ v_1 = variables.VariableV1(50.0, name="v_1")
+ v_2 = variables.VariableV1(-50.0, name="v_2")
delta_1 = constant_op.constant(5.0, name="delta_1")
delta_2 = constant_op.constant(-5.0, name="delta_2")
inc_v_1 = state_ops.assign_add(v_1, delta_1, name="inc_v_1")
@@ -592,8 +592,8 @@
def testTensorBoardDebuggerWrapperToggleBreakpointsWorks(self):
with session.Session(
config=session_debug_testlib.no_rewrite_session_config()) as sess:
- v_1 = variables.Variable(50.0, name="v_1")
- v_2 = variables.Variable(-50.0, name="v_2")
+ v_1 = variables.VariableV1(50.0, name="v_1")
+ v_2 = variables.VariableV1(-50.0, name="v_2")
delta_1 = constant_op.constant(5.0, name="delta_1")
delta_2 = constant_op.constant(-5.0, name="delta_2")
inc_v_1 = state_ops.assign_add(v_1, delta_1, name="inc_v_1")
@@ -665,8 +665,8 @@
def testTensorBoardDebuggerWrapperDisablingTracebackSourceSendingWorks(self):
with session.Session(
config=session_debug_testlib.no_rewrite_session_config()) as sess:
- v_1 = variables.Variable(50.0, name="v_1")
- v_2 = variables.Variable(-50.0, name="v_2")
+ v_1 = variables.VariableV1(50.0, name="v_1")
+ v_2 = variables.VariableV1(-50.0, name="v_2")
delta_1 = constant_op.constant(5.0, name="delta_1")
delta_2 = constant_op.constant(-5.0, name="delta_2")
inc_v_1 = state_ops.assign_add(v_1, delta_1, name="inc_v_1")
@@ -699,7 +699,7 @@
def testGetGrpcDebugWatchesReturnsCorrectAnswer(self):
with session.Session() as sess:
- v = variables.Variable(50.0, name="v")
+ v = variables.VariableV1(50.0, name="v")
delta = constant_op.constant(5.0, name="delta")
inc_v = state_ops.assign_add(v, delta, name="inc_v")
@@ -743,7 +743,7 @@
with self.cached_session() as sess:
a_init = constant_op.constant(42.0, name="a_init")
- a = variables.Variable(a_init, name="a")
+ a = variables.VariableV1(a_init, name="a")
def watch_fn(fetches, feeds):
del fetches, feeds
diff --git a/tensorflow/python/debug/lib/session_debug_testlib.py b/tensorflow/python/debug/lib/session_debug_testlib.py
index 070d9c4..25ef91b 100644
--- a/tensorflow/python/debug/lib/session_debug_testlib.py
+++ b/tensorflow/python/debug/lib/session_debug_testlib.py
@@ -70,7 +70,7 @@
def __init__(self, input_output_size, state_size):
self._input_output_size = input_output_size
self._state_size = state_size
- self._w = variables.Variable(1.0, dtype=dtypes.float32, name="w")
+ self._w = variables.VariableV1(1.0, dtype=dtypes.float32, name="w")
@property
def output_size(self):
@@ -182,9 +182,9 @@
w_name = "w"
u_init = constant_op.constant(u_init_val, shape=[2, 2])
- u = variables.Variable(u_init, name=u_name)
+ u = variables.VariableV1(u_init, name=u_name)
v_init = constant_op.constant(v_init_val, shape=[2, 1])
- v = variables.Variable(v_init, name=v_name)
+ v = variables.VariableV1(v_init, name=v_name)
w = math_ops.matmul(u, v, name=w_name)
@@ -221,8 +221,8 @@
def testCopyNodesHaveCorrectDebugOpsAndURLsAttributeValues(self):
with session.Session() as sess:
- u = variables.Variable(2.1, name="u")
- v = variables.Variable(20.0, name="v")
+ u = variables.VariableV1(2.1, name="u")
+ v = variables.VariableV1(20.0, name="v")
w = math_ops.multiply(u, v, name="w")
sess.run(variables.global_variables_initializer())
@@ -324,8 +324,8 @@
str1_name = "str1"
str2_name = "str2"
- str1 = variables.Variable(str1_init, name=str1_name)
- str2 = variables.Variable(str2_init, name=str2_name)
+ str1 = variables.VariableV1(str1_init, name=str1_name)
+ str2 = variables.VariableV1(str2_init, name=str2_name)
# Concatenate str1 and str2
str_concat = math_ops.add(str1, str2, name="str_concat")
@@ -387,9 +387,9 @@
s_name = "%s/s" % op_namespace
u_init = constant_op.constant(u_init_val, shape=[2, 2])
- u = variables.Variable(u_init, name=u_name)
+ u = variables.VariableV1(u_init, name=u_name)
s_init = constant_op.constant(s_init_val)
- s = variables.Variable(s_init, name=s_name)
+ s = variables.VariableV1(s_init, name=s_name)
run_options = config_pb2.RunOptions(output_partition_graphs=True)
debug_urls = self._debug_urls()
@@ -439,7 +439,7 @@
u_init_val = np.array(11.0)
u_init = constant_op.constant(u_init_val)
- u = variables.Variable(u_init, name=u_name)
+ u = variables.VariableV1(u_init, name=u_name)
# "v" is the increment.
v_name = "testDumpToFileWhileLoop/v"
@@ -447,7 +447,7 @@
v_init_val = np.array(2.0)
v_init = constant_op.constant(v_init_val)
- v = variables.Variable(v_init, name=v_name)
+ v = variables.VariableV1(v_init, name=v_name)
u.initializer.run()
v.initializer.run()
@@ -605,8 +605,8 @@
def testDebugCondWatchingWholeGraphWorks(self):
with session.Session() as sess:
- x = variables.Variable(10.0, name="x")
- y = variables.Variable(20.0, name="y")
+ x = variables.VariableV1(10.0, name="x")
+ y = variables.VariableV1(20.0, name="y")
cond = control_flow_ops.cond(
x > y, lambda: math_ops.add(x, 1), lambda: math_ops.add(y, 1))
@@ -628,9 +628,9 @@
z_name = "testFindNodesWithBadTensorValues/z"
u_init = constant_op.constant([2.0, 4.0])
- u = variables.Variable(u_init, name=u_name)
+ u = variables.VariableV1(u_init, name=u_name)
v_init = constant_op.constant([2.0, 1.0])
- v = variables.Variable(v_init, name=v_name)
+ v = variables.VariableV1(v_init, name=v_name)
# Expected output: [0.0, 3.0]
w = math_ops.subtract(u, v, name=w_name)
@@ -679,9 +679,9 @@
z_name = "testFindInfOrNanWithOpNameExclusion/z"
u_init = constant_op.constant([2.0, 4.0])
- u = variables.Variable(u_init, name=u_name)
+ u = variables.VariableV1(u_init, name=u_name)
v_init = constant_op.constant([2.0, 1.0])
- v = variables.Variable(v_init, name=v_name)
+ v = variables.VariableV1(v_init, name=v_name)
# Expected output: [0.0, 3.0]
w = math_ops.subtract(u, v, name=w_name)
@@ -725,7 +725,7 @@
w_name = "testDumpGraphStructureLookup/w"
u_init = constant_op.constant([2.0, 4.0])
- u = variables.Variable(u_init, name=u_name)
+ u = variables.VariableV1(u_init, name=u_name)
v = math_ops.add(u, u, name=v_name)
w = math_ops.add(v, v, name=w_name)
@@ -859,9 +859,9 @@
def testGraphPathFindingOnControlEdgesWorks(self):
with session.Session(config=no_rewrite_session_config()) as sess:
- v1 = variables.Variable(1.0, name="v1")
- v2 = variables.Variable(2.0, name="v2")
- v3 = variables.Variable(3.0, name="v3")
+ v1 = variables.VariableV1(1.0, name="v1")
+ v2 = variables.VariableV1(2.0, name="v2")
+ v3 = variables.VariableV1(3.0, name="v3")
a = math_ops.add(v1, v2, name="a")
with ops.control_dependencies([a]):
c = math_ops.subtract(v3, v3, name="c")
@@ -875,8 +875,8 @@
def testGraphPathFindingReverseRefEdgeWorks(self):
with session.Session(config=no_rewrite_session_config()) as sess:
- v = variables.Variable(10.0, name="v")
- delta = variables.Variable(1.0, name="delta")
+ v = variables.VariableV1(10.0, name="v")
+ delta = variables.VariableV1(1.0, name="delta")
inc_v = state_ops.assign_add(v, delta, name="inc_v")
sess.run(variables.global_variables_initializer())
@@ -894,7 +894,7 @@
w_name = "testDumpCausalityCheck/w"
u_init = constant_op.constant([2.0, 4.0])
- u = variables.Variable(u_init, name=u_name)
+ u = variables.VariableV1(u_init, name=u_name)
v = math_ops.add(u, u, name=v_name)
w = math_ops.add(v, v, name=w_name)
@@ -980,7 +980,7 @@
w_name = "oneOfTwoSlots/w"
y_name = "oneOfTwoSlots/y"
- x = variables.Variable([1, 3, 3, 7], dtype=dtypes.int32, name=x_name)
+ x = variables.VariableV1([1, 3, 3, 7], dtype=dtypes.int32, name=x_name)
sess.run(x.initializer)
unique_x, indices, _ = array_ops.unique_with_counts(x, name=u_name)
@@ -1039,9 +1039,9 @@
with session.Session(config=no_rewrite_session_config()) as sess:
u_init = constant_op.constant(10.0)
- u = variables.Variable(u_init, name="gdo/u")
+ u = variables.VariableV1(u_init, name="gdo/u")
v_init = constant_op.constant(20.0)
- v = variables.Variable(v_init, name="gdo/v")
+ v = variables.VariableV1(v_init, name="gdo/v")
w = math_ops.multiply(u, v, name="gdo/w")
# gdo stands for GradientDescentOptimizer.
@@ -1085,7 +1085,7 @@
with session.Session() as sess:
x_init = constant_op.constant([2, 2, 3, 5, 5])
- x = variables.Variable(x_init, name="unconnected/x")
+ x = variables.VariableV1(x_init, name="unconnected/x")
# The UniqueOp (tf.unique) has two output slots. Use only slot 0 in the
# graph. Let the debugger watch the unused slot 1.
@@ -1225,14 +1225,14 @@
def testDebugNumericSummaryOnInitializedTensorGivesCorrectResult(self):
with session.Session(config=no_rewrite_session_config()) as sess:
- a = variables.Variable(
+ a = variables.VariableV1(
[
np.nan, np.nan, 0.0, 0.0, 0.0, -1.0, -3.0, 3.0, 7.0, -np.inf,
-np.inf, np.inf, np.inf, np.inf, np.inf, np.inf, np.nan, np.nan
],
dtype=np.float32,
name="numeric_summary/a")
- b = variables.Variable(
+ b = variables.VariableV1(
[0.0] * 18, dtype=np.float32, name="numeric_summary/b")
c = math_ops.add(a, b, name="numeric_summary/c")
@@ -1249,7 +1249,7 @@
def testDebugNumericSummaryOnUninitializedTensorGivesCorrectResult(self):
with session.Session() as sess:
- a = variables.Variable(
+ a = variables.VariableV1(
[42], dtype=np.float32, name="numeric_summary_uninit/a")
_, dump = self._debug_run_and_get_dump(
@@ -1275,9 +1275,9 @@
def testDebugNumericSummaryFailureIsToleratedWhenOrdered(self):
with session.Session() as sess:
- a = variables.Variable("1", name="a")
- b = variables.Variable("3", name="b")
- c = variables.Variable("2", name="c")
+ a = variables.VariableV1("1", name="a")
+ b = variables.VariableV1("3", name="b")
+ c = variables.VariableV1("2", name="c")
d = math_ops.add(a, b, name="d")
e = math_ops.add(d, c, name="e")
@@ -1313,9 +1313,9 @@
def testDebugNumericSummaryInvalidAttributesStringAreCaught(self):
with session.Session(config=no_rewrite_session_config()) as sess:
- a = variables.Variable(10.0, name="a")
- b = variables.Variable(0.0, name="b")
- c = variables.Variable(0.0, name="c")
+ a = variables.VariableV1(10.0, name="a")
+ b = variables.VariableV1(0.0, name="b")
+ c = variables.VariableV1(0.0, name="c")
x = math_ops.divide(a, b, name="x")
y = math_ops.multiply(x, c, name="y")
@@ -1361,9 +1361,9 @@
def testDebugNumericSummaryMuteOnHealthyMutesOnlyHealthyTensorDumps(self):
with session.Session(config=no_rewrite_session_config()) as sess:
- a = variables.Variable(10.0, name="a")
- b = variables.Variable(0.0, name="b")
- c = variables.Variable(0.0, name="c")
+ a = variables.VariableV1(10.0, name="a")
+ b = variables.VariableV1(0.0, name="b")
+ c = variables.VariableV1(0.0, name="c")
x = math_ops.divide(a, b, name="x")
y = math_ops.multiply(x, c, name="y")
@@ -1396,8 +1396,8 @@
def testDebugNumericSummaryMuteOnHealthyAndCustomBoundsWork(self):
with session.Session() as sess:
- a = variables.Variable([10.0, 10.0], name="a")
- b = variables.Variable([10.0, 2.0], name="b")
+ a = variables.VariableV1([10.0, 10.0], name="a")
+ b = variables.VariableV1([10.0, 2.0], name="b")
x = math_ops.add(a, b, name="x") # [20.0, 12.0]
y = math_ops.divide(x, b, name="y") # [2.0, 6.0]
@@ -1436,9 +1436,9 @@
def testLookUpNodePythonTracebackWorks(self):
with session.Session() as sess:
u_init = constant_op.constant(10.0)
- u = variables.Variable(u_init, name="traceback/u")
+ u = variables.VariableV1(u_init, name="traceback/u")
v_init = constant_op.constant(20.0)
- v = variables.Variable(v_init, name="traceback/v")
+ v = variables.VariableV1(v_init, name="traceback/v")
w = math_ops.multiply(u, v, name="traceback/w")
@@ -1487,7 +1487,7 @@
self.skipTest("No testing concurrent runs on a single GPU.")
with session.Session() as sess:
- v = variables.Variable(30.0, name="v")
+ v = variables.VariableV1(30.0, name="v")
constants = []
for i in xrange(self._num_concurrent_runs):
constants.append(constant_op.constant(1.0, name="c%d" % i))
diff --git a/tensorflow/python/debug/lib/stepper_test.py b/tensorflow/python/debug/lib/stepper_test.py
index 9a3d0ef..3839c67 100644
--- a/tensorflow/python/debug/lib/stepper_test.py
+++ b/tensorflow/python/debug/lib/stepper_test.py
@@ -36,8 +36,8 @@
class StepperTest(test_util.TensorFlowTestCase):
def setUp(self):
- self.a = variables.Variable(2.0, name="a")
- self.b = variables.Variable(3.0, name="b")
+ self.a = variables.VariableV1(2.0, name="a")
+ self.b = variables.VariableV1(3.0, name="b")
self.c = math_ops.multiply(self.a, self.b, name="c") # Should be 6.0.
self.d = math_ops.multiply(self.a, self.a, name="d") # Should be 4.0.
@@ -49,7 +49,7 @@
# The there nodes x, y and z form a graph with "cross-links" in. I.e., x
# and y are both direct inputs to z, but x is also a direct input to y.
- self.x = variables.Variable(2.0, name="x") # Should be 2.0
+ self.x = variables.VariableV1(2.0, name="x") # Should be 2.0
self.y = math_ops.negative(self.x, name="y") # Should be -2.0.
self.z = math_ops.multiply(self.x, self.y, name="z") # Should be -4.0.
@@ -580,7 +580,7 @@
class StepperAssignAddTest(test_util.TensorFlowTestCase):
def setUp(self):
- self.v = variables.Variable(10.0, name="v")
+ self.v = variables.VariableV1(10.0, name="v")
self.p = math_ops.add(self.v, self.v, name="p")
self.q = math_ops.multiply(self.p, self.p, name="q")
self.delta = constant_op.constant(2.0, name="delta")
@@ -711,9 +711,9 @@
Construct a backward graph using the GradientDescentOptimizer.
"""
- self.a = variables.Variable(1.0, name="a")
- self.b = variables.Variable(2.0, name="b")
- self.c = variables.Variable(4.0, name="c")
+ self.a = variables.VariableV1(1.0, name="a")
+ self.b = variables.VariableV1(2.0, name="b")
+ self.c = variables.VariableV1(4.0, name="c")
self.d = math_ops.multiply(self.a, self.b, name="d")
self.e = math_ops.multiply(self.b, self.c, name="e")
self.f = math_ops.multiply(self.d, self.e, name="f")
diff --git a/tensorflow/python/debug/wrappers/dumping_wrapper_test.py b/tensorflow/python/debug/wrappers/dumping_wrapper_test.py
index 254201c..11011a5 100644
--- a/tensorflow/python/debug/wrappers/dumping_wrapper_test.py
+++ b/tensorflow/python/debug/wrappers/dumping_wrapper_test.py
@@ -46,7 +46,7 @@
def setUp(self):
self.session_root = tempfile.mkdtemp()
- self.v = variables.Variable(10.0, dtype=dtypes.float32, name="v")
+ self.v = variables.VariableV1(10.0, dtype=dtypes.float32, name="v")
self.delta = constant_op.constant(1.0, dtype=dtypes.float32, name="delta")
self.eta = constant_op.constant(-1.4, dtype=dtypes.float32, name="eta")
self.inc_v = state_ops.assign_add(self.v, self.delta, name="inc_v")
diff --git a/tensorflow/python/debug/wrappers/local_cli_wrapper_test.py b/tensorflow/python/debug/wrappers/local_cli_wrapper_test.py
index 05c9eaa..149a749 100644
--- a/tensorflow/python/debug/wrappers/local_cli_wrapper_test.py
+++ b/tensorflow/python/debug/wrappers/local_cli_wrapper_test.py
@@ -132,8 +132,8 @@
def setUp(self):
self._tmp_dir = tempfile.mktemp()
- self.v = variables.Variable(10.0, name="v")
- self.w = variables.Variable(21.0, name="w")
+ self.v = variables.VariableV1(10.0, name="v")
+ self.w = variables.VariableV1(21.0, name="w")
self.delta = constant_op.constant(1.0, name="delta")
self.inc_v = state_ops.assign_add(self.v, self.delta, name="inc_v")
@@ -358,7 +358,7 @@
def testDebuggingMakeCallableTensorRunnerWorks(self):
wrapped_sess = LocalCLIDebuggerWrapperSessionForTest(
[["run"], ["run"]], self.sess, dump_root=self._tmp_dir)
- v = variables.Variable(42)
+ v = variables.VariableV1(42)
tensor_runner = wrapped_sess.make_callable(v)
self.sess.run(v.initializer)
@@ -382,7 +382,7 @@
def testDebuggingMakeCallableOperationRunnerWorks(self):
wrapped_sess = LocalCLIDebuggerWrapperSessionForTest(
[["run"], ["run"]], self.sess, dump_root=self._tmp_dir)
- v = variables.Variable(10.0)
+ v = variables.VariableV1(10.0)
inc_v = state_ops.assign_add(v, 1.0)
op_runner = wrapped_sess.make_callable(inc_v.op)
self.sess.run(v.initializer)
@@ -403,7 +403,7 @@
self.assertEqual(1, len(wrapped_sess.observers["debug_dumps"]))
def testDebuggingMakeCallableFromOptionsWithZeroFeedWorks(self):
- variable_1 = variables.Variable(
+ variable_1 = variables.VariableV1(
10.5, dtype=dtypes.float32, name="variable_1")
a = math_ops.add(variable_1, variable_1, "callable_a")
math_ops.add(a, a, "callable_b")
@@ -480,7 +480,7 @@
self.assertItemsEqual(["callable_a", "callable_b"], node_names)
def testDebugMakeCallableFromOptionsWithCustomOptionsAndMetadataWorks(self):
- variable_1 = variables.Variable(
+ variable_1 = variables.VariableV1(
10.5, dtype=dtypes.float32, name="variable_1")
a = math_ops.add(variable_1, variable_1, "callable_a")
math_ops.add(a, a, "callable_b")
@@ -528,7 +528,7 @@
def testRuntimeErrorBeforeGraphExecutionIsRaised(self):
# Use an impossible device name to cause an error before graph execution.
with ops.device("/device:GPU:1337"):
- w = variables.Variable([1.0] * 10, name="w")
+ w = variables.VariableV1([1.0] * 10, name="w")
wrapped_sess = LocalCLIDebuggerWrapperSessionForTest(
[["run"]], self.sess, dump_root=self._tmp_dir)
diff --git a/tensorflow/python/distribute/distribute_coordinator.py b/tensorflow/python/distribute/distribute_coordinator.py
index bd3562f..b9b77d4 100644
--- a/tensorflow/python/distribute/distribute_coordinator.py
+++ b/tensorflow/python/distribute/distribute_coordinator.py
@@ -126,7 +126,7 @@
replicated training.
task_id: an integer indicating id of the corresponding task. It can be
None if it is local training or in-graph replicated training.
- session_config: an optional @{tf.ConfigProto} object.
+ session_config: an optional `tf.ConfigProto` object.
rpc_layer: optional string specifying the RPC protocol for communication
with worker masters. If None or empty, hosts in the `cluster_spec` will
be used directly.
@@ -685,7 +685,7 @@
in a cluster. If not set or empty, fall back to local training.
task_type: the current task type, optional if this is a client.
task_id: the current task id, optional if this is a client.
- session_config: an optional @{tf.ConfigProto} object which will be passed
+ session_config: an optional `tf.ConfigProto` object which will be passed
to `strategy`'s `configure` method and used to create a session.
rpc_layer: optional string, the protocol for RPC, e.g. "grpc".
diff --git a/tensorflow/python/distribute/estimator_training.py b/tensorflow/python/distribute/estimator_training.py
index 8daa34c..02896891 100644
--- a/tensorflow/python/distribute/estimator_training.py
+++ b/tensorflow/python/distribute/estimator_training.py
@@ -62,7 +62,7 @@
# Sort task names in cluster by "chief"/"master", "evaluator", "worker"
# and "ps". More details can be found at the documentation of
- # @{tf.estimator.RunConfig.global_id_in_cluster}.
+ # `tf.estimator.RunConfig.global_id_in_cluster`.
task_type_ordered_list = []
if chief_task_type in cluster_spec.jobs:
task_type_ordered_list = [chief_task_type]
diff --git a/tensorflow/python/eager/BUILD b/tensorflow/python/eager/BUILD
index d3d997e..d0c1a93 100644
--- a/tensorflow/python/eager/BUILD
+++ b/tensorflow/python/eager/BUILD
@@ -37,6 +37,7 @@
"//tensorflow/python:safe_ptr",
"//third_party/py/numpy:headers",
"//third_party/python_runtime:headers",
+ "@com_google_absl//absl/strings",
"@com_google_absl//absl/types:variant",
],
)
diff --git a/tensorflow/python/eager/backprop.py b/tensorflow/python/eager/backprop.py
index 78f3198..deac291 100644
--- a/tensorflow/python/eager/backprop.py
+++ b/tensorflow/python/eager/backprop.py
@@ -619,7 +619,7 @@
def _handle_or_self(x):
"""If x is ResourceVariable, return its handle, else x."""
- if isinstance(x, resource_variable_ops.ResourceVariable):
+ if resource_variable_ops.is_resource_variable(x):
x = x.handle
return x
diff --git a/tensorflow/python/eager/function.py b/tensorflow/python/eager/function.py
index b28befe..f261d92 100644
--- a/tensorflow/python/eager/function.py
+++ b/tensorflow/python/eager/function.py
@@ -21,6 +21,7 @@
import collections
import functools
+import re
import sys
import threading
import weakref
@@ -61,9 +62,15 @@
# This is to avoid a circular dependency with gradients_impl
gradients_impl._function = sys.modules[__name__] # pylint: disable=protected-access
+FORWARD_FUNCTION_ATTRIBUTE_NAME = "forward_function_name"
+BACKWARD_FUNCTION_ATTRIBUTE_NAME = "backward_function_name"
# TODO(scottzhu): Update this to allow arbitrary attribute names in future.
-WHITELIST_FUNCTION_ATTRIBUTE_PREFIX = "experimental_"
+WHITELIST_FUNCTION_ATTRIBUTE_REGEX = [
+ "experimental_.*",
+ FORWARD_FUNCTION_ATTRIBUTE_NAME,
+ BACKWARD_FUNCTION_ATTRIBUTE_NAME
+]
def _create_substitute_placeholder(value, name=None, dtype=None):
@@ -140,10 +147,11 @@
"""
attrs = {}
for key, value in attributes.items():
- if not key.startswith(WHITELIST_FUNCTION_ATTRIBUTE_PREFIX):
+ if not any([re.match(reg, key)
+ for reg in WHITELIST_FUNCTION_ATTRIBUTE_REGEX]):
raise ValueError("Attribute name is not whitelisted. "
"Whitelisted: prefix %s, got: %s" %
- (WHITELIST_FUNCTION_ATTRIBUTE_PREFIX, key))
+ (WHITELIST_FUNCTION_ATTRIBUTE_REGEX, key))
if isinstance(value, attr_value_pb2.AttrValue):
attrs[key] = value
@@ -154,7 +162,7 @@
attrs[key] = attr_value_pb2.AttrValue(i=value)
elif isinstance(value, float):
attrs[key] = attr_value_pb2.AttrValue(f=value)
- elif isinstance(value, str):
+ elif isinstance(value, (str, bytes)):
attrs[key] = attr_value_pb2.AttrValue(s=compat.as_bytes(value))
else:
raise ValueError("Unsupported attribute type for %s with type %s" %
@@ -261,6 +269,15 @@
def variables(self, var_list):
self._weak_variables = [weakref.ref(v) for v in var_list]
+ def control_dependencies(self, control_inputs):
+ # Drop control dependencies to outside of the graph. TODO(b/117109273)
+ # unclear how to capture an op, not a tensor.
+ if not control_inputs:
+ return super(FuncGraph, self).control_dependencies(control_inputs)
+ return super(FuncGraph, self).control_dependencies(
+ [c for c in control_inputs
+ if getattr(c, "graph", None) is self])
+
def create_op(
self,
op_type,
@@ -705,6 +722,7 @@
def _construct_backprop_function(self):
"""Constructs the backprop function object for this function."""
backwards_graph = FuncGraph(_backward_name(self._func_graph.name))
+ forward_function_name = _forward_name(self._func_graph.name)
with backwards_graph.as_default():
gradients_wrt_outputs = [
graph_placeholder(x.dtype, x.shape) for x in self._func_graph.outputs
@@ -715,11 +733,11 @@
grad_ys=gradients_wrt_outputs,
src_graph=self._func_graph)
- self._forward_function = _EagerDefinedFunction(
- _forward_name(
- self._func_graph.name), self._func_graph, self._func_graph.inputs,
- self._func_graph.outputs + list(backwards_graph.captures.keys()),
- self._attrs)
+ backwards_graph_captures = list(backwards_graph.captures.keys())
+
+ backward_function_attr = _parse_func_attrs(
+ {FORWARD_FUNCTION_ATTRIBUTE_NAME: forward_function_name})
+ backward_function_attr.update(self._attrs)
# The ordering of `backwards_graph.inputs` is important: inputs of
# `self._backward_graph_function` correspond to outputs of
@@ -732,7 +750,17 @@
grad for grad in _flatten(gradients_wrt_inputs) if grad is not None)
backwards_graph.structured_outputs = gradients_wrt_inputs
self._backward_graph_function = Function(
- backwards_graph, attrs=self._attrs)
+ backwards_graph, attrs=backward_function_attr)
+
+ forward_function_attr = _parse_func_attrs({
+ BACKWARD_FUNCTION_ATTRIBUTE_NAME:
+ self._backward_graph_function._inference_function.name}) # pylint: disable=protected-access
+ forward_function_attr.update(self._attrs)
+
+ self._forward_function = _EagerDefinedFunction(
+ forward_function_name, self._func_graph, self._func_graph.inputs,
+ self._func_graph.outputs + backwards_graph_captures,
+ forward_function_attr)
def _backprop_call(self, args):
"""Calls the forward function and records the result on a tape.
@@ -986,52 +1014,8 @@
return func_graph
-_TensorType = collections.namedtuple("_TensorType", ["dtype", "shape"])
-
-
-def _encode_arg(arg):
- """A canonical representation for this argument, for use in a cache key."""
-
- # `defun` uses dtypes and shapes instead of `Tensors` as cache keys. Dtypes
- # are used because TensorFlow graphs are not parametric w.r.t. dtypes. Shapes
- # are used for both performance reasons, as much TensorFlow code specializes
- # on known shapes to produce slimmer graphs, and correctness, as some
- # high-level APIs require shapes to be fully-known.
- #
- # TODO(akshayka): Add support for sparse tensors.
- #
- # pylint: disable=protected-access
- if isinstance(arg, ops.Tensor):
- return _TensorType(arg.dtype, arg._shape_tuple())
- elif isinstance(arg, ops.IndexedSlices):
- if arg.dense_shape is not None:
- return tuple([
- _TensorType(arg.values.dtype, arg.values._shape_tuple()),
- _TensorType(arg.indices.dtype, arg.indices._shape_tuple()),
- _TensorType(arg.dense_shape.dtype, arg.dense_shape._shape_tuple()),
- ])
- else:
- return tuple([
- _TensorType(arg.values.dtype, arg.values._shape_tuple()),
- _TensorType(arg.indices.dtype, arg.indices._shape_tuple()),
- ])
- # pylint: enable=protected-access
- elif isinstance(arg, (list, tuple)):
- return tuple([_encode_arg(elem) for elem in arg])
- elif isinstance(arg, dict):
- return tuple(
- (_encode_arg(key), _encode_arg(arg[key])) for key in sorted(arg))
- else:
- try:
- # If possible, keep only a weak reference to Python objects. Weak
- # references hash to the same value as the original object.
- # TODO(allenl): Clean up dead functions and their cache keys if the cache
- # gets large. Right now creating objects with a defunned method, calling
- # the method, and losing a reference to the object in a loop will leak
- # memory here.
- return weakref.ref(arg)
- except TypeError:
- return arg
+pywrap_tensorflow.RegisterType("Tensor", ops.Tensor)
+pywrap_tensorflow.RegisterType("IndexedSlices", ops.IndexedSlices)
def _deterministic_dict_values(dictionary):
@@ -1101,6 +1085,8 @@
offset + index: default
for index, default in enumerate(fullargspec.defaults or [])
}
+ self._default_values = fullargspec.defaults
+ self._default_values_start_index = offset
if input_signature is None:
self._input_signature = None
else:
@@ -1161,7 +1147,7 @@
"""Computes the cache key given inputs and execution context."""
if self._input_signature is None:
inputs = (args, kwargs) if kwargs else args
- cache_key = tuple(_encode_arg(arg) for arg in inputs)
+ cache_key = pywrap_tensorflow.TFE_Py_EncodeArg(inputs)
else:
del args, kwargs
cache_key = self._flat_input_signature
@@ -1184,7 +1170,7 @@
colocation_stack = (() if executing_eagerly else
tuple(default_graph._colocation_stack.peek_objs())) # pylint: disable=protected-access
- return cache_key + (execution_context, device_functions, colocation_stack)
+ return (cache_key, execution_context, device_functions, colocation_stack)
def _canonicalize_function_inputs(self, *args, **kwargs):
"""Canonicalizes `args` and `kwargs`.
@@ -1212,26 +1198,32 @@
# Maps from index of arg to its corresponding value, according to `args`
# and `kwargs`; seeded with the default values for the named args that
# aren't in `args`.
- arg_indices_to_values = {
- index: default
- for index, default in six.iteritems(self._arg_indices_to_default_values)
- if index >= len(args)
- }
- consumed_args = []
- for arg, value in six.iteritems(kwargs):
- index = self._args_to_indices.get(arg, None)
- if index is not None:
- arg_indices_to_values[index] = value
- consumed_args.append(arg)
- elif self._input_signature is not None:
- raise ValueError("Cannot define a TensorFlow function from a Python "
- "function with keyword arguments when "
- "input_signature is provided.")
- for arg in consumed_args:
- # After this loop, `kwargs` will only contain true keyword arguments, as
- # opposed to named arguments called in a keyword-like fashion.
- kwargs.pop(arg)
- inputs = args + _deterministic_dict_values(arg_indices_to_values)
+ if not kwargs:
+ if self._default_values:
+ inputs = args + self._default_values[len(args) -
+ self._default_values_start_index:]
+ else:
+ inputs = args
+ else:
+ arg_indices_to_values = {
+ index: default for index, default in six.iteritems(
+ self._arg_indices_to_default_values) if index >= len(args)
+ }
+ consumed_args = []
+ for arg, value in six.iteritems(kwargs):
+ index = self._args_to_indices.get(arg, None)
+ if index is not None:
+ arg_indices_to_values[index] = value
+ consumed_args.append(arg)
+ elif self._input_signature is not None:
+ raise ValueError("Cannot define a TensorFlow function from a Python "
+ "function with keyword arguments when "
+ "input_signature is provided.")
+ for arg in consumed_args:
+ # After this loop, `kwargs` will only contain true keyword arguments, as
+ # opposed to named arguments called in a keyword-like fashion.
+ kwargs.pop(arg)
+ inputs = args + _deterministic_dict_values(arg_indices_to_values)
flat_inputs = nest.flatten(inputs)
# Check for NumPy arrays in arguments and convert them to Tensors.
@@ -1328,8 +1320,25 @@
"Got type: %s" % type(func))
concrete_func = func.get_concrete_function(*args, **kwargs)
graph = ops.get_default_graph()
- concrete_func._inference_function.add_to_graph(graph) # pylint: disable=protected-access
- # TODO(scottzhu): support concrete_func._backward_graph_function in future.
+
+ # There are two situations for the actual call of a defun:
+ # 1. If none of the input args are resource variables or watch by any tape,
+ # it will run the _inference_function of concrete_func for forward pass, and
+ # the gradient will be generated by standard mechanism.
+ # 2. Otherwise, defun will create two functions, one for forward pass, and the
+ # backward pass will be created via tape.
+ # When registering the function, we put both cases into graph.
+ # pylint: disable=protected-access
+ concrete_func._inference_function.add_to_graph(graph)
+
+ if concrete_func._backward_graph_function is None:
+ concrete_func._construct_backprop_function()
+ forward_function = concrete_func._forward_function
+ backward_function = concrete_func._backward_graph_function._inference_function
+ forward_function.add_to_graph(graph)
+ backward_function.add_to_graph(graph)
+ # pylint: enable=protected-access
+
return concrete_func
diff --git a/tensorflow/python/eager/function_test.py b/tensorflow/python/eager/function_test.py
index 59faf96..9ce367a 100644
--- a/tensorflow/python/eager/function_test.py
+++ b/tensorflow/python/eager/function_test.py
@@ -1237,6 +1237,24 @@
x = constant_op.constant([1.0, 2.0])
self.assertAllEqual([2., 4.], self.evaluate(defined(x)))
+ def testCacheObjectHashCollisions(self):
+
+ class Foo(object):
+
+ def __hash__(self):
+ return 42
+
+ def func(foo):
+ del foo
+ return
+
+ defined = function.defun(func)
+ defined(Foo())
+ self.assertEqual(len(defined._function_cache), 1)
+
+ defined(Foo())
+ self.assertEqual(len(defined._function_cache), 2)
+
def testPythonFunctionWithDefaultArgs(self):
def func(foo, bar=1, baz=2):
@@ -1250,20 +1268,20 @@
def cache_keys():
"""Sanitizes cache keys of non-input metadata."""
- return tuple(key[:3] for key in defined._function_cache)
+ return tuple(key[0] for key in defined._function_cache)
# `True` corresponds to the fact that we're executing eagerly
- self.assertIn((0, 1, 20), cache_keys())
+ self.assertIn(('tRRR', (0, 1, 20)), cache_keys())
defined(1) # bar=1, baz=2
- self.assertIn((1, 1, 2), cache_keys())
+ self.assertIn(('tRRR', (1, 1, 2)), cache_keys())
# This matches the previous call.
defined(foo=1)
self.assertEqual(len(defined._function_cache), 2)
defined(1, 2, 3)
- self.assertIn((1, 2, 3), cache_keys())
+ self.assertIn(('tRRR', (1, 2, 3)), cache_keys())
# This matches the previous call.
defined(1, bar=2, baz=3)
@@ -1669,12 +1687,38 @@
graph = ops.get_default_graph()
# pylint: disable=protected-access
- self.assertEqual(len(graph._functions), 2)
+ self.assertEqual(len(graph._functions), 6)
+ # two sets of functions, each of them are (inference, forward, backward)
functions = list(graph._functions.values())
- pre_register_matmul_func_name = functions[0].definition.signature.name
- self.assertRegexpMatches(pre_register_matmul_func_name, '.*matmul.*')
- pre_register_add_func_name = functions[1].definition.signature.name
- self.assertRegexpMatches(pre_register_add_func_name, '.*add.*')
+ captured_function_names = [
+ f.definition.signature.name for f in functions
+ ]
+ expected_func_name_regex = [
+ '.*inference.*matmul.*',
+ '.*forward.*matmul.*',
+ '.*inference.*backward.*matmul.*',
+ '.*inference.*add.*',
+ '.*forward.*add.*',
+ '.*inference.*backward.*add.*',
+ ]
+ for i in range(len(functions)):
+ self.assertRegexpMatches(captured_function_names[i],
+ expected_func_name_regex[i])
+
+ # Check the forward and backward function has the correct attributes.
+ self.assertEquals(
+ functions[1].definition.attr['backward_function_name'].s,
+ functions[2].name)
+ self.assertEquals(
+ functions[2].definition.attr['forward_function_name'].s,
+ functions[1].name)
+
+ self.assertEquals(
+ functions[4].definition.attr['backward_function_name'].s,
+ functions[5].name)
+ self.assertEquals(
+ functions[5].definition.attr['forward_function_name'].s,
+ functions[4].name)
sq = defun_matmul(t, t)
double = add(t, t)
@@ -1682,12 +1726,11 @@
self.assertAllEqual(double.eval().reshape(-1), [2, 4, 6, 8])
# Make sure the pre registered function is used, and no other function
# is added.
- self.assertEqual(len(graph._functions), 2)
+ self.assertEqual(len(graph._functions), 6)
functions = list(graph._functions.values())
- called_func_name = functions[0].definition.signature.name
- self.assertEqual(pre_register_matmul_func_name, called_func_name)
- called_func_name = functions[1].definition.signature.name
- self.assertEqual(pre_register_add_func_name, called_func_name)
+ for i in range(len(functions)):
+ self.assertEquals(captured_function_names[i],
+ functions[i].definition.signature.name)
def testRegisterFunctionWithInputSignature(self):
def matmul(x, y):
@@ -1705,7 +1748,7 @@
graph = ops.get_default_graph()
# pylint: disable=protected-access
- self.assertEqual(len(graph._functions), 1)
+ self.assertEqual(len(graph._functions), 3)
# Test input param shape mismatch
t2 = constant_op.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
@@ -1728,7 +1771,7 @@
graph = ops.get_default_graph()
# Only one function is registered since the input param are in same type
# pylint: disable=protected-access
- self.assertEqual(len(graph._functions), 1)
+ self.assertEqual(len(graph._functions), 3)
def testCallingFunctionWithDifferentVariables(self):
@@ -1767,7 +1810,8 @@
'be Tensors;.*'):
graph_function('Not a Tensor.')
- def testSwapImplementationWithGrapplerPlugin(self):
+ # TODO(scottzhu): Revive the test once the grappler plugin is updated.
+ def disabled_testSwapImplementationWithGrapplerPlugin(self):
rewrites = rewriter_config_pb2.RewriterConfig()
# function_optimizer has to be turn off, otherwise it will delete the
# registered function if it does not get called.
diff --git a/tensorflow/python/eager/pywrap_tfe.h b/tensorflow/python/eager/pywrap_tfe.h
index f1b4042..decd635 100755
--- a/tensorflow/python/eager/pywrap_tfe.h
+++ b/tensorflow/python/eager/pywrap_tfe.h
@@ -224,4 +224,8 @@
// The shape is represented as a Python tuple of integers.
PyObject* TFE_Py_TensorShapeOnDevice(PyObject* tensor);
+// Encodes the object as a tuple that is meant to be used as part of the key
+// for the defun function cache.
+PyObject* TFE_Py_EncodeArg(PyObject*);
+
#endif // TENSORFLOW_PYTHON_EAGER_PYWRAP_TFE_H_
diff --git a/tensorflow/python/eager/pywrap_tfe_src.cc b/tensorflow/python/eager/pywrap_tfe_src.cc
index 196e20e..ae1e12f 100644
--- a/tensorflow/python/eager/pywrap_tfe_src.cc
+++ b/tensorflow/python/eager/pywrap_tfe_src.cc
@@ -17,6 +17,7 @@
#include "tensorflow/python/eager/pywrap_tfe.h"
+#include "absl/strings/str_cat.h"
#include "absl/types/variant.h"
#include "tensorflow/c/c_api.h"
#include "tensorflow/c/c_api_internal.h"
@@ -567,11 +568,8 @@
return false;
}
}
- TFE_Op* func = TFE_NewOp(
- ctx, string(func_name.data(), func_name.size()).c_str(), status);
- if (TF_GetCode(status) != TF_OK) return false;
- TFE_OpSetAttrFunction(op, key, func);
- TFE_DeleteOp(func);
+ TF_SetStatus(status, TF_OK, "");
+ TFE_OpSetAttrFunctionName(op, key, func_name.data(), func_name.size());
} else {
TF_SetStatus(
status, TF_UNIMPLEMENTED,
@@ -1569,9 +1567,8 @@
}
for (TFE_Py_Tape* tape : SafeTapeSet()) {
- auto* function = backward_function_getter();
tape->tape->RecordOperation(op_type_str, output_info, input_ids,
- input_dtypes, function,
+ input_dtypes, backward_function_getter,
backward_function_killer);
}
}
@@ -2748,3 +2745,218 @@
return RecordGradient(op_name, inputs, attrs, results, name);
}
+
+namespace {
+
+tensorflow::int64 GetPyNoneHash() {
+ tensorflow::int64 py_none_hash = PyObject_Hash(Py_None);
+ return py_none_hash;
+}
+
+struct EncodeResult {
+ string str;
+ std::vector<PyObject*> objects;
+
+ PyObject* ToPyTuple() {
+ PyObject* result = PyTuple_New(2);
+
+ PyTuple_SET_ITEM(result, 0, GetPythonObjectFromString(str.c_str()));
+
+ if (objects.empty()) {
+ Py_INCREF(Py_None);
+ PyTuple_SET_ITEM(result, 1, Py_None);
+ } else {
+ PyObject* objects_tuple = PyTuple_New(objects.size());
+
+ for (int i = 0; i < objects.size(); i++) {
+ PyTuple_SET_ITEM(objects_tuple, i, objects[i]);
+ }
+
+ PyTuple_SET_ITEM(result, 1, objects_tuple);
+ }
+
+ return result;
+ }
+};
+
+tensorflow::Status TFE_Py_EncodeTensor(PyObject* arg, EncodeResult* result) {
+ if (EagerTensor_CheckExact(arg)) {
+ TFE_TensorHandle* t = EagerTensor_Handle(arg);
+ tensorflow::TensorShape tensor_shape;
+ TF_RETURN_IF_ERROR(t->handle->Shape(&tensor_shape));
+ absl::StrAppend(&result->str, t->handle->dtype);
+
+ for (tensorflow::int64 dim_size : tensor_shape.dim_sizes()) {
+ absl::StrAppend(&result->str, dim_size);
+ }
+
+ return tensorflow::Status::OK();
+ }
+
+ tensorflow::Safe_PyObjectPtr dtype_object(
+ PyObject_GetAttrString(arg, "dtype"));
+
+ if (dtype_object == nullptr) {
+ return tensorflow::errors::InvalidArgument(
+ "ops.Tensor object doesn't have dtype() attr.");
+ }
+
+ tensorflow::Safe_PyObjectPtr dtype_enum(
+ PyObject_GetAttrString(dtype_object.get(), "_type_enum"));
+
+ if (dtype_enum == nullptr) {
+ return tensorflow::errors::InvalidArgument(
+ "ops.Tensor's dtype object doesn't have _type_enum() attr.");
+ }
+
+ tensorflow::DataType dtype =
+ static_cast<tensorflow::DataType>(MakeInt(dtype_enum.get()));
+
+ absl::StrAppend(&result->str, dtype);
+ static char _shape_tuple[] = "_shape_tuple";
+ tensorflow::Safe_PyObjectPtr shape_tuple(
+ PyObject_CallMethod(arg, _shape_tuple, nullptr));
+
+ if (shape_tuple == nullptr) {
+ return tensorflow::errors::InvalidArgument(
+ "ops.Tensor object doesn't have _shape_tuple() method.");
+ }
+
+ if (shape_tuple.get() == Py_None) {
+ // Unknown shape, encode that directly.
+ absl::StrAppend(&result->str, GetPyNoneHash());
+ return tensorflow::Status::OK();
+ }
+
+ tensorflow::Safe_PyObjectPtr shape_seq(PySequence_Fast(
+ shape_tuple.get(), "shape_tuple didn't return a sequence"));
+
+ int len = PySequence_Fast_GET_SIZE(shape_seq.get());
+ for (int i = 0; i < len; ++i) {
+ PyObject* item = PySequence_Fast_GET_ITEM(shape_seq.get(), i);
+ if (item == Py_None) {
+ absl::StrAppend(&result->str, GetPyNoneHash());
+ } else {
+ absl::StrAppend(&result->str, MakeInt(item));
+ }
+ }
+
+ return tensorflow::Status::OK();
+}
+
+const char kTensor[] = "T";
+const char kIndexedSlices[] = "I";
+const char kList[] = "L";
+const char kTuple[] = "t";
+const char kDict[] = "D";
+const char kRaw[] = "R";
+
+tensorflow::Status TFE_Py_EncodeArgHelper(PyObject* arg, EncodeResult* result);
+
+// This function doesn't set the type of sequence before
+tensorflow::Status TFE_Py_EncodeSequence(PyObject* arg, const char* type,
+ EncodeResult* result) {
+ tensorflow::Safe_PyObjectPtr arg_seq(
+ PySequence_Fast(arg, "unable to create seq from list/tuple"));
+
+ absl::StrAppend(&result->str, type);
+ int len = PySequence_Fast_GET_SIZE(arg_seq.get());
+ for (int i = 0; i < len; ++i) {
+ PyObject* item = PySequence_Fast_GET_ITEM(arg_seq.get(), i);
+ if (item == Py_None) {
+ absl::StrAppend(&result->str, GetPyNoneHash());
+ } else {
+ TF_RETURN_IF_ERROR(TFE_Py_EncodeArgHelper(item, result));
+ }
+ }
+
+ return tensorflow::Status::OK();
+}
+
+tensorflow::Status TFE_Py_EncodeArgHelper(PyObject* arg, EncodeResult* result) {
+ if (tensorflow::swig::IsTensor(arg)) {
+ absl::StrAppend(&result->str, kTensor);
+ TF_RETURN_IF_ERROR(TFE_Py_EncodeTensor(arg, result));
+ } else if (tensorflow::swig::IsIndexedSlices(arg)) {
+ absl::StrAppend(&result->str, kIndexedSlices);
+ tensorflow::Safe_PyObjectPtr values(PyObject_GetAttrString(arg, "values"));
+ if (values == nullptr) {
+ PyErr_Clear();
+ return tensorflow::errors::InvalidArgument(
+ "IndexedSlices does not have a values attr");
+ }
+ TF_RETURN_IF_ERROR(TFE_Py_EncodeTensor(values.get(), result));
+
+ tensorflow::Safe_PyObjectPtr indices(
+ PyObject_GetAttrString(arg, "indices"));
+ if (indices == nullptr) {
+ PyErr_Clear();
+ return tensorflow::errors::InvalidArgument(
+ "IndexedSlices does not have a indices attr");
+ }
+ TF_RETURN_IF_ERROR(TFE_Py_EncodeTensor(indices.get(), result));
+
+ tensorflow::Safe_PyObjectPtr dense_shape(
+ PyObject_GetAttrString(arg, "dense_shape"));
+ if (dense_shape == nullptr) {
+ PyErr_Clear();
+ return tensorflow::errors::InvalidArgument(
+ "IndexedSlices does not have a dense_shape attr");
+ }
+ if (dense_shape.get() != Py_None) {
+ TF_RETURN_IF_ERROR(TFE_Py_EncodeTensor(dense_shape.get(), result));
+ }
+ } else if (PyList_Check(arg)) {
+ TF_RETURN_IF_ERROR(TFE_Py_EncodeSequence(arg, kList, result));
+ } else if (PyTuple_Check(arg)) {
+ TF_RETURN_IF_ERROR(TFE_Py_EncodeSequence(arg, kTuple, result));
+ } else if (PyDict_Check(arg)) {
+ tensorflow::Safe_PyObjectPtr keys(PyDict_Keys(arg));
+ if (PyList_Sort(keys.get()) == -1) {
+ return tensorflow::errors::Internal("Unable to sort keys");
+ }
+
+ absl::StrAppend(&result->str, kDict);
+ int len = PyList_Size(keys.get());
+
+ for (int i = 0; i < len; i++) {
+ PyObject* key = PyList_GetItem(keys.get(), i);
+ TF_RETURN_IF_ERROR(TFE_Py_EncodeArgHelper(key, result));
+ PyObject* value = PyDict_GetItem(arg, key);
+ TF_RETURN_IF_ERROR(TFE_Py_EncodeArgHelper(value, result));
+ }
+ } else {
+ PyObject* object = PyWeakref_NewRef(arg, nullptr);
+
+ if (object == nullptr) {
+ PyErr_Clear();
+
+ object = arg;
+ Py_INCREF(object);
+ }
+
+ absl::StrAppend(&result->str, kRaw);
+ result->objects.push_back(object);
+ }
+
+ return tensorflow::Status::OK();
+}
+
+} // namespace
+
+// `defun` uses dtypes and shapes instead of `Tensors` as cache keys. Dtypes
+// are used because TensorFlow graphs are not parametric w.r.t. dtypes. Shapes
+// are used for both performance reasons, as much TensorFlow code specializes
+// on known shapes to produce slimmer graphs, and correctness, as some
+// high-level APIs require shapes to be fully-known.
+//
+// TODO(nareshmodi): Add support for sparse tensors.
+PyObject* TFE_Py_EncodeArg(PyObject* arg) {
+ EncodeResult result;
+ const auto status = TFE_Py_EncodeArgHelper(arg, &result);
+ if (MaybeRaiseExceptionFromStatus(status, nullptr)) {
+ return nullptr;
+ }
+
+ return result.ToPyTuple();
+}
diff --git a/tensorflow/python/estimator/BUILD b/tensorflow/python/estimator/BUILD
index 7f23499..1c4c595 100644
--- a/tensorflow/python/estimator/BUILD
+++ b/tensorflow/python/estimator/BUILD
@@ -281,6 +281,7 @@
":prediction_keys",
"//tensorflow:tensorflow_py_no_contrib",
"//third_party/py/numpy",
+ "@absl_py//absl/testing:parameterized",
"@six_archive//:six",
],
)
@@ -303,6 +304,7 @@
":pandas_io",
":prediction_keys",
"//tensorflow:tensorflow_py_no_contrib",
+ "@absl_py//absl/testing:parameterized",
"@six_archive//:six",
],
)
@@ -342,6 +344,7 @@
":pandas_io",
":prediction_keys",
"//tensorflow:tensorflow_py_no_contrib",
+ "@absl_py//absl/testing:parameterized",
"@six_archive//:six",
],
)
diff --git a/tensorflow/python/estimator/canned/dnn.py b/tensorflow/python/estimator/canned/dnn.py
index 1c0c458..a6c2aaa 100644
--- a/tensorflow/python/estimator/canned/dnn.py
+++ b/tensorflow/python/estimator/canned/dnn.py
@@ -24,7 +24,10 @@
from tensorflow.python.estimator import model_fn
from tensorflow.python.estimator.canned import head as head_lib
from tensorflow.python.estimator.canned import optimizers
-from tensorflow.python.feature_column import feature_column as feature_column_lib
+from tensorflow.python.feature_column import feature_column
+from tensorflow.python.feature_column import feature_column_v2
+from tensorflow.python.framework import ops
+from tensorflow.python.keras.engine import training
from tensorflow.python.layers import core as core_layers
from tensorflow.python.layers import normalization
from tensorflow.python.ops import init_ops
@@ -45,8 +48,14 @@
summary.histogram('%s/activation' % tag, value)
-def _dnn_logit_fn_builder(units, hidden_units, feature_columns, activation_fn,
- dropout, input_layer_partitioner, batch_norm):
+def _dnn_logit_fn_builder(units,
+ hidden_units,
+ feature_columns,
+ activation_fn,
+ dropout,
+ input_layer_partitioner,
+ batch_norm,
+ shared_state_manager=None):
"""Function builder for a dnn logit_fn.
Args:
@@ -60,6 +69,8 @@
coordinate.
input_layer_partitioner: Partitioner for input layer.
batch_norm: Whether to use batch normalization after each hidden layer.
+ shared_state_manager: A SharedEmbeddingStateManager object to hold the
+ shared state for SharedEmbeddingColumn's.
Returns:
A logit_fn (see below).
@@ -85,50 +96,129 @@
A `Tensor` representing the logits, or a list of `Tensor`'s representing
multiple logits in the MultiHead case.
"""
- is_training = mode == model_fn.ModeKeys.TRAIN
- with variable_scope.variable_scope(
- 'input_from_feature_columns',
- values=tuple(six.itervalues(features)),
- partitioner=input_layer_partitioner):
- net = feature_column_lib.input_layer(
- features=features, feature_columns=feature_columns)
+ dnn_model = _DNNModel(
+ units,
+ hidden_units,
+ feature_columns,
+ activation_fn,
+ dropout,
+ input_layer_partitioner,
+ batch_norm,
+ shared_state_manager,
+ name='dnn')
+ return dnn_model(features, mode)
+
+ return dnn_logit_fn
+
+
+def _get_previous_name_scope():
+ current_name_scope = ops.get_name_scope()
+ return current_name_scope.rsplit('/', 1)[0] + '/'
+
+
+class _DNNModel(training.Model):
+ """A DNN Model."""
+
+ def __init__(self,
+ units,
+ hidden_units,
+ feature_columns,
+ activation_fn,
+ dropout,
+ input_layer_partitioner,
+ batch_norm,
+ shared_state_manager,
+ name=None,
+ **kwargs):
+ super(_DNNModel, self).__init__(name=name, **kwargs)
+ if feature_column_v2.is_feature_column_v2(feature_columns):
+ self._input_layer = feature_column_v2.FeatureLayer(
+ feature_columns=feature_columns,
+ name='input_layer',
+ shared_state_manager=shared_state_manager)
+ else:
+ self._input_layer = feature_column.InputLayer(
+ feature_columns=feature_columns,
+ name='input_layer',
+ create_scope_now=False)
+
+ self._add_layer(self._input_layer, 'input_layer')
+
+ self._dropout = dropout
+ self._batch_norm = batch_norm
+
+ self._hidden_layers = []
+ self._dropout_layers = []
+ self._batch_norm_layers = []
+ self._hidden_layer_scope_names = []
for layer_id, num_hidden_units in enumerate(hidden_units):
with variable_scope.variable_scope(
- 'hiddenlayer_%d' % layer_id, values=(net,)) as hidden_layer_scope:
- net = core_layers.dense(
- net,
+ 'hiddenlayer_%d' % layer_id) as hidden_layer_scope:
+ hidden_layer = core_layers.Dense(
units=num_hidden_units,
activation=activation_fn,
kernel_initializer=init_ops.glorot_uniform_initializer(),
- name=hidden_layer_scope)
- if dropout is not None and is_training:
- net = core_layers.dropout(net, rate=dropout, training=True)
- if batch_norm:
- # TODO(hjm): In future, if this becomes popular, we can enable
- # customization of the batch normalization params by accepting a
- # list of `BatchNormalization` instances as `batch_norm`.
- net = normalization.batch_normalization(
- net,
+ name=hidden_layer_scope,
+ _scope=hidden_layer_scope)
+ self._add_layer(hidden_layer, hidden_layer_scope.name)
+ self._hidden_layer_scope_names.append(hidden_layer_scope.name)
+ self._hidden_layers.append(hidden_layer)
+ if self._dropout is not None:
+ dropout_layer = core_layers.Dropout(rate=self._dropout)
+ self._add_layer(dropout_layer, dropout_layer.name)
+ self._dropout_layers.append(dropout_layer)
+ if self._batch_norm:
+ batch_norm_layer = normalization.BatchNormalization(
# The default momentum 0.99 actually crashes on certain
# problem, so here we use 0.999, which is the default of
# tf.contrib.layers.batch_norm.
momentum=0.999,
- training=is_training,
- name='batchnorm_%d' % layer_id)
- _add_hidden_layer_summary(net, hidden_layer_scope.name)
+ trainable=True,
+ name='batchnorm_%d' % layer_id,
+ _scope='batchnorm_%d' % layer_id)
+ self._add_layer(batch_norm_layer, batch_norm_layer.name)
+ self._batch_norm_layers.append(batch_norm_layer)
- with variable_scope.variable_scope('logits', values=(net,)) as logits_scope:
- logits = core_layers.dense(
- net,
+ with variable_scope.variable_scope('logits') as logits_scope:
+ self._logits_layer = core_layers.Dense(
units=units,
activation=None,
kernel_initializer=init_ops.glorot_uniform_initializer(),
- name=logits_scope)
- _add_hidden_layer_summary(logits, logits_scope.name)
+ name=logits_scope,
+ _scope=logits_scope)
+ self._add_layer(self._logits_layer, logits_scope.name)
+ self._logits_scope_name = logits_scope.name
+ self._input_layer_partitioner = input_layer_partitioner
- return logits
+ def call(self, features, mode):
+ is_training = mode == model_fn.ModeKeys.TRAIN
+ # The Keras training.Model adds a name_scope with the name of the model
+ # which modifies the constructed graph. Hence we add another name_scope
+ # here which is the one before the training.Model one was applied.
+ # TODO(rohanj): Remove this in TF 2.0 (b/116728605)
+ with ops.name_scope(name=_get_previous_name_scope()):
+ # TODO(rohanj): Remove dependence on variable scope for partitioning.
+ with variable_scope.variable_scope(
+ 'input_from_feature_columns',
+ partitioner=self._input_layer_partitioner):
+ net = self._input_layer(features)
+ for i in range(len(self._hidden_layers)):
+ net = self._hidden_layers[i](net)
+ if self._dropout is not None and is_training:
+ net = self._dropout_layers[i](net, training=True)
+ if self._batch_norm:
+ net = self._batch_norm_layers[i](net, training=is_training)
+ _add_hidden_layer_summary(net, self._hidden_layer_scope_names[i])
- return dnn_logit_fn
+ logits = self._logits_layer(net)
+ _add_hidden_layer_summary(logits, self._logits_scope_name)
+ return logits
+
+ def _add_layer(self, layer, layer_name):
+ # "Magic" required for keras.Model classes to track all the variables in
+ # a list of layers.Layer objects.
+ # TODO(ashankar): Figure out API so user code doesn't have to do this.
+ setattr(self, layer_name, layer)
def _dnn_model_fn(features,
@@ -143,7 +233,8 @@
input_layer_partitioner=None,
config=None,
use_tpu=False,
- batch_norm=False):
+ batch_norm=False,
+ shared_state_manager=None):
"""Deep Neural Net model_fn.
Args:
@@ -167,6 +258,8 @@
use_tpu: Whether to make a DNN model able to run on TPU. Will make function
return a `_TPUEstimatorSpec` instance and disable variable partitioning.
batch_norm: Whether to use batch normalization after each hidden layer.
+ shared_state_manager: A SharedEmbeddingStateManager object to hold the
+ shared state for SharedEmbeddingColumn's.
Returns:
An `EstimatorSpec` instance.
@@ -202,7 +295,8 @@
activation_fn=activation_fn,
dropout=dropout,
input_layer_partitioner=input_layer_partitioner,
- batch_norm=batch_norm)
+ batch_norm=batch_norm,
+ shared_state_manager=shared_state_manager)
logits = logit_fn(features=features, mode=mode)
if use_tpu:
@@ -370,6 +464,10 @@
"""
head = head_lib._binary_logistic_or_multi_class_head( # pylint: disable=protected-access
n_classes, weight_column, label_vocabulary, loss_reduction)
+
+ shared_state_manager = feature_column_v2.maybe_create_shared_state_manager(
+ feature_columns)
+
def _model_fn(features, labels, mode, config):
"""Call the defined shared _dnn_model_fn."""
return _dnn_model_fn(
@@ -384,7 +482,8 @@
dropout=dropout,
input_layer_partitioner=input_layer_partitioner,
config=config,
- batch_norm=batch_norm)
+ batch_norm=batch_norm,
+ shared_state_manager=shared_state_manager)
super(DNNClassifier, self).__init__(
model_fn=_model_fn, model_dir=model_dir, config=config,
@@ -532,6 +631,10 @@
batch_norm: Whether to use batch normalization after each hidden layer.
"""
+ shared_state_manager = None
+ if feature_column_v2.is_feature_column_v2(feature_columns):
+ shared_state_manager = feature_column_v2.SharedEmbeddingStateManager()
+
def _model_fn(features, labels, mode, config):
"""Call the defined shared _dnn_model_fn."""
return _dnn_model_fn(
@@ -539,7 +642,8 @@
labels=labels,
mode=mode,
head=head_lib._regression_head( # pylint: disable=protected-access
- label_dimension=label_dimension, weight_column=weight_column,
+ label_dimension=label_dimension,
+ weight_column=weight_column,
loss_reduction=loss_reduction),
hidden_units=hidden_units,
feature_columns=tuple(feature_columns or []),
@@ -548,7 +652,8 @@
dropout=dropout,
input_layer_partitioner=input_layer_partitioner,
config=config,
- batch_norm=batch_norm)
+ batch_norm=batch_norm,
+ shared_state_manager=shared_state_manager)
super(DNNRegressor, self).__init__(
model_fn=_model_fn, model_dir=model_dir, config=config,
diff --git a/tensorflow/python/estimator/canned/dnn_linear_combined.py b/tensorflow/python/estimator/canned/dnn_linear_combined.py
index 9799cf9..f712244 100644
--- a/tensorflow/python/estimator/canned/dnn_linear_combined.py
+++ b/tensorflow/python/estimator/canned/dnn_linear_combined.py
@@ -27,6 +27,7 @@
from tensorflow.python.estimator.canned import head as head_lib
from tensorflow.python.estimator.canned import linear
from tensorflow.python.estimator.canned import optimizers
+from tensorflow.python.feature_column import feature_column_v2
from tensorflow.python.framework import ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import nn
@@ -142,6 +143,9 @@
max_partitions=num_ps_replicas,
min_slice_size=64 << 20))
+ shared_state_manager = feature_column_v2.maybe_create_shared_state_manager(
+ list(linear_feature_columns) + list(dnn_feature_columns))
+
# Build DNN Logits.
dnn_parent_scope = 'dnn'
@@ -169,8 +173,9 @@
feature_columns=dnn_feature_columns,
activation_fn=dnn_activation_fn,
dropout=dnn_dropout,
+ batch_norm=batch_norm,
input_layer_partitioner=input_layer_partitioner,
- batch_norm=batch_norm)
+ shared_state_manager=shared_state_manager)
dnn_logits = dnn_logit_fn(features=features, mode=mode)
linear_parent_scope = 'linear'
diff --git a/tensorflow/python/estimator/canned/dnn_linear_combined_test.py b/tensorflow/python/estimator/canned/dnn_linear_combined_test.py
index d163186..ae968e7 100644
--- a/tensorflow/python/estimator/canned/dnn_linear_combined_test.py
+++ b/tensorflow/python/estimator/canned/dnn_linear_combined_test.py
@@ -21,6 +21,7 @@
import shutil
import tempfile
+from absl.testing import parameterized
import numpy as np
import six
@@ -35,6 +36,7 @@
from tensorflow.python.estimator.inputs import numpy_io
from tensorflow.python.estimator.inputs import pandas_io
from tensorflow.python.feature_column import feature_column
+from tensorflow.python.feature_column import feature_column_v2
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import nn
@@ -119,7 +121,16 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
linear_testing_utils.BaseLinearRegressorPartitionerTest.__init__(
- self, _linear_regressor_fn)
+ self, _linear_regressor_fn, fc_lib=feature_column)
+
+
+class LinearOnlyRegressorPartitionerV2Test(
+ linear_testing_utils.BaseLinearRegressorPartitionerTest, test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ linear_testing_utils.BaseLinearRegressorPartitionerTest.__init__(
+ self, _linear_regressor_fn, fc_lib=feature_column_v2)
class LinearOnlyRegressorEvaluationTest(
@@ -128,7 +139,16 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
linear_testing_utils.BaseLinearRegressorEvaluationTest.__init__(
- self, _linear_regressor_fn)
+ self, _linear_regressor_fn, fc_lib=feature_column)
+
+
+class LinearOnlyRegressorEvaluationV2Test(
+ linear_testing_utils.BaseLinearRegressorEvaluationTest, test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ linear_testing_utils.BaseLinearRegressorEvaluationTest.__init__(
+ self, _linear_regressor_fn, fc_lib=feature_column_v2)
class LinearOnlyRegressorPredictTest(
@@ -137,7 +157,16 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
linear_testing_utils.BaseLinearRegressorPredictTest.__init__(
- self, _linear_regressor_fn)
+ self, _linear_regressor_fn, fc_lib=feature_column)
+
+
+class LinearOnlyRegressorPredictV2Test(
+ linear_testing_utils.BaseLinearRegressorPredictTest, test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ linear_testing_utils.BaseLinearRegressorPredictTest.__init__(
+ self, _linear_regressor_fn, fc_lib=feature_column_v2)
class LinearOnlyRegressorIntegrationTest(
@@ -146,7 +175,16 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
linear_testing_utils.BaseLinearRegressorIntegrationTest.__init__(
- self, _linear_regressor_fn)
+ self, _linear_regressor_fn, fc_lib=feature_column)
+
+
+class LinearOnlyRegressorIntegrationV2Test(
+ linear_testing_utils.BaseLinearRegressorIntegrationTest, test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ linear_testing_utils.BaseLinearRegressorIntegrationTest.__init__(
+ self, _linear_regressor_fn, fc_lib=feature_column_v2)
class LinearOnlyRegressorTrainingTest(
@@ -155,7 +193,16 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
linear_testing_utils.BaseLinearRegressorTrainingTest.__init__(
- self, _linear_regressor_fn)
+ self, _linear_regressor_fn, fc_lib=feature_column)
+
+
+class LinearOnlyRegressorTrainingV2Test(
+ linear_testing_utils.BaseLinearRegressorTrainingTest, test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ linear_testing_utils.BaseLinearRegressorTrainingTest.__init__(
+ self, _linear_regressor_fn, fc_lib=feature_column_v2)
def _linear_classifier_fn(feature_columns,
@@ -185,7 +232,18 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
linear_testing_utils.BaseLinearClassifierTrainingTest.__init__(
- self, linear_classifier_fn=_linear_classifier_fn)
+ self, linear_classifier_fn=_linear_classifier_fn, fc_lib=feature_column)
+
+
+class LinearOnlyClassifierTrainingV2Test(
+ linear_testing_utils.BaseLinearClassifierTrainingTest, test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ linear_testing_utils.BaseLinearClassifierTrainingTest.__init__(
+ self,
+ linear_classifier_fn=_linear_classifier_fn,
+ fc_lib=feature_column_v2)
class LinearOnlyClassifierClassesEvaluationTest(
@@ -194,7 +252,18 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
linear_testing_utils.BaseLinearClassifierEvaluationTest.__init__(
- self, linear_classifier_fn=_linear_classifier_fn)
+ self, linear_classifier_fn=_linear_classifier_fn, fc_lib=feature_column)
+
+
+class LinearOnlyClassifierClassesEvaluationV2Test(
+ linear_testing_utils.BaseLinearClassifierEvaluationTest, test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ linear_testing_utils.BaseLinearClassifierEvaluationTest.__init__(
+ self,
+ linear_classifier_fn=_linear_classifier_fn,
+ fc_lib=feature_column_v2)
class LinearOnlyClassifierPredictTest(
@@ -203,7 +272,18 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
linear_testing_utils.BaseLinearClassifierPredictTest.__init__(
- self, linear_classifier_fn=_linear_classifier_fn)
+ self, linear_classifier_fn=_linear_classifier_fn, fc_lib=feature_column)
+
+
+class LinearOnlyClassifierPredictV2Test(
+ linear_testing_utils.BaseLinearClassifierPredictTest, test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ linear_testing_utils.BaseLinearClassifierPredictTest.__init__(
+ self,
+ linear_classifier_fn=_linear_classifier_fn,
+ fc_lib=feature_column_v2)
class LinearOnlyClassifierIntegrationTest(
@@ -212,9 +292,21 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
linear_testing_utils.BaseLinearClassifierIntegrationTest.__init__(
- self, linear_classifier_fn=_linear_classifier_fn)
+ self, linear_classifier_fn=_linear_classifier_fn, fc_lib=feature_column)
+class LinearOnlyClassifierIntegrationV2Test(
+ linear_testing_utils.BaseLinearClassifierIntegrationTest, test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ linear_testing_utils.BaseLinearClassifierIntegrationTest.__init__(
+ self,
+ linear_classifier_fn=_linear_classifier_fn,
+ fc_lib=feature_column_v2)
+
+
+@parameterized.parameters((feature_column,), (feature_column_v2,))
class DNNLinearCombinedRegressorIntegrationTest(test.TestCase):
def setUp(self):
@@ -225,13 +317,15 @@
writer_cache.FileWriterCache.clear()
shutil.rmtree(self._model_dir)
- def _test_complete_flow(
- self, train_input_fn, eval_input_fn, predict_input_fn, input_dimension,
- label_dimension, batch_size):
+ def _test_complete_flow(self, train_input_fn, eval_input_fn, predict_input_fn,
+ input_dimension, label_dimension, batch_size,
+ fc_impl):
linear_feature_columns = [
- feature_column.numeric_column('x', shape=(input_dimension,))]
+ fc_impl.numeric_column('x', shape=(input_dimension,))
+ ]
dnn_feature_columns = [
- feature_column.numeric_column('x', shape=(input_dimension,))]
+ fc_impl.numeric_column('x', shape=(input_dimension,))
+ ]
feature_columns = linear_feature_columns + dnn_feature_columns
est = dnn_linear_combined.DNNLinearCombinedRegressor(
linear_feature_columns=linear_feature_columns,
@@ -257,14 +351,14 @@
self.assertAllEqual((batch_size, label_dimension), predictions.shape)
# EXPORT
- feature_spec = feature_column.make_parse_example_spec(feature_columns)
+ feature_spec = fc_impl.make_parse_example_spec(feature_columns)
serving_input_receiver_fn = export.build_parsing_serving_input_receiver_fn(
feature_spec)
export_dir = est.export_savedmodel(tempfile.mkdtemp(),
serving_input_receiver_fn)
self.assertTrue(gfile.Exists(export_dir))
- def test_numpy_input_fn(self):
+ def test_numpy_input_fn(self, fc_impl):
"""Tests complete flow with numpy_input_fn."""
label_dimension = 2
batch_size = 10
@@ -293,9 +387,10 @@
predict_input_fn=predict_input_fn,
input_dimension=label_dimension,
label_dimension=label_dimension,
- batch_size=batch_size)
+ batch_size=batch_size,
+ fc_impl=fc_impl)
- def test_pandas_input_fn(self):
+ def test_pandas_input_fn(self, fc_impl):
"""Tests complete flow with pandas_input_fn."""
if not HAS_PANDAS:
return
@@ -326,9 +421,10 @@
predict_input_fn=predict_input_fn,
input_dimension=label_dimension,
label_dimension=label_dimension,
- batch_size=batch_size)
+ batch_size=batch_size,
+ fc_impl=fc_impl)
- def test_input_fn_from_parse_example(self):
+ def test_input_fn_from_parse_example(self, fc_impl):
"""Tests complete flow with input_fn constructed from parse_example."""
label_dimension = 2
batch_size = 10
@@ -376,7 +472,8 @@
predict_input_fn=_predict_input_fn,
input_dimension=label_dimension,
label_dimension=label_dimension,
- batch_size=batch_size)
+ batch_size=batch_size,
+ fc_impl=fc_impl)
# A function to mimic dnn-classifier init reuse same tests.
@@ -407,7 +504,16 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
dnn_testing_utils.BaseDNNClassifierEvaluateTest.__init__(
- self, _dnn_classifier_fn)
+ self, _dnn_classifier_fn, fc_impl=feature_column)
+
+
+class DNNOnlyClassifierEvaluateV2Test(
+ dnn_testing_utils.BaseDNNClassifierEvaluateTest, test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ dnn_testing_utils.BaseDNNClassifierEvaluateTest.__init__(
+ self, _dnn_classifier_fn, fc_impl=feature_column_v2)
class DNNOnlyClassifierPredictTest(
@@ -416,7 +522,16 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
dnn_testing_utils.BaseDNNClassifierPredictTest.__init__(
- self, _dnn_classifier_fn)
+ self, _dnn_classifier_fn, fc_impl=feature_column)
+
+
+class DNNOnlyClassifierPredictV2Test(
+ dnn_testing_utils.BaseDNNClassifierPredictTest, test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ dnn_testing_utils.BaseDNNClassifierPredictTest.__init__(
+ self, _dnn_classifier_fn, fc_impl=feature_column_v2)
class DNNOnlyClassifierTrainTest(
@@ -425,7 +540,16 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
dnn_testing_utils.BaseDNNClassifierTrainTest.__init__(
- self, _dnn_classifier_fn)
+ self, _dnn_classifier_fn, fc_impl=feature_column)
+
+
+class DNNOnlyClassifierTrainV2Test(dnn_testing_utils.BaseDNNClassifierTrainTest,
+ test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ dnn_testing_utils.BaseDNNClassifierTrainTest.__init__(
+ self, _dnn_classifier_fn, fc_impl=feature_column_v2)
# A function to mimic dnn-regressor init reuse same tests.
@@ -454,7 +578,16 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
dnn_testing_utils.BaseDNNRegressorEvaluateTest.__init__(
- self, _dnn_regressor_fn)
+ self, _dnn_regressor_fn, fc_impl=feature_column)
+
+
+class DNNOnlyRegressorEvaluateV2Test(
+ dnn_testing_utils.BaseDNNRegressorEvaluateTest, test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ dnn_testing_utils.BaseDNNRegressorEvaluateTest.__init__(
+ self, _dnn_regressor_fn, fc_impl=feature_column_v2)
class DNNOnlyRegressorPredictTest(
@@ -463,7 +596,16 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
dnn_testing_utils.BaseDNNRegressorPredictTest.__init__(
- self, _dnn_regressor_fn)
+ self, _dnn_regressor_fn, fc_impl=feature_column)
+
+
+class DNNOnlyRegressorPredictV2Test(
+ dnn_testing_utils.BaseDNNRegressorPredictTest, test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ dnn_testing_utils.BaseDNNRegressorPredictTest.__init__(
+ self, _dnn_regressor_fn, fc_impl=feature_column_v2)
class DNNOnlyRegressorTrainTest(
@@ -472,9 +614,19 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
dnn_testing_utils.BaseDNNRegressorTrainTest.__init__(
- self, _dnn_regressor_fn)
+ self, _dnn_regressor_fn, fc_impl=feature_column)
+class DNNOnlyRegressorTrainV2Test(dnn_testing_utils.BaseDNNRegressorTrainTest,
+ test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ dnn_testing_utils.BaseDNNRegressorTrainTest.__init__(
+ self, _dnn_regressor_fn, fc_impl=feature_column_v2)
+
+
+@parameterized.parameters((feature_column,), (feature_column_v2,))
class DNNLinearCombinedClassifierIntegrationTest(test.TestCase):
def setUp(self):
@@ -488,13 +640,14 @@
def _as_label(self, data_in_float):
return np.rint(data_in_float).astype(np.int64)
- def _test_complete_flow(
- self, train_input_fn, eval_input_fn, predict_input_fn, input_dimension,
- n_classes, batch_size):
+ def _test_complete_flow(self, train_input_fn, eval_input_fn, predict_input_fn,
+ input_dimension, n_classes, batch_size, fc_impl):
linear_feature_columns = [
- feature_column.numeric_column('x', shape=(input_dimension,))]
+ fc_impl.numeric_column('x', shape=(input_dimension,))
+ ]
dnn_feature_columns = [
- feature_column.numeric_column('x', shape=(input_dimension,))]
+ fc_impl.numeric_column('x', shape=(input_dimension,))
+ ]
feature_columns = linear_feature_columns + dnn_feature_columns
est = dnn_linear_combined.DNNLinearCombinedClassifier(
linear_feature_columns=linear_feature_columns,
@@ -520,14 +673,14 @@
self.assertAllEqual((batch_size, n_classes), predicted_proba.shape)
# EXPORT
- feature_spec = feature_column.make_parse_example_spec(feature_columns)
+ feature_spec = fc_impl.make_parse_example_spec(feature_columns)
serving_input_receiver_fn = export.build_parsing_serving_input_receiver_fn(
feature_spec)
export_dir = est.export_savedmodel(tempfile.mkdtemp(),
serving_input_receiver_fn)
self.assertTrue(gfile.Exists(export_dir))
- def test_numpy_input_fn(self):
+ def test_numpy_input_fn(self, fc_impl):
"""Tests complete flow with numpy_input_fn."""
n_classes = 3
input_dimension = 2
@@ -559,9 +712,10 @@
predict_input_fn=predict_input_fn,
input_dimension=input_dimension,
n_classes=n_classes,
- batch_size=batch_size)
+ batch_size=batch_size,
+ fc_impl=fc_impl)
- def test_pandas_input_fn(self):
+ def test_pandas_input_fn(self, fc_impl):
"""Tests complete flow with pandas_input_fn."""
if not HAS_PANDAS:
return
@@ -593,9 +747,10 @@
predict_input_fn=predict_input_fn,
input_dimension=input_dimension,
n_classes=n_classes,
- batch_size=batch_size)
+ batch_size=batch_size,
+ fc_impl=fc_impl)
- def test_input_fn_from_parse_example(self):
+ def test_input_fn_from_parse_example(self, fc_impl):
"""Tests complete flow with input_fn constructed from parse_example."""
input_dimension = 2
n_classes = 3
@@ -647,9 +802,11 @@
predict_input_fn=_predict_input_fn,
input_dimension=input_dimension,
n_classes=n_classes,
- batch_size=batch_size)
+ batch_size=batch_size,
+ fc_impl=fc_impl)
+@parameterized.parameters((feature_column,), (feature_column_v2,))
class DNNLinearCombinedTests(test.TestCase):
def setUp(self):
@@ -681,9 +838,9 @@
return optimizer_mock
- def test_train_op_calls_both_dnn_and_linear(self):
+ def test_train_op_calls_both_dnn_and_linear(self, fc_impl):
opt = gradient_descent.GradientDescentOptimizer(1.)
- x_column = feature_column.numeric_column('x')
+ x_column = fc_impl.numeric_column('x')
input_fn = numpy_io.numpy_input_fn(
x={'x': np.array([[0.], [1.]])},
y=np.array([[0.], [1.]]),
@@ -708,7 +865,7 @@
checkpoint_utils.load_variable(
self._model_dir, 'dnn_called'))
- def test_dnn_and_linear_logits_are_added(self):
+ def test_dnn_and_linear_logits_are_added(self, fc_impl):
with ops.Graph().as_default():
variables_lib.Variable([[1.0]], name='linear/linear_model/x/weights')
variables_lib.Variable([2.0], name='linear/linear_model/bias_weights')
@@ -719,7 +876,7 @@
variables_lib.Variable(1, name='global_step', dtype=dtypes.int64)
linear_testing_utils.save_variables_to_ckpt(self._model_dir)
- x_column = feature_column.numeric_column('x')
+ x_column = fc_impl.numeric_column('x')
est = dnn_linear_combined.DNNLinearCombinedRegressor(
linear_feature_columns=[x_column],
dnn_hidden_units=[1],
@@ -737,6 +894,7 @@
next(est.predict(input_fn=input_fn)))
+@parameterized.parameters((feature_column,), (feature_column_v2,))
class DNNLinearCombinedWarmStartingTest(test.TestCase):
def setUp(self):
@@ -758,11 +916,11 @@
writer_cache.FileWriterCache.clear()
shutil.rmtree(self._ckpt_and_vocab_dir)
- def test_classifier_basic_warm_starting(self):
+ def test_classifier_basic_warm_starting(self, fc_impl):
"""Tests correctness of DNNLinearCombinedClassifier default warm-start."""
- age = feature_column.numeric_column('age')
- city = feature_column.embedding_column(
- feature_column.categorical_column_with_vocabulary_list(
+ age = fc_impl.numeric_column('age')
+ city = fc_impl.embedding_column(
+ fc_impl.categorical_column_with_vocabulary_list(
'city', vocabulary_list=['Mountain View', 'Palo Alto']),
dimension=5)
@@ -798,11 +956,11 @@
dnn_lc_classifier.get_variable_value(variable_name),
warm_started_dnn_lc_classifier.get_variable_value(variable_name))
- def test_regressor_basic_warm_starting(self):
+ def test_regressor_basic_warm_starting(self, fc_impl):
"""Tests correctness of DNNLinearCombinedRegressor default warm-start."""
- age = feature_column.numeric_column('age')
- city = feature_column.embedding_column(
- feature_column.categorical_column_with_vocabulary_list(
+ age = fc_impl.numeric_column('age')
+ city = fc_impl.embedding_column(
+ fc_impl.categorical_column_with_vocabulary_list(
'city', vocabulary_list=['Mountain View', 'Palo Alto']),
dimension=5)
@@ -836,11 +994,11 @@
dnn_lc_regressor.get_variable_value(variable_name),
warm_started_dnn_lc_regressor.get_variable_value(variable_name))
- def test_warm_starting_selective_variables(self):
+ def test_warm_starting_selective_variables(self, fc_impl):
"""Tests selecting variables to warm-start."""
- age = feature_column.numeric_column('age')
- city = feature_column.embedding_column(
- feature_column.categorical_column_with_vocabulary_list(
+ age = fc_impl.numeric_column('age')
+ city = fc_impl.embedding_column(
+ fc_impl.categorical_column_with_vocabulary_list(
'city', vocabulary_list=['Mountain View', 'Palo Alto']),
dimension=5)
diff --git a/tensorflow/python/estimator/canned/dnn_test.py b/tensorflow/python/estimator/canned/dnn_test.py
index fc90b7c..756696c 100644
--- a/tensorflow/python/estimator/canned/dnn_test.py
+++ b/tensorflow/python/estimator/canned/dnn_test.py
@@ -21,6 +21,7 @@
import shutil
import tempfile
+from absl.testing import parameterized
import numpy as np
import six
@@ -33,6 +34,7 @@
from tensorflow.python.estimator.inputs import numpy_io
from tensorflow.python.estimator.inputs import pandas_io
from tensorflow.python.feature_column import feature_column
+from tensorflow.python.feature_column import feature_column_v2
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import data_flow_ops
@@ -62,15 +64,32 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
- dnn_testing_utils.BaseDNNModelFnTest.__init__(self, dnn._dnn_model_fn)
+ dnn_testing_utils.BaseDNNModelFnTest.__init__(
+ self, dnn._dnn_model_fn, fc_impl=feature_column)
+
+
+class DNNModelFnV2Test(dnn_testing_utils.BaseDNNModelFnTest, test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ dnn_testing_utils.BaseDNNModelFnTest.__init__(
+ self, dnn._dnn_model_fn, fc_impl=feature_column_v2)
class DNNLogitFnTest(dnn_testing_utils.BaseDNNLogitFnTest, test.TestCase):
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
- dnn_testing_utils.BaseDNNLogitFnTest.__init__(self,
- dnn._dnn_logit_fn_builder)
+ dnn_testing_utils.BaseDNNLogitFnTest.__init__(
+ self, dnn._dnn_logit_fn_builder, fc_impl=feature_column)
+
+
+class DNNLogitFnV2Test(dnn_testing_utils.BaseDNNLogitFnTest, test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ dnn_testing_utils.BaseDNNLogitFnTest.__init__(
+ self, dnn._dnn_logit_fn_builder, fc_impl=feature_column_v2)
class DNNWarmStartingTest(dnn_testing_utils.BaseDNNWarmStartingTest,
@@ -78,8 +97,17 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
- dnn_testing_utils.BaseDNNWarmStartingTest.__init__(self, _dnn_classifier_fn,
- _dnn_regressor_fn)
+ dnn_testing_utils.BaseDNNWarmStartingTest.__init__(
+ self, _dnn_classifier_fn, _dnn_regressor_fn, fc_impl=feature_column)
+
+
+class DNNWarmStartingV2Test(dnn_testing_utils.BaseDNNWarmStartingTest,
+ test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ dnn_testing_utils.BaseDNNWarmStartingTest.__init__(
+ self, _dnn_classifier_fn, _dnn_regressor_fn, fc_impl=feature_column_v2)
class DNNClassifierEvaluateTest(
@@ -88,7 +116,16 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
dnn_testing_utils.BaseDNNClassifierEvaluateTest.__init__(
- self, _dnn_classifier_fn)
+ self, _dnn_classifier_fn, fc_impl=feature_column)
+
+
+class DNNClassifierEvaluateV2Test(
+ dnn_testing_utils.BaseDNNClassifierEvaluateTest, test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ dnn_testing_utils.BaseDNNClassifierEvaluateTest.__init__(
+ self, _dnn_classifier_fn, fc_impl=feature_column_v2)
class DNNClassifierPredictTest(
@@ -97,7 +134,16 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
dnn_testing_utils.BaseDNNClassifierPredictTest.__init__(
- self, _dnn_classifier_fn)
+ self, _dnn_classifier_fn, fc_impl=feature_column)
+
+
+class DNNClassifierPredictV2Test(dnn_testing_utils.BaseDNNClassifierPredictTest,
+ test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ dnn_testing_utils.BaseDNNClassifierPredictTest.__init__(
+ self, _dnn_classifier_fn, fc_impl=feature_column_v2)
class DNNClassifierTrainTest(
@@ -106,7 +152,16 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
dnn_testing_utils.BaseDNNClassifierTrainTest.__init__(
- self, _dnn_classifier_fn)
+ self, _dnn_classifier_fn, fc_impl=feature_column)
+
+
+class DNNClassifierTrainV2Test(dnn_testing_utils.BaseDNNClassifierTrainTest,
+ test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ dnn_testing_utils.BaseDNNClassifierTrainTest.__init__(
+ self, _dnn_classifier_fn, fc_impl=feature_column_v2)
def _dnn_regressor_fn(*args, **kwargs):
@@ -119,7 +174,16 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
dnn_testing_utils.BaseDNNRegressorEvaluateTest.__init__(
- self, _dnn_regressor_fn)
+ self, _dnn_regressor_fn, fc_impl=feature_column)
+
+
+class DNNRegressorEvaluateV2Test(dnn_testing_utils.BaseDNNRegressorEvaluateTest,
+ test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ dnn_testing_utils.BaseDNNRegressorEvaluateTest.__init__(
+ self, _dnn_regressor_fn, fc_impl=feature_column_v2)
class DNNRegressorPredictTest(
@@ -128,7 +192,16 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
dnn_testing_utils.BaseDNNRegressorPredictTest.__init__(
- self, _dnn_regressor_fn)
+ self, _dnn_regressor_fn, fc_impl=feature_column)
+
+
+class DNNRegressorPredictV2Test(dnn_testing_utils.BaseDNNRegressorPredictTest,
+ test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ dnn_testing_utils.BaseDNNRegressorPredictTest.__init__(
+ self, _dnn_regressor_fn, fc_impl=feature_column_v2)
class DNNRegressorTrainTest(
@@ -137,7 +210,16 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
dnn_testing_utils.BaseDNNRegressorTrainTest.__init__(
- self, _dnn_regressor_fn)
+ self, _dnn_regressor_fn, fc_impl=feature_column)
+
+
+class DNNRegressorTrainV2Test(dnn_testing_utils.BaseDNNRegressorTrainTest,
+ test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ dnn_testing_utils.BaseDNNRegressorTrainTest.__init__(
+ self, _dnn_regressor_fn, fc_impl=feature_column_v2)
def _queue_parsed_features(feature_map):
@@ -156,7 +238,8 @@
return {keys[i]: dequeued_tensors[i] for i in range(len(dequeued_tensors))}
-class DNNRegressorIntegrationTest(test.TestCase):
+@parameterized.parameters((feature_column,), (feature_column_v2,))
+class DNNRegressorIntegrationTest(test.TestCase, parameterized.TestCase):
def setUp(self):
self._model_dir = tempfile.mkdtemp()
@@ -166,11 +249,11 @@
writer_cache.FileWriterCache.clear()
shutil.rmtree(self._model_dir)
- def _test_complete_flow(
- self, train_input_fn, eval_input_fn, predict_input_fn, input_dimension,
- label_dimension, batch_size):
- feature_columns = [
- feature_column.numeric_column('x', shape=(input_dimension,))]
+ def _test_complete_flow(self, train_input_fn, eval_input_fn, predict_input_fn,
+ input_dimension, label_dimension, batch_size,
+ fc_impl):
+ feature_columns = [fc_impl.numeric_column('x', shape=(input_dimension,))]
+
est = dnn.DNNRegressor(
hidden_units=(2, 2),
feature_columns=feature_columns,
@@ -194,14 +277,14 @@
self.assertAllEqual((batch_size, label_dimension), predictions.shape)
# EXPORT
- feature_spec = feature_column.make_parse_example_spec(feature_columns)
+ feature_spec = fc_impl.make_parse_example_spec(feature_columns)
serving_input_receiver_fn = export.build_parsing_serving_input_receiver_fn(
feature_spec)
export_dir = est.export_savedmodel(tempfile.mkdtemp(),
serving_input_receiver_fn)
self.assertTrue(gfile.Exists(export_dir))
- def test_numpy_input_fn(self):
+ def test_numpy_input_fn(self, fc_impl):
"""Tests complete flow with numpy_input_fn."""
label_dimension = 2
batch_size = 10
@@ -230,9 +313,10 @@
predict_input_fn=predict_input_fn,
input_dimension=label_dimension,
label_dimension=label_dimension,
- batch_size=batch_size)
+ batch_size=batch_size,
+ fc_impl=fc_impl)
- def test_pandas_input_fn(self):
+ def test_pandas_input_fn(self, fc_impl):
"""Tests complete flow with pandas_input_fn."""
if not HAS_PANDAS:
return
@@ -263,9 +347,10 @@
predict_input_fn=predict_input_fn,
input_dimension=label_dimension,
label_dimension=label_dimension,
- batch_size=batch_size)
+ batch_size=batch_size,
+ fc_impl=fc_impl)
- def test_input_fn_from_parse_example(self):
+ def test_input_fn_from_parse_example(self, fc_impl):
"""Tests complete flow with input_fn constructed from parse_example."""
label_dimension = 2
batch_size = 10
@@ -313,9 +398,11 @@
predict_input_fn=_predict_input_fn,
input_dimension=label_dimension,
label_dimension=label_dimension,
- batch_size=batch_size)
+ batch_size=batch_size,
+ fc_impl=fc_impl)
+@parameterized.parameters((feature_column,), (feature_column_v2,))
class DNNClassifierIntegrationTest(test.TestCase):
def setUp(self):
@@ -329,11 +416,10 @@
def _as_label(self, data_in_float):
return np.rint(data_in_float).astype(np.int64)
- def _test_complete_flow(
- self, train_input_fn, eval_input_fn, predict_input_fn, input_dimension,
- n_classes, batch_size):
- feature_columns = [
- feature_column.numeric_column('x', shape=(input_dimension,))]
+ def _test_complete_flow(self, train_input_fn, eval_input_fn, predict_input_fn,
+ input_dimension, n_classes, batch_size, fc_impl):
+ feature_columns = [fc_impl.numeric_column('x', shape=(input_dimension,))]
+
est = dnn.DNNClassifier(
hidden_units=(2, 2),
feature_columns=feature_columns,
@@ -357,14 +443,14 @@
self.assertAllEqual((batch_size, n_classes), predicted_proba.shape)
# EXPORT
- feature_spec = feature_column.make_parse_example_spec(feature_columns)
+ feature_spec = fc_impl.make_parse_example_spec(feature_columns)
serving_input_receiver_fn = export.build_parsing_serving_input_receiver_fn(
feature_spec)
export_dir = est.export_savedmodel(tempfile.mkdtemp(),
serving_input_receiver_fn)
self.assertTrue(gfile.Exists(export_dir))
- def test_numpy_input_fn(self):
+ def test_numpy_input_fn(self, fc_impl):
"""Tests complete flow with numpy_input_fn."""
n_classes = 3
input_dimension = 2
@@ -396,9 +482,10 @@
predict_input_fn=predict_input_fn,
input_dimension=input_dimension,
n_classes=n_classes,
- batch_size=batch_size)
+ batch_size=batch_size,
+ fc_impl=fc_impl)
- def test_pandas_input_fn(self):
+ def test_pandas_input_fn(self, fc_impl):
"""Tests complete flow with pandas_input_fn."""
if not HAS_PANDAS:
return
@@ -430,9 +517,10 @@
predict_input_fn=predict_input_fn,
input_dimension=input_dimension,
n_classes=n_classes,
- batch_size=batch_size)
+ batch_size=batch_size,
+ fc_impl=fc_impl)
- def test_input_fn_from_parse_example(self):
+ def test_input_fn_from_parse_example(self, fc_impl):
"""Tests complete flow with input_fn constructed from parse_example."""
input_dimension = 2
n_classes = 3
@@ -484,7 +572,8 @@
predict_input_fn=_predict_input_fn,
input_dimension=input_dimension,
n_classes=n_classes,
- batch_size=batch_size)
+ batch_size=batch_size,
+ fc_impl=fc_impl)
if __name__ == '__main__':
diff --git a/tensorflow/python/estimator/canned/dnn_testing_utils.py b/tensorflow/python/estimator/canned/dnn_testing_utils.py
index 11f1e93..cd66d0a 100644
--- a/tensorflow/python/estimator/canned/dnn_testing_utils.py
+++ b/tensorflow/python/estimator/canned/dnn_testing_utils.py
@@ -104,6 +104,7 @@
weights_and_biases: Iterable of tuples of weight and bias values.
global_step: Initial global step to save in checkpoint.
model_dir: Directory into which checkpoint is saved.
+ batch_norm_vars: Variables used for batch normalization.
"""
weights, biases = zip(*weights_and_biases)
if batch_norm_vars:
@@ -244,8 +245,9 @@
class BaseDNNModelFnTest(object):
"""Tests that _dnn_model_fn passes expected logits to mock head."""
- def __init__(self, dnn_model_fn):
+ def __init__(self, dnn_model_fn, fc_impl=feature_column):
self._dnn_model_fn = dnn_model_fn
+ self._fc_impl = fc_impl
def setUp(self):
self._model_dir = tempfile.mkdtemp()
@@ -272,7 +274,7 @@
head=head,
hidden_units=hidden_units,
feature_columns=[
- feature_column.numeric_column(
+ self._fc_impl.numeric_column(
'age', shape=np.array(inputs).shape[1:])
],
optimizer=mock_optimizer(self, hidden_units))
@@ -462,8 +464,8 @@
head=head,
hidden_units=hidden_units,
feature_columns=[
- feature_column.numeric_column('age'),
- feature_column.numeric_column('height')
+ self._fc_impl.numeric_column('age'),
+ self._fc_impl.numeric_column('height')
],
optimizer=mock_optimizer(self, hidden_units))
with monitored_session.MonitoredTrainingSession(
@@ -499,7 +501,7 @@
head=head,
hidden_units=hidden_units,
feature_columns=[
- feature_column.numeric_column(
+ self._fc_impl.numeric_column(
'age', shape=np.array(inputs).shape[1:])
],
optimizer=mock_optimizer(self, hidden_units))
@@ -508,8 +510,9 @@
class BaseDNNLogitFnTest(object):
"""Tests correctness of logits calculated from _dnn_logit_fn_builder."""
- def __init__(self, dnn_logit_fn_builder):
+ def __init__(self, dnn_logit_fn_builder, fc_impl=feature_column):
self._dnn_logit_fn_builder = dnn_logit_fn_builder
+ self._fc_impl = fc_impl
def setUp(self):
self._model_dir = tempfile.mkdtemp()
@@ -541,7 +544,7 @@
units=logits_dimension,
hidden_units=hidden_units,
feature_columns=[
- feature_column.numeric_column(
+ self._fc_impl.numeric_column(
'age', shape=np.array(inputs).shape[1:])
],
activation_fn=nn.relu,
@@ -786,8 +789,8 @@
units=logits_dimension,
hidden_units=hidden_units,
feature_columns=[
- feature_column.numeric_column('age'),
- feature_column.numeric_column('height')
+ self._fc_impl.numeric_column('age'),
+ self._fc_impl.numeric_column('height')
],
activation_fn=nn.relu,
dropout=None,
@@ -806,9 +809,13 @@
class BaseDNNWarmStartingTest(object):
- def __init__(self, _dnn_classifier_fn, _dnn_regressor_fn):
+ def __init__(self,
+ _dnn_classifier_fn,
+ _dnn_regressor_fn,
+ fc_impl=feature_column):
self._dnn_classifier_fn = _dnn_classifier_fn
self._dnn_regressor_fn = _dnn_regressor_fn
+ self._fc_impl = fc_impl
def setUp(self):
# Create a directory to save our old checkpoint and vocabularies to.
@@ -843,8 +850,8 @@
def test_classifier_basic_warm_starting(self):
"""Tests correctness of DNNClassifier default warm-start."""
- city = feature_column.embedding_column(
- feature_column.categorical_column_with_vocabulary_list(
+ city = self._fc_impl.embedding_column(
+ self._fc_impl.categorical_column_with_vocabulary_list(
'city', vocabulary_list=['Mountain View', 'Palo Alto']),
dimension=5)
@@ -875,8 +882,8 @@
def test_regressor_basic_warm_starting(self):
"""Tests correctness of DNNRegressor default warm-start."""
- city = feature_column.embedding_column(
- feature_column.categorical_column_with_vocabulary_list(
+ city = self._fc_impl.embedding_column(
+ self._fc_impl.categorical_column_with_vocabulary_list(
'city', vocabulary_list=['Mountain View', 'Palo Alto']),
dimension=5)
@@ -905,8 +912,8 @@
def test_warm_starting_selective_variables(self):
"""Tests selecting variables to warm-start."""
- city = feature_column.embedding_column(
- feature_column.categorical_column_with_vocabulary_list(
+ city = self._fc_impl.embedding_column(
+ self._fc_impl.categorical_column_with_vocabulary_list(
'city', vocabulary_list=['Mountain View', 'Palo Alto']),
dimension=5)
@@ -958,8 +965,8 @@
vocab_file = os.path.join(self._ckpt_and_vocab_dir, 'occupation_vocab')
with open(vocab_file, 'w') as f:
f.write('\n'.join(vocab_list))
- occupation = feature_column.embedding_column(
- feature_column.categorical_column_with_vocabulary_file(
+ occupation = self._fc_impl.embedding_column(
+ self._fc_impl.categorical_column_with_vocabulary_file(
'occupation',
vocabulary_file=vocab_file,
vocabulary_size=len(vocab_list)),
@@ -985,8 +992,8 @@
'new_occupation_vocab')
with open(new_vocab_file, 'w') as f:
f.write('\n'.join(new_vocab_list))
- new_occupation = feature_column.embedding_column(
- feature_column.categorical_column_with_vocabulary_file(
+ new_occupation = self._fc_impl.embedding_column(
+ self._fc_impl.categorical_column_with_vocabulary_file(
'occupation',
vocabulary_file=new_vocab_file,
vocabulary_size=len(new_vocab_list)),
@@ -1051,8 +1058,8 @@
def test_warm_starting_with_naming_change(self):
"""Tests warm-starting with a Tensor name remapping."""
- locality = feature_column.embedding_column(
- feature_column.categorical_column_with_vocabulary_list(
+ locality = self._fc_impl.embedding_column(
+ self._fc_impl.categorical_column_with_vocabulary_list(
'locality', vocabulary_list=['Mountain View', 'Palo Alto']),
dimension=5)
@@ -1068,8 +1075,8 @@
# Create a second DNNClassifier, warm-started from the first. Use a
# learning_rate = 0.0 optimizer to check values (use SGD so we don't have
# accumulator values that change).
- city = feature_column.embedding_column(
- feature_column.categorical_column_with_vocabulary_list(
+ city = self._fc_impl.embedding_column(
+ self._fc_impl.categorical_column_with_vocabulary_list(
'city', vocabulary_list=['Mountain View', 'Palo Alto']),
dimension=5)
warm_started_dnn_classifier = self._dnn_classifier_fn(
@@ -1101,8 +1108,9 @@
class BaseDNNClassifierEvaluateTest(object):
- def __init__(self, dnn_classifier_fn):
+ def __init__(self, dnn_classifier_fn, fc_impl=feature_column):
self._dnn_classifier_fn = dnn_classifier_fn
+ self._fc_impl = fc_impl
def setUp(self):
self._model_dir = tempfile.mkdtemp()
@@ -1121,7 +1129,7 @@
dnn_classifier = self._dnn_classifier_fn(
hidden_units=(2, 2),
- feature_columns=[feature_column.numeric_column('age')],
+ feature_columns=[self._fc_impl.numeric_column('age')],
model_dir=self._model_dir)
def _input_fn():
# batch_size = 2, one false label, and one true.
@@ -1161,7 +1169,7 @@
dnn_classifier = self._dnn_classifier_fn(
hidden_units=(2, 2),
- feature_columns=[feature_column.numeric_column('age', shape=[2])],
+ feature_columns=[self._fc_impl.numeric_column('age', shape=[2])],
n_classes=n_classes,
model_dir=self._model_dir)
def _input_fn():
@@ -1192,7 +1200,7 @@
dnn_classifier = self._dnn_classifier_fn(
hidden_units=(2, 2),
- feature_columns=[feature_column.numeric_column('age')],
+ feature_columns=[self._fc_impl.numeric_column('age')],
model_dir=self._model_dir)
def _input_fn():
# batch_size = 2, one false label, and one true.
@@ -1218,7 +1226,7 @@
dnn_classifier = self._dnn_classifier_fn(
hidden_units=(2, 2),
- feature_columns=[feature_column.numeric_column('age', shape=[2])],
+ feature_columns=[self._fc_impl.numeric_column('age', shape=[2])],
n_classes=n_classes,
weight_column='w',
model_dir=self._model_dir)
@@ -1238,8 +1246,9 @@
class BaseDNNRegressorEvaluateTest(object):
- def __init__(self, dnn_regressor_fn):
+ def __init__(self, dnn_regressor_fn, fc_impl=feature_column):
self._dnn_regressor_fn = dnn_regressor_fn
+ self._fc_impl = fc_impl
def setUp(self):
self._model_dir = tempfile.mkdtemp()
@@ -1259,7 +1268,7 @@
dnn_regressor = self._dnn_regressor_fn(
hidden_units=(2, 2),
- feature_columns=[feature_column.numeric_column('age')],
+ feature_columns=[self._fc_impl.numeric_column('age')],
model_dir=self._model_dir)
def _input_fn():
return {'age': [[10.]]}, [[1.]]
@@ -1289,7 +1298,7 @@
dnn_regressor = self._dnn_regressor_fn(
hidden_units=(2, 2),
- feature_columns=[feature_column.numeric_column('age', shape=[2])],
+ feature_columns=[self._fc_impl.numeric_column('age', shape=[2])],
label_dimension=label_dimension,
model_dir=self._model_dir)
def _input_fn():
@@ -1320,7 +1329,7 @@
dnn_regressor = self._dnn_regressor_fn(
hidden_units=(2, 2),
- feature_columns=[feature_column.numeric_column('age', shape=[2])],
+ feature_columns=[self._fc_impl.numeric_column('age', shape=[2])],
label_dimension=label_dimension,
weight_column='w',
model_dir=self._model_dir)
@@ -1339,8 +1348,9 @@
class BaseDNNClassifierPredictTest(object):
- def __init__(self, dnn_classifier_fn):
+ def __init__(self, dnn_classifier_fn, fc_impl=feature_column):
self._dnn_classifier_fn = dnn_classifier_fn
+ self._fc_impl = fc_impl
def setUp(self):
self._model_dir = tempfile.mkdtemp()
@@ -1361,7 +1371,7 @@
dnn_classifier = self._dnn_classifier_fn(
hidden_units=(2, 2),
label_vocabulary=label_vocabulary,
- feature_columns=(feature_column.numeric_column('x'),),
+ feature_columns=(self._fc_impl.numeric_column('x'),),
model_dir=self._model_dir)
input_fn = numpy_io.numpy_input_fn(
x={'x': np.array([[10.]])}, batch_size=1, shuffle=False)
@@ -1405,7 +1415,7 @@
dnn_classifier = self._dnn_classifier_fn(
hidden_units=(2, 2),
- feature_columns=(feature_column.numeric_column('x', shape=(2,)),),
+ feature_columns=(self._fc_impl.numeric_column('x', shape=(2,)),),
label_vocabulary=label_vocabulary,
n_classes=3,
model_dir=self._model_dir)
@@ -1453,8 +1463,9 @@
class BaseDNNRegressorPredictTest(object):
- def __init__(self, dnn_regressor_fn):
+ def __init__(self, dnn_regressor_fn, fc_impl=feature_column):
self._dnn_regressor_fn = dnn_regressor_fn
+ self._fc_impl = fc_impl
def setUp(self):
self._model_dir = tempfile.mkdtemp()
@@ -1475,7 +1486,7 @@
dnn_regressor = self._dnn_regressor_fn(
hidden_units=(2, 2),
- feature_columns=(feature_column.numeric_column('x'),),
+ feature_columns=(self._fc_impl.numeric_column('x'),),
model_dir=self._model_dir)
input_fn = numpy_io.numpy_input_fn(
x={'x': np.array([[10.]])}, batch_size=1, shuffle=False)
@@ -1497,7 +1508,7 @@
dnn_regressor = self._dnn_regressor_fn(
hidden_units=(2, 2),
- feature_columns=(feature_column.numeric_column('x', shape=(2,)),),
+ feature_columns=(self._fc_impl.numeric_column('x', shape=(2,)),),
label_dimension=3,
model_dir=self._model_dir)
input_fn = numpy_io.numpy_input_fn(
@@ -1594,8 +1605,9 @@
class BaseDNNClassifierTrainTest(object):
- def __init__(self, dnn_classifier_fn):
+ def __init__(self, dnn_classifier_fn, fc_impl=feature_column):
self._dnn_classifier_fn = dnn_classifier_fn
+ self._fc_impl = fc_impl
def setUp(self):
self._model_dir = tempfile.mkdtemp()
@@ -1609,7 +1621,7 @@
hidden_units = (2, 2)
dnn_classifier = self._dnn_classifier_fn(
hidden_units=hidden_units,
- feature_columns=(feature_column.numeric_column('age'),),
+ feature_columns=(self._fc_impl.numeric_column('age'),),
model_dir=self._model_dir)
# Train for a few steps, then validate final checkpoint.
@@ -1625,7 +1637,7 @@
n_classes = 3
dnn_classifier = self._dnn_classifier_fn(
hidden_units=hidden_units,
- feature_columns=(feature_column.numeric_column('age'),),
+ feature_columns=(self._fc_impl.numeric_column('age'),),
n_classes=n_classes,
model_dir=self._model_dir)
@@ -1643,7 +1655,7 @@
self, hidden_units=hidden_units)
dnn_classifier = self._dnn_classifier_fn(
hidden_units=hidden_units,
- feature_columns=(feature_column.numeric_column('age'),),
+ feature_columns=(self._fc_impl.numeric_column('age'),),
optimizer=opt,
model_dir=self._model_dir)
self.assertEqual(0, opt.minimize.call_count)
@@ -1682,7 +1694,7 @@
self, hidden_units=hidden_units, expected_loss=expected_loss)
dnn_classifier = self._dnn_classifier_fn(
hidden_units=hidden_units,
- feature_columns=(feature_column.numeric_column('age'),),
+ feature_columns=(self._fc_impl.numeric_column('age'),),
optimizer=opt,
model_dir=self._model_dir)
self.assertEqual(0, opt.minimize.call_count)
@@ -1728,7 +1740,7 @@
self, hidden_units=hidden_units, expected_loss=expected_loss)
dnn_classifier = self._dnn_classifier_fn(
hidden_units=hidden_units,
- feature_columns=(feature_column.numeric_column('age'),),
+ feature_columns=(self._fc_impl.numeric_column('age'),),
optimizer=opt,
model_dir=self._model_dir)
self.assertEqual(0, opt.minimize.call_count)
@@ -1759,7 +1771,7 @@
dnn_classifier = self._dnn_classifier_fn(
n_classes=n_classes,
hidden_units=hidden_units,
- feature_columns=(feature_column.numeric_column('age'),),
+ feature_columns=(self._fc_impl.numeric_column('age'),),
optimizer=opt,
model_dir=self._model_dir)
self.assertEqual(0, opt.minimize.call_count)
@@ -1793,8 +1805,9 @@
class BaseDNNRegressorTrainTest(object):
- def __init__(self, dnn_regressor_fn):
+ def __init__(self, dnn_regressor_fn, fc_impl=feature_column):
self._dnn_regressor_fn = dnn_regressor_fn
+ self._fc_impl = fc_impl
def setUp(self):
self._model_dir = tempfile.mkdtemp()
@@ -1808,7 +1821,7 @@
hidden_units = (2, 2)
dnn_regressor = self._dnn_regressor_fn(
hidden_units=hidden_units,
- feature_columns=(feature_column.numeric_column('age'),),
+ feature_columns=(self._fc_impl.numeric_column('age'),),
model_dir=self._model_dir)
# Train for a few steps, then validate final checkpoint.
@@ -1824,7 +1837,7 @@
opt = mock_optimizer(self, hidden_units=hidden_units)
dnn_regressor = self._dnn_regressor_fn(
hidden_units=hidden_units,
- feature_columns=(feature_column.numeric_column('age'),),
+ feature_columns=(self._fc_impl.numeric_column('age'),),
optimizer=opt,
model_dir=self._model_dir)
self.assertEqual(0, opt.minimize.call_count)
@@ -1864,7 +1877,7 @@
self, hidden_units=hidden_units, expected_loss=expected_loss)
dnn_regressor = self._dnn_regressor_fn(
hidden_units=hidden_units,
- feature_columns=(feature_column.numeric_column('age'),),
+ feature_columns=(self._fc_impl.numeric_column('age'),),
optimizer=opt,
model_dir=self._model_dir)
self.assertEqual(0, opt.minimize.call_count)
@@ -1917,7 +1930,8 @@
dnn_regressor = self._dnn_regressor_fn(
hidden_units=hidden_units,
feature_columns=[
- feature_column.numeric_column('age', shape=[input_dimension])],
+ self._fc_impl.numeric_column('age', shape=[input_dimension])
+ ],
label_dimension=label_dimension,
optimizer=opt,
model_dir=self._model_dir)
diff --git a/tensorflow/python/estimator/canned/linear.py b/tensorflow/python/estimator/canned/linear.py
index 115dd18..8b96284 100644
--- a/tensorflow/python/estimator/canned/linear.py
+++ b/tensorflow/python/estimator/canned/linear.py
@@ -25,14 +25,18 @@
from tensorflow.python.estimator import estimator
from tensorflow.python.estimator.canned import head as head_lib
from tensorflow.python.estimator.canned import optimizers
-from tensorflow.python.feature_column import feature_column as feature_column_lib
+from tensorflow.python.feature_column import feature_column
+from tensorflow.python.feature_column import feature_column_v2
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import nn
from tensorflow.python.ops import partitioned_variables
+from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import variables as variable_ops
from tensorflow.python.ops.losses import losses
from tensorflow.python.summary import summary
from tensorflow.python.training import ftrl
+from tensorflow.python.util import nest
from tensorflow.python.util.tf_export import estimator_export
@@ -46,23 +50,42 @@
return ftrl.FtrlOptimizer(learning_rate=learning_rate)
-def _compute_fraction_of_zero(cols_to_vars):
- """Given a linear cols_to_vars dict, compute the fraction of zero weights.
+def _get_expanded_variable_list(var_list):
+ """Given a list of variables, expands them if they are partitioned.
Args:
- cols_to_vars: A dictionary mapping FeatureColumns to lists of tf.Variables
- like one returned from feature_column_lib.linear_model.
+ var_list: A list of variables.
+
+ Returns:
+ A list of variables where each partitioned variable is expanded to its
+ components.
+ """
+ returned_list = []
+ for variable in var_list:
+ if (isinstance(variable, variable_ops.Variable) or
+ resource_variable_ops.is_resource_variable(variable)):
+ returned_list.append(variable) # Single variable case.
+ else: # Must be a PartitionedVariable, so convert into a list.
+ returned_list.extend(list(variable))
+ return returned_list
+
+
+# TODO(rohanj): Consider making this a public utility method.
+def _compute_fraction_of_zero(variables):
+ """Given a linear variables list, compute the fraction of zero weights.
+
+ Args:
+ variables: A list or list of list of variables
Returns:
The fraction of zeros (sparsity) in the linear model.
"""
all_weight_vars = []
- for var_or_var_list in cols_to_vars.values():
+ for var_or_var_list in variables:
+ var_list = nest.flatten(var_or_var_list)
# Skip empty-lists associated with columns that created no Variables.
- if var_or_var_list:
- all_weight_vars += [
- array_ops.reshape(var, [-1]) for var in var_or_var_list
- ]
+ if var_list:
+ all_weight_vars += [array_ops.reshape(var, [-1]) for var in var_list]
return nn.zero_fraction(array_ops.concat(all_weight_vars, axis=0))
@@ -92,14 +115,36 @@
Returns:
A `Tensor` representing the logits.
"""
- cols_to_vars = {}
- logits = feature_column_lib.linear_model(
- features=features,
- feature_columns=feature_columns,
- units=units,
- sparse_combiner=sparse_combiner,
- cols_to_vars=cols_to_vars)
- bias = cols_to_vars.pop('bias')
+ if feature_column_v2.is_feature_column_v2(feature_columns):
+ shared_state_manager = feature_column_v2.SharedEmbeddingStateManager()
+ linear_model = feature_column_v2.LinearModel(
+ feature_columns=feature_columns,
+ units=units,
+ sparse_combiner=sparse_combiner,
+ shared_state_manager=shared_state_manager)
+ logits = linear_model(features)
+ bias = linear_model.bias_variable
+
+ # We'd like to get all the non-bias variables associated with this
+ # LinearModel. This includes the shared embedding variables as well.
+ variables = linear_model.variables
+ variables.remove(bias)
+ variables.extend(shared_state_manager.variables)
+
+ # Expand (potential) Partitioned variables
+ bias = _get_expanded_variable_list([bias])
+ variables = _get_expanded_variable_list(variables)
+ else:
+ linear_model = feature_column._LinearModel( # pylint: disable=protected-access
+ feature_columns=feature_columns,
+ units=units,
+ sparse_combiner=sparse_combiner,
+ name='linear_model')
+ logits = linear_model(features)
+ cols_to_vars = linear_model.cols_to_vars()
+ bias = cols_to_vars.pop('bias')
+ variables = cols_to_vars.values()
+
if units > 1:
summary.histogram('bias', bias)
else:
@@ -107,7 +152,7 @@
# so we should provide a scalar summary.
summary.scalar('bias', bias[0][0])
summary.scalar('fraction_of_zero_weights',
- _compute_fraction_of_zero(cols_to_vars))
+ _compute_fraction_of_zero(variables))
return logits
return linear_logit_fn
diff --git a/tensorflow/python/estimator/canned/linear_test.py b/tensorflow/python/estimator/canned/linear_test.py
index 59a2304..3e6da5d 100644
--- a/tensorflow/python/estimator/canned/linear_test.py
+++ b/tensorflow/python/estimator/canned/linear_test.py
@@ -20,6 +20,8 @@
from tensorflow.python.estimator.canned import linear
from tensorflow.python.estimator.canned import linear_testing_utils
+from tensorflow.python.feature_column import feature_column
+from tensorflow.python.feature_column import feature_column_v2
from tensorflow.python.platform import test
@@ -40,7 +42,16 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
linear_testing_utils.BaseLinearRegressorPartitionerTest.__init__(
- self, _linear_regressor_fn)
+ self, _linear_regressor_fn, fc_lib=feature_column)
+
+
+class LinearRegressorPartitionerV2Test(
+ linear_testing_utils.BaseLinearRegressorPartitionerTest, test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ linear_testing_utils.BaseLinearRegressorPartitionerTest.__init__(
+ self, _linear_regressor_fn, fc_lib=feature_column_v2)
class LinearRegressorEvaluationTest(
@@ -49,7 +60,16 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
linear_testing_utils.BaseLinearRegressorEvaluationTest.__init__(
- self, _linear_regressor_fn)
+ self, _linear_regressor_fn, fc_lib=feature_column)
+
+
+class LinearRegressorEvaluationV2Test(
+ linear_testing_utils.BaseLinearRegressorEvaluationTest, test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ linear_testing_utils.BaseLinearRegressorEvaluationTest.__init__(
+ self, _linear_regressor_fn, fc_lib=feature_column_v2)
class LinearRegressorPredictTest(
@@ -58,7 +78,16 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
linear_testing_utils.BaseLinearRegressorPredictTest.__init__(
- self, _linear_regressor_fn)
+ self, _linear_regressor_fn, fc_lib=feature_column)
+
+
+class LinearRegressorPredictV2Test(
+ linear_testing_utils.BaseLinearRegressorPredictTest, test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ linear_testing_utils.BaseLinearRegressorPredictTest.__init__(
+ self, _linear_regressor_fn, fc_lib=feature_column_v2)
class LinearRegressorIntegrationTest(
@@ -67,7 +96,16 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
linear_testing_utils.BaseLinearRegressorIntegrationTest.__init__(
- self, _linear_regressor_fn)
+ self, _linear_regressor_fn, fc_lib=feature_column)
+
+
+class LinearRegressorIntegrationV2Test(
+ linear_testing_utils.BaseLinearRegressorIntegrationTest, test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ linear_testing_utils.BaseLinearRegressorIntegrationTest.__init__(
+ self, _linear_regressor_fn, fc_lib=feature_column_v2)
class LinearRegressorTrainingTest(
@@ -76,19 +114,37 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
linear_testing_utils.BaseLinearRegressorTrainingTest.__init__(
- self, _linear_regressor_fn)
+ self, _linear_regressor_fn, fc_lib=feature_column)
+
+
+class LinearRegressorTrainingV2Test(
+ linear_testing_utils.BaseLinearRegressorTrainingTest, test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ linear_testing_utils.BaseLinearRegressorTrainingTest.__init__(
+ self, _linear_regressor_fn, fc_lib=feature_column_v2)
# Tests for Linear Classifier.
-
-
class LinearClassifierTrainingTest(
linear_testing_utils.BaseLinearClassifierTrainingTest, test.TestCase):
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
linear_testing_utils.BaseLinearClassifierTrainingTest.__init__(
- self, linear_classifier_fn=_linear_classifier_fn)
+ self, linear_classifier_fn=_linear_classifier_fn, fc_lib=feature_column)
+
+
+class LinearClassifierTrainingV2Test(
+ linear_testing_utils.BaseLinearClassifierTrainingTest, test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ linear_testing_utils.BaseLinearClassifierTrainingTest.__init__(
+ self,
+ linear_classifier_fn=_linear_classifier_fn,
+ fc_lib=feature_column_v2)
class LinearClassifierEvaluationTest(
@@ -97,7 +153,18 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
linear_testing_utils.BaseLinearClassifierEvaluationTest.__init__(
- self, linear_classifier_fn=_linear_classifier_fn)
+ self, linear_classifier_fn=_linear_classifier_fn, fc_lib=feature_column)
+
+
+class LinearClassifierEvaluationV2Test(
+ linear_testing_utils.BaseLinearClassifierEvaluationTest, test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ linear_testing_utils.BaseLinearClassifierEvaluationTest.__init__(
+ self,
+ linear_classifier_fn=_linear_classifier_fn,
+ fc_lib=feature_column_v2)
class LinearClassifierPredictTest(
@@ -106,7 +173,18 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
linear_testing_utils.BaseLinearClassifierPredictTest.__init__(
- self, linear_classifier_fn=_linear_classifier_fn)
+ self, linear_classifier_fn=_linear_classifier_fn, fc_lib=feature_column)
+
+
+class LinearClassifierPredictV2Test(
+ linear_testing_utils.BaseLinearClassifierPredictTest, test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ linear_testing_utils.BaseLinearClassifierPredictTest.__init__(
+ self,
+ linear_classifier_fn=_linear_classifier_fn,
+ fc_lib=feature_column_v2)
class LinearClassifierIntegrationTest(
@@ -115,7 +193,18 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
linear_testing_utils.BaseLinearClassifierIntegrationTest.__init__(
- self, linear_classifier_fn=_linear_classifier_fn)
+ self, linear_classifier_fn=_linear_classifier_fn, fc_lib=feature_column)
+
+
+class LinearClassifierIntegrationV2Test(
+ linear_testing_utils.BaseLinearClassifierIntegrationTest, test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ linear_testing_utils.BaseLinearClassifierIntegrationTest.__init__(
+ self,
+ linear_classifier_fn=_linear_classifier_fn,
+ fc_lib=feature_column_v2)
# Tests for Linear logit_fn.
@@ -124,7 +213,17 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
- linear_testing_utils.BaseLinearLogitFnTest.__init__(self)
+ linear_testing_utils.BaseLinearLogitFnTest.__init__(
+ self, fc_lib=feature_column)
+
+
+class LinearLogitFnV2Test(linear_testing_utils.BaseLinearLogitFnTest,
+ test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ linear_testing_utils.BaseLinearLogitFnTest.__init__(
+ self, fc_lib=feature_column_v2)
# Tests for warm-starting with Linear logit_fn.
@@ -134,7 +233,22 @@
def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
test.TestCase.__init__(self, methodName)
linear_testing_utils.BaseLinearWarmStartingTest.__init__(
- self, _linear_classifier_fn, _linear_regressor_fn)
+ self,
+ _linear_classifier_fn,
+ _linear_regressor_fn,
+ fc_lib=feature_column)
+
+
+class LinearWarmStartingV2Test(linear_testing_utils.BaseLinearWarmStartingTest,
+ test.TestCase):
+
+ def __init__(self, methodName='runTest'): # pylint: disable=invalid-name
+ test.TestCase.__init__(self, methodName)
+ linear_testing_utils.BaseLinearWarmStartingTest.__init__(
+ self,
+ _linear_classifier_fn,
+ _linear_regressor_fn,
+ fc_lib=feature_column_v2)
if __name__ == '__main__':
diff --git a/tensorflow/python/estimator/canned/linear_testing_utils.py b/tensorflow/python/estimator/canned/linear_testing_utils.py
index 65cdd50..827352a 100644
--- a/tensorflow/python/estimator/canned/linear_testing_utils.py
+++ b/tensorflow/python/estimator/canned/linear_testing_utils.py
@@ -37,7 +37,8 @@
from tensorflow.python.estimator.export import export
from tensorflow.python.estimator.inputs import numpy_io
from tensorflow.python.estimator.inputs import pandas_io
-from tensorflow.python.feature_column import feature_column as feature_column_lib
+from tensorflow.python.feature_column import feature_column
+from tensorflow.python.feature_column import feature_column_v2
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import sparse_tensor
@@ -152,8 +153,9 @@
class BaseLinearRegressorPartitionerTest(object):
- def __init__(self, linear_regressor_fn):
+ def __init__(self, linear_regressor_fn, fc_lib=feature_column):
self._linear_regressor_fn = linear_regressor_fn
+ self._fc_lib = fc_lib
def setUp(self):
self._model_dir = tempfile.mkdtemp()
@@ -173,7 +175,7 @@
return [partitions, 1] if shape[0] == x_dim else [1]
regressor = self._linear_regressor_fn(
- feature_columns=(feature_column_lib.categorical_column_with_hash_bucket(
+ feature_columns=(self._fc_lib.categorical_column_with_hash_bucket(
'language', hash_bucket_size=x_dim),),
partitioner=_partitioner,
model_dir=self._model_dir)
@@ -209,9 +211,8 @@
'_get_replica_device_setter',
return_value=lambda _: '/cpu:0'):
linear_regressor = self._linear_regressor_fn(
- feature_columns=(
- feature_column_lib.categorical_column_with_hash_bucket(
- 'language', hash_bucket_size=x_dim),),
+ feature_columns=(self._fc_lib.categorical_column_with_hash_bucket(
+ 'language', hash_bucket_size=x_dim),),
config=FakeRunConfig(),
model_dir=self._model_dir)
@@ -232,8 +233,9 @@
# TODO(b/36813849): Add tests with dynamic shape inputs using placeholders.
class BaseLinearRegressorEvaluationTest(object):
- def __init__(self, linear_regressor_fn):
+ def __init__(self, linear_regressor_fn, fc_lib=feature_column):
self._linear_regressor_fn = linear_regressor_fn
+ self._fc_lib = fc_lib
def setUp(self):
self._model_dir = tempfile.mkdtemp()
@@ -252,7 +254,7 @@
save_variables_to_ckpt(self._model_dir)
linear_regressor = self._linear_regressor_fn(
- feature_columns=(feature_column_lib.numeric_column('age'),),
+ feature_columns=(self._fc_lib.numeric_column('age'),),
model_dir=self._model_dir)
eval_metrics = linear_regressor.evaluate(
input_fn=lambda: ({'age': ((1,),)}, ((10.,),)), steps=1)
@@ -276,7 +278,7 @@
save_variables_to_ckpt(self._model_dir)
linear_regressor = self._linear_regressor_fn(
- feature_columns=(feature_column_lib.numeric_column('age'),),
+ feature_columns=(self._fc_lib.numeric_column('age'),),
model_dir=self._model_dir)
eval_metrics = linear_regressor.evaluate(
input_fn=lambda: ({'age': ((1,), (1,))}, ((10.,), (10.,))), steps=1)
@@ -308,7 +310,7 @@
return features, labels
linear_regressor = self._linear_regressor_fn(
- feature_columns=(feature_column_lib.numeric_column('age'),),
+ feature_columns=(self._fc_lib.numeric_column('age'),),
weight_column='weights',
model_dir=self._model_dir)
eval_metrics = linear_regressor.evaluate(input_fn=_input_fn, steps=1)
@@ -336,8 +338,7 @@
save_variables_to_ckpt(self._model_dir)
linear_regressor = self._linear_regressor_fn(
- feature_columns=(feature_column_lib.numeric_column(
- 'age', shape=(x_dim,)),),
+ feature_columns=(self._fc_lib.numeric_column('age', shape=(x_dim,)),),
label_dimension=label_dim,
model_dir=self._model_dir)
input_fn = numpy_io.numpy_input_fn(
@@ -374,8 +375,8 @@
batch_size = 2
feature_columns = [
- feature_column_lib.numeric_column('age'),
- feature_column_lib.numeric_column('height')
+ self._fc_lib.numeric_column('age'),
+ self._fc_lib.numeric_column('height')
]
input_fn = numpy_io.numpy_input_fn(
x={'age': np.array([20, 40]),
@@ -402,8 +403,9 @@
class BaseLinearRegressorPredictTest(object):
- def __init__(self, linear_regressor_fn):
+ def __init__(self, linear_regressor_fn, fc_lib=feature_column):
self._linear_regressor_fn = linear_regressor_fn
+ self._fc_lib = fc_lib
def setUp(self):
self._model_dir = tempfile.mkdtemp()
@@ -422,7 +424,7 @@
save_variables_to_ckpt(self._model_dir)
linear_regressor = self._linear_regressor_fn(
- feature_columns=(feature_column_lib.numeric_column('x'),),
+ feature_columns=(self._fc_lib.numeric_column('x'),),
model_dir=self._model_dir)
predict_input_fn = numpy_io.numpy_input_fn(
@@ -441,7 +443,7 @@
batch_size = 2
label_dimension = 3
x_dim = 4
- feature_columns = (feature_column_lib.numeric_column('x', shape=(x_dim,)),)
+ feature_columns = (self._fc_lib.numeric_column('x', shape=(x_dim,)),)
with ops.Graph().as_default():
variables_lib.Variable( # shape=[x_dim, label_dimension]
[[1., 2., 3.], [2., 3., 4.], [3., 4., 5.], [4., 5., 6.]],
@@ -479,8 +481,8 @@
save_variables_to_ckpt(self._model_dir)
linear_regressor = self._linear_regressor_fn(
- feature_columns=(feature_column_lib.numeric_column('x0'),
- feature_column_lib.numeric_column('x1')),
+ feature_columns=(self._fc_lib.numeric_column('x0'),
+ self._fc_lib.numeric_column('x1')),
model_dir=self._model_dir)
predict_input_fn = numpy_io.numpy_input_fn(
@@ -515,9 +517,8 @@
dense_shape=[2, 2]),
})
- feature_columns = (
- feature_column_lib.categorical_column_with_vocabulary_list(
- 'language', vocabulary_list=['a', 'b', 'c']),)
+ feature_columns = (self._fc_lib.categorical_column_with_vocabulary_list(
+ 'language', vocabulary_list=['a', 'b', 'c']),)
# Check prediction for each sparse_combiner.
# With sparse_combiner = 'sum', we have
@@ -561,8 +562,9 @@
class BaseLinearRegressorIntegrationTest(object):
- def __init__(self, linear_regressor_fn):
+ def __init__(self, linear_regressor_fn, fc_lib=feature_column):
self._linear_regressor_fn = linear_regressor_fn
+ self._fc_lib = fc_lib
def setUp(self):
self._model_dir = tempfile.mkdtemp()
@@ -575,7 +577,7 @@
def _test_complete_flow(self, train_input_fn, eval_input_fn, predict_input_fn,
input_dimension, label_dimension, prediction_length):
feature_columns = [
- feature_column_lib.numeric_column('x', shape=(input_dimension,))
+ self._fc_lib.numeric_column('x', shape=(input_dimension,))
]
est = self._linear_regressor_fn(
feature_columns=feature_columns,
@@ -597,7 +599,7 @@
self.assertAllEqual((prediction_length, label_dimension), predictions.shape)
# EXPORT
- feature_spec = feature_column_lib.make_parse_example_spec(feature_columns)
+ feature_spec = self._fc_lib.make_parse_example_spec(feature_columns)
serving_input_receiver_fn = export.build_parsing_serving_input_receiver_fn(
feature_spec)
export_dir = est.export_savedmodel(tempfile.mkdtemp(),
@@ -729,8 +731,9 @@
class BaseLinearRegressorTrainingTest(object):
- def __init__(self, linear_regressor_fn):
+ def __init__(self, linear_regressor_fn, fc_lib=feature_column):
self._linear_regressor_fn = linear_regressor_fn
+ self._fc_lib = fc_lib
def setUp(self):
self._model_dir = tempfile.mkdtemp()
@@ -808,7 +811,7 @@
label = 5.
age = 17
linear_regressor = self._linear_regressor_fn(
- feature_columns=(feature_column_lib.numeric_column('age'),),
+ feature_columns=(self._fc_lib.numeric_column('age'),),
model_dir=self._model_dir)
# Train for a few steps, and validate final checkpoint.
@@ -820,7 +823,7 @@
def testTrainWithOneDimLabel(self):
label_dimension = 1
batch_size = 20
- feature_columns = [feature_column_lib.numeric_column('age', shape=(1,))]
+ feature_columns = [self._fc_lib.numeric_column('age', shape=(1,))]
est = self._linear_regressor_fn(
feature_columns=feature_columns,
label_dimension=label_dimension,
@@ -840,7 +843,7 @@
def testTrainWithOneDimWeight(self):
label_dimension = 1
batch_size = 20
- feature_columns = [feature_column_lib.numeric_column('age', shape=(1,))]
+ feature_columns = [self._fc_lib.numeric_column('age', shape=(1,))]
est = self._linear_regressor_fn(
feature_columns=feature_columns,
label_dimension=label_dimension,
@@ -867,7 +870,7 @@
# loss = (logits - label)^2 = (0 - 5.)^2 = 25.
mock_optimizer = self._mock_optimizer(expected_loss=25.)
linear_regressor = self._linear_regressor_fn(
- feature_columns=(feature_column_lib.numeric_column('age'),),
+ feature_columns=(self._fc_lib.numeric_column('age'),),
model_dir=self._model_dir,
optimizer=mock_optimizer)
self.assertEqual(0, mock_optimizer.minimize.call_count)
@@ -900,7 +903,7 @@
# loss = (logits - label)^2 = (175 - 5)^2 = 28900
mock_optimizer = self._mock_optimizer(expected_loss=28900.)
linear_regressor = self._linear_regressor_fn(
- feature_columns=(feature_column_lib.numeric_column('age'),),
+ feature_columns=(self._fc_lib.numeric_column('age'),),
model_dir=self._model_dir,
optimizer=mock_optimizer)
self.assertEqual(0, mock_optimizer.minimize.call_count)
@@ -935,7 +938,7 @@
# loss = sum(logits - label)^2 = (175 - 5)^2 + (155 - 3)^2 = 52004
mock_optimizer = self._mock_optimizer(expected_loss=52004.)
linear_regressor = self._linear_regressor_fn(
- feature_columns=(feature_column_lib.numeric_column('age'),),
+ feature_columns=(self._fc_lib.numeric_column('age'),),
model_dir=self._model_dir,
optimizer=mock_optimizer)
self.assertEqual(0, mock_optimizer.minimize.call_count)
@@ -954,8 +957,9 @@
class BaseLinearClassifierTrainingTest(object):
- def __init__(self, linear_classifier_fn):
+ def __init__(self, linear_classifier_fn, fc_lib=feature_column):
self._linear_classifier_fn = linear_classifier_fn
+ self._fc_lib = fc_lib
def setUp(self):
self._model_dir = tempfile.mkdtemp()
@@ -1031,7 +1035,7 @@
label = 0
age = 17
est = linear.LinearClassifier(
- feature_columns=(feature_column_lib.numeric_column('age'),),
+ feature_columns=(self._fc_lib.numeric_column('age'),),
n_classes=n_classes,
model_dir=self._model_dir)
@@ -1051,7 +1055,7 @@
batch_size = 20
est = linear.LinearClassifier(
- feature_columns=(feature_column_lib.numeric_column('age'),),
+ feature_columns=(self._fc_lib.numeric_column('age'),),
n_classes=n_classes,
model_dir=self._model_dir)
data_rank_1 = np.array([0, 1])
@@ -1078,7 +1082,7 @@
batch_size = 20
est = linear.LinearClassifier(
- feature_columns=(feature_column_lib.numeric_column('age'),),
+ feature_columns=(self._fc_lib.numeric_column('age'),),
n_classes=n_classes,
model_dir=self._model_dir)
data_rank_1 = np.array([0, 1])
@@ -1103,7 +1107,7 @@
batch_size = 20
est = linear.LinearClassifier(
- feature_columns=(feature_column_lib.numeric_column('age'),),
+ feature_columns=(self._fc_lib.numeric_column('age'),),
weight_column='w',
n_classes=n_classes,
model_dir=self._model_dir)
@@ -1129,7 +1133,7 @@
batch_size = 20
est = linear.LinearClassifier(
- feature_columns=(feature_column_lib.numeric_column('age'),),
+ feature_columns=(self._fc_lib.numeric_column('age'),),
weight_column='w',
n_classes=n_classes,
model_dir=self._model_dir)
@@ -1166,7 +1170,7 @@
expected_loss=-1 * math.log(1.0/n_classes))
est = linear.LinearClassifier(
- feature_columns=(feature_column_lib.numeric_column('age'),),
+ feature_columns=(self._fc_lib.numeric_column('age'),),
n_classes=n_classes,
optimizer=mock_optimizer,
model_dir=self._model_dir)
@@ -1229,7 +1233,7 @@
mock_optimizer = self._mock_optimizer(expected_loss=expected_loss)
est = linear.LinearClassifier(
- feature_columns=(feature_column_lib.numeric_column('age'),),
+ feature_columns=(self._fc_lib.numeric_column('age'),),
n_classes=n_classes,
optimizer=mock_optimizer,
model_dir=self._model_dir)
@@ -1277,7 +1281,7 @@
mock_optimizer = self._mock_optimizer(expected_loss=1.1132617)
est = linear.LinearClassifier(
- feature_columns=(feature_column_lib.numeric_column('age'),),
+ feature_columns=(self._fc_lib.numeric_column('age'),),
n_classes=n_classes,
optimizer=mock_optimizer,
model_dir=self._model_dir)
@@ -1341,7 +1345,7 @@
mock_optimizer = self._mock_optimizer(expected_loss=expected_loss)
est = linear.LinearClassifier(
- feature_columns=(feature_column_lib.numeric_column('age'),),
+ feature_columns=(self._fc_lib.numeric_column('age'),),
n_classes=n_classes,
optimizer=mock_optimizer,
model_dir=self._model_dir)
@@ -1368,8 +1372,9 @@
class BaseLinearClassifierEvaluationTest(object):
- def __init__(self, linear_classifier_fn):
+ def __init__(self, linear_classifier_fn, fc_lib=feature_column):
self._linear_classifier_fn = linear_classifier_fn
+ self._fc_lib = fc_lib
def setUp(self):
self._model_dir = tempfile.mkdtemp()
@@ -1398,7 +1403,7 @@
save_variables_to_ckpt(self._model_dir)
est = self._linear_classifier_fn(
- feature_columns=(feature_column_lib.numeric_column('age'),),
+ feature_columns=(self._fc_lib.numeric_column('age'),),
n_classes=n_classes,
model_dir=self._model_dir)
eval_metrics = est.evaluate(
@@ -1464,7 +1469,7 @@
save_variables_to_ckpt(self._model_dir)
est = self._linear_classifier_fn(
- feature_columns=(feature_column_lib.numeric_column('age'),),
+ feature_columns=(self._fc_lib.numeric_column('age'),),
n_classes=n_classes,
model_dir=self._model_dir)
eval_metrics = est.evaluate(
@@ -1540,7 +1545,7 @@
save_variables_to_ckpt(self._model_dir)
est = self._linear_classifier_fn(
- feature_columns=(feature_column_lib.numeric_column('age'),),
+ feature_columns=(self._fc_lib.numeric_column('age'),),
n_classes=n_classes,
weight_column='w',
model_dir=self._model_dir)
@@ -1605,8 +1610,9 @@
class BaseLinearClassifierPredictTest(object):
- def __init__(self, linear_classifier_fn):
+ def __init__(self, linear_classifier_fn, fc_lib=feature_column):
self._linear_classifier_fn = linear_classifier_fn
+ self._fc_lib = fc_lib
def setUp(self):
self._model_dir = tempfile.mkdtemp()
@@ -1634,7 +1640,7 @@
save_variables_to_ckpt(self._model_dir)
est = self._linear_classifier_fn(
- feature_columns=(feature_column_lib.numeric_column('age'),),
+ feature_columns=(self._fc_lib.numeric_column('age'),),
label_vocabulary=label_vocabulary,
n_classes=n_classes,
model_dir=self._model_dir)
@@ -1730,9 +1736,8 @@
dense_shape=[2, 2]),
})
- feature_columns = (
- feature_column_lib.categorical_column_with_vocabulary_list(
- 'language', vocabulary_list=['a', 'b', 'c']),)
+ feature_columns = (self._fc_lib.categorical_column_with_vocabulary_list(
+ 'language', vocabulary_list=['a', 'b', 'c']),)
# Check prediction for each sparse_combiner.
# With sparse_combiner = 'sum', we have
@@ -1776,8 +1781,9 @@
class BaseLinearClassifierIntegrationTest(object):
- def __init__(self, linear_classifier_fn):
+ def __init__(self, linear_classifier_fn, fc_lib=feature_column):
self._linear_classifier_fn = linear_classifier_fn
+ self._fc_lib = fc_lib
def setUp(self):
self._model_dir = tempfile.mkdtemp()
@@ -1789,7 +1795,7 @@
def _test_complete_flow(self, n_classes, train_input_fn, eval_input_fn,
predict_input_fn, input_dimension, prediction_length):
feature_columns = [
- feature_column_lib.numeric_column('x', shape=(input_dimension,))
+ self._fc_lib.numeric_column('x', shape=(input_dimension,))
]
est = self._linear_classifier_fn(
feature_columns=feature_columns,
@@ -1811,7 +1817,7 @@
self.assertAllEqual((prediction_length, 1), predictions.shape)
# EXPORT
- feature_spec = feature_column_lib.make_parse_example_spec(feature_columns)
+ feature_spec = self._fc_lib.make_parse_example_spec(feature_columns)
serving_input_receiver_fn = export.build_parsing_serving_input_receiver_fn(
feature_spec)
export_dir = est.export_savedmodel(tempfile.mkdtemp(),
@@ -1961,9 +1967,12 @@
class BaseLinearLogitFnTest(object):
+ def __init__(self, fc_lib=feature_column):
+ self._fc_lib = fc_lib
+
def test_basic_logit_correctness(self):
"""linear_logit_fn simply wraps feature_column_lib.linear_model."""
- age = feature_column_lib.numeric_column('age')
+ age = self._fc_lib.numeric_column('age')
with ops.Graph().as_default():
logit_fn = linear._linear_logit_fn_builder(units=2, feature_columns=[age])
logits = logit_fn(features={'age': [[23.], [31.]]})
@@ -1983,12 +1992,14 @@
def test_compute_fraction_of_zero(self):
"""Tests the calculation of sparsity."""
- age = feature_column_lib.numeric_column('age')
- occupation = feature_column_lib.categorical_column_with_hash_bucket(
+ if self._fc_lib != feature_column:
+ return
+ age = feature_column.numeric_column('age')
+ occupation = feature_column.categorical_column_with_hash_bucket(
'occupation', hash_bucket_size=5)
with ops.Graph().as_default():
cols_to_vars = {}
- feature_column_lib.linear_model(
+ feature_column.linear_model(
features={
'age': [[23.], [31.]],
'occupation': [['doctor'], ['engineer']]
@@ -1997,7 +2008,42 @@
units=3,
cols_to_vars=cols_to_vars)
cols_to_vars.pop('bias')
- fraction_zero = linear._compute_fraction_of_zero(cols_to_vars)
+ fraction_zero = linear._compute_fraction_of_zero(cols_to_vars.values())
+ age_var = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES,
+ 'linear_model/age')[0]
+ with tf_session.Session() as sess:
+ sess.run([variables_lib.global_variables_initializer()])
+ # Upon initialization, all variables will be zero.
+ self.assertAllClose(1, fraction_zero.eval())
+
+ sess.run(age_var.assign([[2.0, 0.0, -1.0]]))
+ # 1 of the 3 age weights are zero, and all of the 15 (5 hash buckets
+ # x 3-dim output) are zero.
+ self.assertAllClose(16. / 18., fraction_zero.eval())
+
+ def test_compute_fraction_of_zero_v2(self):
+ """Tests the calculation of sparsity."""
+ if self._fc_lib != feature_column_v2:
+ return
+
+ age = feature_column_v2.numeric_column('age')
+ occupation = feature_column_v2.categorical_column_with_hash_bucket(
+ 'occupation', hash_bucket_size=5)
+ shared_state_manager = feature_column_v2.SharedEmbeddingStateManager()
+ with ops.Graph().as_default():
+ model = feature_column_v2.LinearModel(
+ feature_columns=[age, occupation],
+ units=3,
+ shared_state_manager=shared_state_manager)
+ features = {
+ 'age': [[23.], [31.]],
+ 'occupation': [['doctor'], ['engineer']]
+ }
+ model(features)
+ variables = model.variables
+ variables.remove(model.bias_variable)
+ variables.extend(shared_state_manager.variables)
+ fraction_zero = linear._compute_fraction_of_zero(variables)
age_var = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES,
'linear_model/age')[0]
with tf_session.Session() as sess:
@@ -2013,9 +2059,13 @@
class BaseLinearWarmStartingTest(object):
- def __init__(self, _linear_classifier_fn, _linear_regressor_fn):
+ def __init__(self,
+ _linear_classifier_fn,
+ _linear_regressor_fn,
+ fc_lib=feature_column):
self._linear_classifier_fn = _linear_classifier_fn
self._linear_regressor_fn = _linear_regressor_fn
+ self._fc_lib = fc_lib
def setUp(self):
# Create a directory to save our old checkpoint and vocabularies to.
@@ -2039,7 +2089,7 @@
def test_classifier_basic_warm_starting(self):
"""Tests correctness of LinearClassifier default warm-start."""
- age = feature_column_lib.numeric_column('age')
+ age = self._fc_lib.numeric_column('age')
# Create a LinearClassifier and train to save a checkpoint.
linear_classifier = self._linear_classifier_fn(
@@ -2066,7 +2116,7 @@
def test_regressor_basic_warm_starting(self):
"""Tests correctness of LinearRegressor default warm-start."""
- age = feature_column_lib.numeric_column('age')
+ age = self._fc_lib.numeric_column('age')
# Create a LinearRegressor and train to save a checkpoint.
linear_regressor = self._linear_regressor_fn(
@@ -2091,7 +2141,7 @@
def test_warm_starting_selective_variables(self):
"""Tests selecting variables to warm-start."""
- age = feature_column_lib.numeric_column('age')
+ age = self._fc_lib.numeric_column('age')
# Create a LinearClassifier and train to save a checkpoint.
linear_classifier = self._linear_classifier_fn(
@@ -2128,7 +2178,7 @@
vocab_file = os.path.join(self._ckpt_and_vocab_dir, 'occupation_vocab')
with open(vocab_file, 'w') as f:
f.write('\n'.join(vocab_list))
- occupation = feature_column_lib.categorical_column_with_vocabulary_file(
+ occupation = self._fc_lib.categorical_column_with_vocabulary_file(
'occupation',
vocabulary_file=vocab_file,
vocabulary_size=len(vocab_list))
@@ -2152,7 +2202,7 @@
'new_occupation_vocab')
with open(new_vocab_file, 'w') as f:
f.write('\n'.join(new_vocab_list))
- new_occupation = feature_column_lib.categorical_column_with_vocabulary_file(
+ new_occupation = self._fc_lib.categorical_column_with_vocabulary_file(
'occupation',
vocabulary_file=new_vocab_file,
vocabulary_size=len(new_vocab_list))
@@ -2205,7 +2255,7 @@
def test_warm_starting_with_naming_change(self):
"""Tests warm-starting with a Tensor name remapping."""
- age_in_years = feature_column_lib.numeric_column('age_in_years')
+ age_in_years = self._fc_lib.numeric_column('age_in_years')
# Create a LinearClassifier and train to save a checkpoint.
linear_classifier = self._linear_classifier_fn(
@@ -2219,7 +2269,7 @@
# learning_rate = 0.0 optimizer to check values (use SGD so we don't have
# accumulator values that change).
warm_started_linear_classifier = self._linear_classifier_fn(
- feature_columns=[feature_column_lib.numeric_column('age')],
+ feature_columns=[self._fc_lib.numeric_column('age')],
n_classes=4,
optimizer=gradient_descent.GradientDescentOptimizer(learning_rate=0.0),
# The 'age' variable correspond to the 'age_in_years' variable in the
diff --git a/tensorflow/python/estimator/estimator.py b/tensorflow/python/estimator/estimator.py
index eec64ad..e6d82f0 100644
--- a/tensorflow/python/estimator/estimator.py
+++ b/tensorflow/python/estimator/estimator.py
@@ -144,7 +144,7 @@
* `labels`: This is the second item returned from the `input_fn`
passed to `train`, `evaluate`, and `predict`. This should be a
single `tf.Tensor` or `dict` of same (for multi-head models).
- If mode is @{tf.estimator.ModeKeys.PREDICT}, `labels=None` will
+ If mode is `tf.estimator.ModeKeys.PREDICT`, `labels=None` will
be passed. If the `model_fn`'s signature does not accept
`mode`, the `model_fn` must still be able to handle
`labels=None`.
@@ -468,17 +468,41 @@
with ops.Graph().as_default():
if self._eval_distribution:
+ # We want to create the iterations variable outside the distribution
+ # scope as that is just stored on the host and mainly used to drive
+ # the loop and doesn't need to be a Mirrored/Device variable.
+ training.get_or_create_steps_per_run_variable()
with self._eval_distribution.scope():
return _evaluate()
else:
return _evaluate()
def _convert_eval_steps_to_hooks(self, steps):
+ """Create hooks to run correct number of steps in evaluation.
+
+ Args:
+ steps: number of steps to run during evaluation.
+
+ Raises:
+ ValueError: if steps is less than or equal to zero.
+
+ Returns:
+ List of hooks to be passed to the estimator.
+ """
if steps is None:
return []
if steps <= 0:
raise ValueError('Must specify steps > 0, given: {}'.format(steps))
+
+ # The hooks are declared as private in evaluation.py discourage the use
+ # by other libraries or open source users. This should be the only usage
+ # of the estimator evaluation hooks.
+ if self._eval_distribution:
+ steps_per_run = getattr(self._eval_distribution, 'steps_per_run', 1)
+ if steps_per_run > 1:
+ return [evaluation._MultiStepStopAfterNEvalsHook( # pylint: disable=protected-access
+ num_evals=steps, steps_per_run=steps_per_run)]
return [evaluation._StopAfterNEvalsHook(num_evals=steps)] # pylint: disable=protected-access
def predict(self,
@@ -783,9 +807,9 @@
those features and labels, and restores the given checkpoint
(or, lacking that, the most recent checkpoint) into the graph.
Only one of the modes is used for saving variables to the `SavedModel`
- (order of preference: @{tf.estimator.ModeKeys#TRAIN$TRAIN},
- @{tf.estimator.ModeKeys#EVAL$EVAL}, then
- @{tf.estimator.ModeKeys#PREDICT$PREDICT}), such that up to three
+ (order of preference: `tf.estimator.ModeKeys.TRAIN`,
+ `tf.estimator.ModeKeys.EVAL`, then
+ `tf.estimator.ModeKeys.PREDICT`), such that up to three
`tf.MetaGraphDefs` are saved with a single set of variables in a single
`SavedModel` directory.
@@ -1081,7 +1105,7 @@
"""Creates the global step tensor in graph.
The global step tensor must be an integer type with name 'global_step' and
- be added to the collection @{tf.GraphKeys#GLOBAL_STEP$GLOBAL_STEP}.
+ be added to the collection `tf.GraphKeys.GLOBAL_STEP`.
Args:
graph: The graph in which to create the global step tensor.
@@ -1394,6 +1418,36 @@
# It is expected to have one CheckpointSaverHook. If multiple, we pick
# up the first one to add listener.
saver_hooks[0]._listeners.extend(saving_listeners) # pylint: disable=protected-access
+
+ # Add summary hooks to worker 0 if we are running with a master, to ensure
+ # that summaries are written at correct intervals even with long-running
+ # evaluations.
+ save_summary_steps = self._config.save_summary_steps
+ log_step_count_steps = self._config.log_step_count_steps
+ if (self._config.cluster_spec and self._config.cluster_spec.jobs and
+ (run_config.TaskType.MASTER in self._config.cluster_spec.jobs)):
+ # Update config values to prevent the default hooks from being created on
+ # the master or other workers.
+ save_summary_steps = 0
+ log_step_count_steps = None
+
+ if (self._config.task_type == run_config.TaskType.WORKER and
+ self._config.task_id == 0):
+ if (self._config.save_summary_steps and
+ self._config.save_summary_steps > 0):
+ worker_hooks.append(
+ training.SummarySaverHook(
+ save_steps=self._config.save_summary_steps,
+ output_dir=self._config.model_dir,
+ scaffold=estimator_spec.scaffold))
+
+ if (self._config.log_step_count_steps and
+ self._config.log_step_count_steps > 0):
+ worker_hooks.append(
+ training.StepCounterHook(
+ every_n_steps=self._config.log_step_count_steps,
+ output_dir=self._config.model_dir))
+
with training.MonitoredTrainingSession(
master=self._config.master,
is_chief=self._config.is_chief,
@@ -1403,9 +1457,9 @@
chief_only_hooks=(
tuple(chief_hooks) + tuple(estimator_spec.training_chief_hooks)),
save_checkpoint_secs=0, # Saving is handled by a hook.
- save_summaries_steps=self._config.save_summary_steps,
+ save_summaries_steps=save_summary_steps,
config=self._session_config,
- log_step_count_steps=self._config.log_step_count_steps) as mon_sess:
+ log_step_count_steps=log_step_count_steps) as mon_sess:
loss = None
while not mon_sess.should_stop():
_, loss = mon_sess.run([estimator_spec.train_op, estimator_spec.loss])
@@ -1474,6 +1528,7 @@
self._eval_distribution.__class__.__name__ == 'TPUStrategy')
if is_tpu_strategy:
+ steps_per_run_variable = training.get_or_create_steps_per_run_variable()
def step_fn(ctx, features, labels=None):
"""Runs one step of the eval computation and captures outputs."""
estimator_spec = self._eval_distribution.call_for_each_tower(
@@ -1490,7 +1545,7 @@
# TODO(priyag): Fix eval step hook to account for steps_per_run.
ctx = self._eval_distribution.run_steps_on_dataset(
- step_fn, iterator, iterations=self._eval_distribution.steps_per_run)
+ step_fn, iterator, iterations=steps_per_run_variable)
update_op = ctx.run_op
eval_dict = ctx.non_tensor_outputs['eval_dict']
grouped_estimator_spec = ctx.non_tensor_outputs['estimator_spec']
diff --git a/tensorflow/python/estimator/estimator_test.py b/tensorflow/python/estimator/estimator_test.py
index 1ed5e30..246dfb1 100644
--- a/tensorflow/python/estimator/estimator_test.py
+++ b/tensorflow/python/estimator/estimator_test.py
@@ -20,6 +20,7 @@
import functools
import glob
+import json
import os
import tempfile
@@ -969,6 +970,99 @@
with self.assertRaisesRegexp(ValueError, 'train_and_evaluate'):
est.train(dummy_input_fn, steps=1)
+ def test_master_distributed_hooks(self):
+ tf_config = json.dumps({
+ 'cluster': {
+ run_config.TaskType.PS: ['localhost:1234'],
+ run_config.TaskType.WORKER: ['localhost:1235'],
+ run_config.TaskType.MASTER: ['localhost:1236']
+ },
+ 'task': {
+ 'type': run_config.TaskType.MASTER,
+ 'index': 0
+ }
+ })
+ with test.mock.patch.dict('os.environ', {'TF_CONFIG': tf_config}):
+ est = estimator.Estimator(
+ model_fn=model_fn_global_step_incrementer,
+ config=run_config.RunConfig())
+
+ with test.mock.patch.object(training,
+ 'MonitoredTrainingSession') as mock_sess:
+ est.train(dummy_input_fn, steps=1)
+ self.assertFalse(
+ any(
+ isinstance(hook, basic_session_run_hooks.SummarySaverHook)
+ for hook in mock_sess.call_args[1]['hooks']))
+ self.assertFalse(
+ any(
+ isinstance(hook, basic_session_run_hooks.StepCounterHook)
+ for hook in mock_sess.call_args[1]['hooks']))
+ self.assertEqual(0, mock_sess.call_args[1]['save_summaries_steps'])
+ self.assertIsNone(mock_sess.call_args[1]['log_step_count_steps'])
+
+ def test_master_distributed_hooks_for_worker_0(self):
+ tf_config = json.dumps({
+ 'cluster': {
+ run_config.TaskType.PS: ['localhost:1234'],
+ run_config.TaskType.WORKER: ['localhost:1235'],
+ run_config.TaskType.MASTER: ['localhost:1236']
+ },
+ 'task': {
+ 'type': run_config.TaskType.WORKER,
+ 'index': 0
+ }
+ })
+ with test.mock.patch.dict('os.environ', {'TF_CONFIG': tf_config}):
+ est = estimator.Estimator(
+ model_fn=model_fn_global_step_incrementer,
+ config=run_config.RunConfig())
+
+ with test.mock.patch.object(training,
+ 'MonitoredTrainingSession') as mock_sess:
+ est.train(dummy_input_fn, steps=1)
+ self.assertTrue(
+ any(
+ isinstance(hook, basic_session_run_hooks.SummarySaverHook)
+ for hook in mock_sess.call_args[1]['hooks']))
+ self.assertTrue(
+ any(
+ isinstance(hook, basic_session_run_hooks.StepCounterHook)
+ for hook in mock_sess.call_args[1]['hooks']))
+ self.assertEqual(0, mock_sess.call_args[1]['save_summaries_steps'])
+ self.assertIsNone(mock_sess.call_args[1]['log_step_count_steps'])
+
+ def test_master_distributed_hooks_for_worker_nonzero(self):
+ tf_config = json.dumps({
+ 'cluster': {
+ run_config.TaskType.PS: ['localhost:1234'],
+ run_config.TaskType.WORKER: ['localhost:1235', 'localhost:1237'],
+ run_config.TaskType.MASTER: ['localhost:1236']
+ },
+ 'task': {
+ 'type': run_config.TaskType.WORKER,
+ 'index': 1
+ }
+ })
+ with test.mock.patch.dict('os.environ', {'TF_CONFIG': tf_config}):
+ est = estimator.Estimator(
+ model_fn=model_fn_global_step_incrementer,
+ config=run_config.RunConfig())
+
+ with test.mock.patch.object(training,
+ 'MonitoredTrainingSession') as mock_sess:
+ est.train(dummy_input_fn, steps=1)
+ self.assertFalse(
+ any(
+ isinstance(hook, basic_session_run_hooks.SummarySaverHook)
+ for hook in mock_sess.call_args[1]['hooks']))
+ self.assertFalse(
+ any(
+ isinstance(hook, basic_session_run_hooks.StepCounterHook)
+ for hook in mock_sess.call_args[1]['hooks']))
+ self.assertEqual(0, mock_sess.call_args[1]['save_summaries_steps'])
+ self.assertIsNone(mock_sess.call_args[1]['log_step_count_steps'])
+
def _model_fn_with_eval_metric_ops(features, labels, mode, params):
_, _ = features, labels
@@ -1017,7 +1111,7 @@
def _model_fn(features, labels, mode):
_, _ = features, labels
- variables.Variable(1., name='one')
+ variables.VariableV1(1., name='one')
return model_fn_lib.EstimatorSpec(
mode=mode,
loss=constant_op.constant(0.),
@@ -1033,8 +1127,8 @@
def _model_fn(features, labels, mode):
_, _ = features, labels
- variables.Variable(1., name='one')
- variables.Variable(3., name='three')
+ variables.VariableV1(1., name='one')
+ variables.VariableV1(3., name='three')
return model_fn_lib.EstimatorSpec(
mode=mode,
loss=constant_op.constant(0.),
@@ -1178,13 +1272,13 @@
def _model_fn(features, labels, mode, params):
del features, labels, params
mean = metrics_module.Mean()
- mean.update_state(variables.Variable(2.) + 1)
+ mean.update_state(variables.VariableV1(2.) + 1)
return model_fn_lib.EstimatorSpec(
mode,
loss=constant_op.constant(1.),
eval_metric_ops={
'mean1': mean,
- 'mean2': metrics_lib.mean(variables.Variable(2.) + 1)
+ 'mean2': metrics_lib.mean(variables.VariableV1(2.) + 1)
})
est = estimator.Estimator(model_fn=_model_fn)
@@ -1332,7 +1426,7 @@
def _model_fn_with_incremental_loss(features, labels, mode):
_, _ = features, labels
- local_weight = variables.Variable(
+ local_weight = variables.VariableV1(
0., name='local_weight', collections=[ops.GraphKeys.LOCAL_VARIABLES])
# Loss will be 2, 4, 6, ...
loss = 2 * state_ops.assign_add(local_weight, 1.)
@@ -1385,7 +1479,7 @@
def _get_model_fn(val=1):
def _model_fn(features, labels, mode):
del features, labels # unused
- variables.Variable(val, name='weight')
+ variables.VariableV1(val, name='weight')
return model_fn_lib.EstimatorSpec(
mode=mode,
predictions=constant_op.constant([[1.]]),
@@ -1409,7 +1503,7 @@
def _model_fn_scaffold(features, labels, mode):
_, _ = features, labels
- variables.Variable(1., name='weight')
+ variables.VariableV1(1., name='weight')
self.mock_saver = get_mock_saver()
return model_fn_lib.EstimatorSpec(
mode=mode,
@@ -1603,7 +1697,7 @@
def test_no_checkpoint_uses_init(self):
def _model_fn(features, labels, mode, params, config):
del features, labels, params, config
- x = variables.Variable([[3.]], name='x')
+ x = variables.VariableV1([[3.]], name='x')
return model_fn_lib.EstimatorSpec(mode, predictions=math_ops.add(x, 1.))
est = estimator.Estimator(model_fn=_model_fn)
# Expected prediction value is 1 + the value of the Variable that is newly
@@ -1614,7 +1708,7 @@
def _make_model_fn(x):
def _variable_creating_and_export_model_fn(features, labels, mode):
_, _ = features, labels
- x_var = variables.Variable([[x]], name='x')
+ x_var = variables.VariableV1([[x]], name='x')
return model_fn_lib.EstimatorSpec(
mode,
predictions=math_ops.add(x_var, 1.),
@@ -1936,7 +2030,7 @@
def _model_fn(features, labels, mode):
_, _ = features, labels
- v = variables.Variable([[16.]], name='weight')
+ v = variables.VariableV1([[16.]], name='weight')
prediction = v * 2
return model_fn_lib.EstimatorSpec(
mode,
@@ -1953,7 +2047,7 @@
def _model_fn(features, labels, mode):
_, _ = features, labels
- v = variables.Variable([[16.]], name='weight')
+ v = variables.VariableV1([[16.]], name='weight')
prediction = v * 2
return model_fn_lib.EstimatorSpec(
mode,
@@ -1974,7 +2068,7 @@
def _model_fn_scaffold(features, labels, mode):
_, _ = features, labels
- variables.Variable(1., name='weight')
+ variables.VariableV1(1., name='weight')
self.mock_saver = get_mock_saver()
return model_fn_lib.EstimatorSpec(
mode=mode,
@@ -2029,7 +2123,7 @@
def _model_fn_for_export_tests(features, labels, mode):
_, _ = features, labels
- variables.Variable(1., name='weight')
+ variables.VariableV1(1., name='weight')
scores = constant_op.constant([3.])
classes = constant_op.constant(['wumpus'])
update_global_step = state_ops.assign_add(training.get_global_step(), 1)
@@ -2052,11 +2146,11 @@
def _model_fn_with_x_y(features, labels, mode):
_ = labels
- variables.Variable(1., name='weight')
+ variables.VariableV1(1., name='weight')
scores = constant_op.constant([3.])
classes = constant_op.constant(['wumpus'])
if mode == model_fn_lib.ModeKeys.PREDICT:
- variables.Variable(36., name='name_collision')
+ variables.VariableV1(36., name='name_collision')
return model_fn_lib.EstimatorSpec(
mode,
predictions=constant_op.constant(10.),
@@ -2076,8 +2170,8 @@
metrics_lib.mean(
features['x'] - features['y'], name='{}mean'.format(prefix))
}
- variables.Variable(1., name='later_var')
- variables.Variable(3., name='name_collision')
+ variables.VariableV1(1., name='later_var')
+ variables.VariableV1(3., name='name_collision')
return model_fn_lib.EstimatorSpec(
mode,
predictions=multiplied,
@@ -2411,9 +2505,9 @@
def _model_fn_with_predict_only_vars(features, labels, mode):
_, _ = features, labels
if mode == model_fn_lib.ModeKeys.PREDICT:
- variables.Variable(1., name='only_in_predict')
+ variables.VariableV1(1., name='only_in_predict')
else:
- variables.Variable(1., name='otherwise')
+ variables.VariableV1(1., name='otherwise')
prediction = constant_op.constant(1.)
return model_fn_lib.EstimatorSpec(
@@ -2684,7 +2778,7 @@
def _model_fn_scaffold(features, labels, mode):
_, _ = features, labels
- variables.Variable(1., name='weight')
+ variables.VariableV1(1., name='weight')
self.mock_saver = get_mock_saver()
scores = constant_op.constant([3.])
return model_fn_lib.EstimatorSpec(
@@ -2717,7 +2811,7 @@
def _model_fn_scaffold(features, labels, mode):
_, _ = features, labels
- variables.Variable(1., name='weight')
+ variables.VariableV1(1., name='weight')
scores = constant_op.constant([3.])
if mode == model_fn_lib.ModeKeys.PREDICT:
@@ -2762,8 +2856,8 @@
def _model_fn_scaffold(features, labels, mode):
_, _ = features, labels
- my_int = variables.Variable(1, name='my_int',
- collections=[ops.GraphKeys.LOCAL_VARIABLES])
+ my_int = variables.VariableV1(1, name='my_int',
+ collections=[ops.GraphKeys.LOCAL_VARIABLES])
_ = training.get_or_create_steps_per_run_variable()
scores = constant_op.constant([3.])
with ops.control_dependencies([
@@ -2808,8 +2902,8 @@
def _model_fn_scaffold(features, labels, mode):
_, _ = features, labels
- my_int = variables.Variable(1, name='my_int',
- collections=[ops.GraphKeys.LOCAL_VARIABLES])
+ my_int = variables.VariableV1(1, name='my_int',
+ collections=[ops.GraphKeys.LOCAL_VARIABLES])
scores = constant_op.constant([3.])
with ops.control_dependencies([
variables.local_variables_initializer(),
@@ -3038,7 +3132,7 @@
def _model_fn(features, labels, mode):
_, _ = features, labels
- variables.Variable(1., name='weight')
+ variables.VariableV1(1., name='weight')
return model_fn_lib.EstimatorSpec(
mode,
predictions=constant_op.constant(10.),
@@ -3081,7 +3175,7 @@
"""A graph that generates NaN's for testing."""
del features, labels
- global_step = variables.Variable(
+ global_step = variables.VariableV1(
0, dtype=dtypes.int64, name='global_step')
inc_global_step = state_ops.assign_add(global_step, 1)
nan_const = constant_op.constant(np.nan, dtype=dtypes.float32)
diff --git a/tensorflow/python/estimator/keras.py b/tensorflow/python/estimator/keras.py
index 6b2765b..5d5ed81 100644
--- a/tensorflow/python/estimator/keras.py
+++ b/tensorflow/python/estimator/keras.py
@@ -21,6 +21,7 @@
import os
import re
+import six
from tensorflow.python.client import session
from tensorflow.python.estimator import estimator as estimator_lib
@@ -31,6 +32,7 @@
from tensorflow.python.framework import sparse_tensor as sparse_tensor_lib
from tensorflow.python.framework import tensor_util
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 check_ops
@@ -214,25 +216,40 @@
if not getattr(model, 'metrics', None):
return None
- # TODO(psv/fchollet): support stateful metrics
eval_metric_ops = {}
+
+ def get_metric_name(metric):
+ if isinstance(metric, metrics.Metric):
+ return metric.name
+ if callable(metric):
+ return metric.__name__
+ assert isinstance(metric, six.string_types)
+ return metric
+
# When each metric maps to an output
if isinstance(model.metrics, dict):
for i, output_name in enumerate(model.metrics.keys()):
- metric_name = model.metrics[output_name]
- if callable(metric_name):
- metric_name = metric_name.__name__
+ # `metric` is the user given metric value in `compile`. This can be
+ # metric name (`acc`), metric function (binary_accuracy) or a metric
+ # object (BinaryAccuracy()).
+ metric = model.metrics[output_name]
+ metric_name = get_metric_name(metric)
# When some outputs use the same metric
if list(model.metrics.values()).count(metric_name) > 1:
metric_name += '_' + output_name
- eval_metric_ops[metric_name] = metrics_module.mean(
- model.metrics_tensors[i - len(model.metrics)])
+ if isinstance(metric, metrics.Metric):
+ eval_metric_ops[metric_name] = metric
+ else:
+ eval_metric_ops[metric_name] = metrics_module.mean(
+ model.metrics_tensors[i - len(model.metrics)])
else:
- for i, metric_name in enumerate(model.metrics):
- if callable(metric_name):
- metric_name = metric_name.__name__
- eval_metric_ops[metric_name] = metrics_module.mean(
- model.metrics_tensors[i])
+ for i, metric in enumerate(model.metrics):
+ metric_name = get_metric_name(metric)
+ if isinstance(metric, metrics.Metric):
+ eval_metric_ops[metric_name] = metric
+ else:
+ eval_metric_ops[metric_name] = metrics_module.mean(
+ model.metrics_tensors[i])
return eval_metric_ops
@@ -351,6 +368,44 @@
return latest_path
+def _get_file_from_google_storage(keras_model_path, model_dir):
+ """Get file from google storage and download to local file.
+
+ Args:
+ keras_model_path: a google storage path for compiled keras model.
+ model_dir: the directory from estimator config.
+
+ Returns:
+ The path where keras model is saved.
+
+ Raises:
+ ValueError: if storage object name does not end with .h5.
+ """
+ try:
+ from google.cloud import storage # pylint:disable=g-import-not-at-top
+ except ImportError:
+ raise TypeError('Could not save model to Google cloud storage; please '
+ 'install `google-cloud-storage` via '
+ '`pip install google-cloud-storage`.')
+ storage_client = storage.Client()
+ path, blob_name = os.path.split(keras_model_path)
+ _, bucket_name = os.path.split(path)
+ keras_model_dir = os.path.join(model_dir, 'keras')
+ if not gfile.Exists(keras_model_dir):
+ gfile.MakeDirs(keras_model_dir)
+ file_name = os.path.join(keras_model_dir, 'keras_model.h5')
+ try:
+ blob = storage_client.get_bucket(bucket_name).blob(blob_name)
+ blob.download_to_filename(file_name)
+ except:
+ raise ValueError('Failed to download keras model, please check '
+ 'environment variable GOOGLE_APPLICATION_CREDENTIALS '
+ 'and model path storage.googleapis.com/{bucket}/{object}.')
+ logging.info('Saving model to {}'.format(file_name))
+ del storage_client
+ return file_name
+
+
def model_to_estimator(keras_model=None,
keras_model_path=None,
custom_objects=None,
@@ -390,12 +445,13 @@
'Please specity either `keras_model` or `keras_model_path`, '
'but not both.')
+ config = estimator_lib.maybe_overwrite_model_dir_and_session_config(
+ config, model_dir)
if not keras_model:
if keras_model_path.startswith(
'gs://') or 'storage.googleapis.com' in keras_model_path:
- raise ValueError(
- '%s is not a local path. Please copy the model locally first.' %
- keras_model_path)
+ keras_model_path = _get_file_from_google_storage(keras_model_path,
+ config.model_dir)
logging.info('Loading models from %s', keras_model_path)
keras_model = models.load_model(keras_model_path)
else:
@@ -408,9 +464,6 @@
'Please compile the model with `model.compile()` '
'before calling `model_to_estimator()`.')
- config = estimator_lib.maybe_overwrite_model_dir_and_session_config(config,
- model_dir)
-
keras_model_fn = _create_keras_model_fn(keras_model, custom_objects)
if _any_weight_initialized(keras_model):
# Warn if config passed to estimator tries to update GPUOptions. If a
diff --git a/tensorflow/python/estimator/keras_test.py b/tensorflow/python/estimator/keras_test.py
index 3758243..4e285fa 100644
--- a/tensorflow/python/estimator/keras_test.py
+++ b/tensorflow/python/estimator/keras_test.py
@@ -257,7 +257,7 @@
keras_model.compile(
loss='categorical_crossentropy',
optimizer='rmsprop',
- metrics=['mse', keras.metrics.categorical_accuracy])
+ metrics=['mse', keras.metrics.CategoricalAccuracy()])
with self.cached_session():
est_keras = keras_lib.model_to_estimator(
@@ -281,7 +281,7 @@
keras_model.compile(
loss='categorical_crossentropy',
optimizer=rmsprop.RMSPropOptimizer(1e-3),
- metrics=['mse', keras.metrics.categorical_accuracy])
+ metrics=['mse', keras.metrics.CategoricalAccuracy()])
my_hook = MyHook()
with self.cached_session():
@@ -306,7 +306,7 @@
keras_model.compile(
loss='categorical_crossentropy',
optimizer=rmsprop.RMSPropOptimizer(1e-3),
- metrics=['mse', keras.metrics.categorical_accuracy])
+ metrics=['mse', keras.metrics.CategoricalAccuracy()])
my_hook = MyHook()
with self.cached_session():
keras_model.fit(x_train, y_train, epochs=1)
@@ -328,7 +328,7 @@
keras_model.compile(
loss='categorical_crossentropy',
optimizer=rmsprop.RMSPropOptimizer(1e-3),
- metrics=['mse', keras.metrics.categorical_accuracy])
+ metrics=['mse', keras.metrics.CategoricalAccuracy()])
with self.cached_session():
est_keras = keras_lib.model_to_estimator(
@@ -351,7 +351,7 @@
keras_model.compile(
loss='categorical_crossentropy',
optimizer=rmsprop.RMSPropOptimizer(1e-3),
- metrics=['mse', keras.metrics.categorical_accuracy])
+ metrics=['mse', keras.metrics.CategoricalAccuracy()])
with self.cached_session():
est_keras = keras_lib.model_to_estimator(
@@ -370,7 +370,7 @@
keras_model.compile(
loss='categorical_crossentropy',
optimizer=rmsprop.RMSPropOptimizer(1e-3),
- metrics=['mse', keras.metrics.categorical_accuracy])
+ metrics=['mse', keras.metrics.CategoricalAccuracy()])
with self.cached_session():
# Create state
@@ -581,12 +581,6 @@
with self.assertRaisesRegexp(ValueError, 'compiled'):
keras_lib.model_to_estimator(keras_model=keras_model)
- with self.cached_session():
- keras_model = simple_sequential_model()
- with self.assertRaisesRegexp(ValueError, 'not a local path'):
- keras_lib.model_to_estimator(
- keras_model_path='gs://bucket/object')
-
def test_invalid_ionames_error(self):
(x_train, y_train), (_, _) = testing_utils.get_test_data(
train_samples=_TRAIN_SIZE,
@@ -662,7 +656,7 @@
keras_model.compile(
loss='categorical_crossentropy',
optimizer='rmsprop',
- metrics=['mse', keras.metrics.categorical_accuracy])
+ metrics=['mse', keras.metrics.CategoricalAccuracy()])
tf_config = json.dumps({
'cluster': {
@@ -687,7 +681,7 @@
keras_model.compile(
loss='categorical_crossentropy',
optimizer='rmsprop',
- metrics=['mse', keras.metrics.categorical_accuracy])
+ metrics=['mse', keras.metrics.CategoricalAccuracy()])
gpu_options = config_pb2.GPUOptions(per_process_gpu_memory_fraction=0.3)
sess_config = config_pb2.ConfigProto(gpu_options=gpu_options)
@@ -706,7 +700,7 @@
keras_model.compile(
loss='categorical_crossentropy',
optimizer='rmsprop',
- metrics=['mse', keras.metrics.categorical_accuracy])
+ metrics=['mse', keras.metrics.CategoricalAccuracy()])
with self.cached_session():
est_keras = keras_lib.model_to_estimator(
@@ -736,7 +730,7 @@
keras_model.compile(
loss='categorical_crossentropy',
optimizer='rmsprop',
- metrics=['mse', keras.metrics.categorical_accuracy])
+ metrics=['mse', keras.metrics.CategoricalAccuracy()])
with self.cached_session():
with test.mock.patch.object(tempfile, 'mkdtemp', return_value=_TMP_DIR):
@@ -751,7 +745,7 @@
keras_model.compile(
loss='categorical_crossentropy',
optimizer='rmsprop',
- metrics=['mse', keras.metrics.categorical_accuracy])
+ metrics=['mse', keras.metrics.CategoricalAccuracy()])
with self.cached_session():
with self.assertRaisesRegexp(ValueError, '`model_dir` are set both in '
@@ -765,7 +759,7 @@
keras_model.compile(
loss='categorical_crossentropy',
optimizer=rmsprop.RMSPropOptimizer(1e-3),
- metrics=['mse', keras.metrics.categorical_accuracy])
+ metrics=['mse', keras.metrics.CategoricalAccuracy()])
with self.cached_session():
keras_model.train_on_batch(
np.random.random((10,) + _INPUT_SIZE),
@@ -776,7 +770,7 @@
keras_model.compile(
loss='categorical_crossentropy',
optimizer=SGD(lr=0.0001, momentum=0.9),
- metrics=['mse', keras.metrics.categorical_accuracy])
+ metrics=['mse', keras.metrics.CategoricalAccuracy()])
keras_lib.model_to_estimator(
keras_model=keras_model, config=self._config)
@@ -786,7 +780,7 @@
keras_model.compile(
loss='categorical_crossentropy',
optimizer=optimizer,
- metrics=['mse', keras.metrics.categorical_accuracy])
+ metrics=['mse', keras.metrics.CategoricalAccuracy()])
with self.cached_session() as sess:
keras_model_fn = keras_lib._create_keras_model_fn(keras_model)
global_step = training_util.create_global_step()
diff --git a/tensorflow/python/estimator/util.py b/tensorflow/python/estimator/util.py
index 31e4778..fb110c4 100644
--- a/tensorflow/python/estimator/util.py
+++ b/tensorflow/python/estimator/util.py
@@ -22,7 +22,6 @@
import os
import time
-from tensorflow.core.protobuf import config_pb2
from tensorflow.python.platform import gfile
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.training import training
@@ -144,14 +143,11 @@
self._finalize_fn = finalize_fn
def begin(self):
+ # We only create the init ops, but don't run it. We rely on SessionManager
+ # to run it for us.
self._init_ops = self._initialization_fn()
self._finalize_ops = self._finalize_fn()
- def after_create_session(self, session, coord):
- logging.info('Initialize system')
- session.run(self._init_ops,
- options=config_pb2.RunOptions(timeout_in_ms=5 * 60 * 1000))
-
def end(self, session):
logging.info('Finalize system.')
session.run(self._finalize_ops)
diff --git a/tensorflow/python/feature_column/BUILD b/tensorflow/python/feature_column/BUILD
index 5800b69..ac53a84 100644
--- a/tensorflow/python/feature_column/BUILD
+++ b/tensorflow/python/feature_column/BUILD
@@ -156,7 +156,7 @@
"//tensorflow/python:variables",
"//tensorflow/python/eager:backprop",
"//tensorflow/python/eager:context",
- "//tensorflow/python/estimator:estimator_py",
+ "//tensorflow/python/estimator:numpy_io",
"//third_party/py/numpy",
],
)
diff --git a/tensorflow/python/feature_column/feature_column.py b/tensorflow/python/feature_column/feature_column.py
index 9984379..618e70f 100644
--- a/tensorflow/python/feature_column/feature_column.py
+++ b/tensorflow/python/feature_column/feature_column.py
@@ -170,7 +170,8 @@
trainable=True,
cols_to_vars=None,
scope=None,
- cols_to_output_tensors=None):
+ cols_to_output_tensors=None,
+ from_template=False):
"""See input_layer. `scope` is a name or variable scope to use."""
feature_columns = _normalize_feature_columns(feature_columns)
@@ -186,10 +187,7 @@
if ops.GraphKeys.MODEL_VARIABLES not in weight_collections:
weight_collections.append(ops.GraphKeys.MODEL_VARIABLES)
- # a non-None `scope` can allow for variable reuse, when, e.g., this function
- # is wrapped by a `make_template`.
- with variable_scope.variable_scope(
- scope, default_name='input_layer', values=features.values()):
+ def _get_logits(): # pylint: disable=missing-docstring
builder = _LazyBuilder(features)
output_tensors = []
ordered_columns = []
@@ -217,6 +215,16 @@
_verify_static_batch_size_equality(output_tensors, ordered_columns)
return array_ops.concat(output_tensors, 1)
+ # If we're constructing from the `make_template`, that by default adds a
+ # variable scope with the name of the layer. In that case, we dont want to
+ # add another `variable_scope` as that would break checkpoints.
+ if from_template:
+ return _get_logits()
+ else:
+ with variable_scope.variable_scope(
+ scope, default_name='input_layer', values=features.values()):
+ return _get_logits()
+
@tf_export('feature_column.input_layer')
def input_layer(features,
@@ -301,17 +309,18 @@
feature_columns,
weight_collections=None,
trainable=True,
- cols_to_vars=None):
+ cols_to_vars=None,
+ name='feature_column_input_layer',
+ create_scope_now=True):
"""See `input_layer`."""
self._feature_columns = feature_columns
self._weight_collections = weight_collections
self._trainable = trainable
self._cols_to_vars = cols_to_vars
+ self._name = name
self._input_layer_template = template.make_template(
- 'feature_column_input_layer',
- _internal_input_layer,
- create_scope_now_=True)
+ self._name, _internal_input_layer, create_scope_now_=create_scope_now)
self._scope = self._input_layer_template.variable_scope
def __call__(self, features):
@@ -321,7 +330,11 @@
weight_collections=self._weight_collections,
trainable=self._trainable,
cols_to_vars=None,
- scope=self._scope)
+ from_template=True)
+
+ @property
+ def name(self):
+ return self._name
@property
def non_trainable_variables(self):
@@ -2305,7 +2318,7 @@
# Input_tensor must have rank 1.
if isinstance(input_tensor, sparse_tensor_lib.SparseTensor):
return sparse_ops.sparse_reshape(
- input_tensor, [array_ops.shape(input_tensor)[0], -1])
+ input_tensor, [array_ops.shape(input_tensor)[0], 1])
else:
return array_ops.expand_dims(input_tensor, -1)
diff --git a/tensorflow/python/feature_column/feature_column_test.py b/tensorflow/python/feature_column/feature_column_test.py
index abb79ef..1ae5102 100644
--- a/tensorflow/python/feature_column/feature_column_test.py
+++ b/tensorflow/python/feature_column/feature_column_test.py
@@ -169,6 +169,18 @@
TypeError, '"key" must be either a "str" or "_FeatureColumn".'):
builder.get(NotAFeatureColumn())
+ def test_expand_dim_rank_1_sparse_tensor_empty_batch(self):
+ # empty 1-D sparse tensor:
+ builder = _LazyBuilder(features={'a': sparse_tensor.SparseTensor(
+ indices=np.reshape(np.array([], dtype=np.int64), (0, 1)),
+ dense_shape=[0],
+ values=np.array([]))})
+ with self.cached_session():
+ spv = builder.get('a').eval()
+ self.assertAllEqual(np.array([0, 1], dtype=np.int64), spv.dense_shape)
+ self.assertAllEqual(
+ np.reshape(np.array([], dtype=np.int64), (0, 2)), spv.indices)
+
class NumericColumnTest(test.TestCase):
diff --git a/tensorflow/python/feature_column/feature_column_v2.py b/tensorflow/python/feature_column/feature_column_v2.py
index 57f7af7..b79373c 100644
--- a/tensorflow/python/feature_column/feature_column_v2.py
+++ b/tensorflow/python/feature_column/feature_column_v2.py
@@ -136,14 +136,11 @@
from tensorflow.python.eager import context
-from tensorflow.python.feature_column import feature_column as fc_old
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import sparse_tensor as sparse_tensor_lib
from tensorflow.python.framework import tensor_shape
-from tensorflow.python.keras.engine import training
from tensorflow.python.keras.engine.base_layer import Layer
-from tensorflow.python.layers import base
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import check_ops
from tensorflow.python.ops import control_flow_ops
@@ -153,7 +150,6 @@
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn_ops
from tensorflow.python.ops import parsing_ops
-from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import sparse_ops
from tensorflow.python.ops import string_ops
from tensorflow.python.ops import variable_scope
@@ -245,28 +241,19 @@
raise NotImplementedError('StateManager.get_resource')
-class _InputLayerStateManager(StateManager):
- """Manages the state of InputLayer."""
+class _StateManagerImpl(StateManager):
+ """Manages the state of FeatureLayer and LinearModel."""
- def __init__(self, layer, feature_columns, trainable):
- """Creates an _InputLayerStateManager object.
+ def __init__(self, layer, trainable):
+ """Creates an _StateManagerImpl object.
Args:
layer: The input layer this state manager is associated with.
- feature_columns: List of feature columns for the input layer
trainable: Whether by default, variables created are trainable or not.
"""
self._trainable = trainable
self._layer = layer
- self._cols_to_vars_map = {}
- self._cols_to_names_map = {}
- for column in sorted(feature_columns, key=lambda x: x.name):
- self._cols_to_vars_map[column] = {}
- base_name = column.name
- if isinstance(column, SharedEmbeddingColumn):
- base_name = column.shared_collection_name
- with variable_scope.variable_scope(base_name) as vs:
- self._cols_to_names_map[column] = _strip_leading_slashes(vs.name)
+ self._cols_to_vars_map = collections.defaultdict(lambda: {})
def create_variable(self,
feature_column,
@@ -277,19 +264,20 @@
initializer=None):
if name in self._cols_to_vars_map[feature_column]:
raise ValueError('Variable already exists.')
- with variable_scope.variable_scope(self._cols_to_names_map[feature_column]):
- var = self._layer.add_variable(
- name=name,
- shape=shape,
- dtype=dtype,
- initializer=initializer,
- trainable=self._trainable and trainable,
- # TODO(rohanj): Get rid of this hack once we have a mechanism for
- # specifying a default partitioner for an entire layer. In that case,
- # the default getter for Layers should work.
- getter=variable_scope.get_variable)
- self._cols_to_vars_map[feature_column][name] = var
- return var
+
+ var = self._layer.add_variable(
+ name=name,
+ shape=shape,
+ dtype=dtype,
+ initializer=initializer,
+ trainable=self._trainable and trainable,
+ use_resource=True,
+ # TODO(rohanj): Get rid of this hack once we have a mechanism for
+ # specifying a default partitioner for an entire layer. In that case,
+ # the default getter for Layers should work.
+ getter=variable_scope.get_variable)
+ self._cols_to_vars_map[feature_column][name] = var
+ return var
def get_variable(self, feature_column, name):
if name in self._cols_to_vars_map[feature_column]:
@@ -313,12 +301,15 @@
keywords_embedded = embedding_column(
categorical_column_with_hash_bucket("keywords", 10K), dimensions=16)
columns = [price, keywords_embedded, ...]
- features = tf.parse_example(..., features=make_parse_example_spec(columns))
feature_layer = FeatureLayer(columns)
+
+ features = tf.parse_example(..., features=make_parse_example_spec(columns))
dense_tensor = feature_layer(features)
for units in [128, 64, 32]:
dense_tensor = tf.layers.dense(dense_tensor, units, tf.nn.relu)
- prediction = tf.layers.dense(dense_tensor, 1)."""
+ prediction = tf.layers.dense(dense_tensor, 1).
+ ```
+ """
def __init__(self,
feature_columns,
@@ -375,8 +366,7 @@
super(FeatureLayer, self).__init__(name=name, trainable=trainable, **kwargs)
self._feature_columns = _normalize_feature_columns(feature_columns)
- self._state_manager = _InputLayerStateManager(self, self._feature_columns,
- self.trainable)
+ self._state_manager = _StateManagerImpl(self, self.trainable)
self._shared_state_manager = shared_state_manager
for column in sorted(self._feature_columns, key=lambda x: x.name):
if not isinstance(column, DenseColumn):
@@ -394,8 +384,9 @@
if isinstance(column, SharedEmbeddingColumn):
column.create_state(self._shared_state_manager)
else:
- with variable_scope.variable_scope(None, default_name=self.name):
- column.create_state(self._state_manager)
+ with variable_scope._pure_variable_scope(self.name): # pylint: disable=protected-access
+ with variable_scope._pure_variable_scope(column.name): # pylint: disable=protected-access
+ column.create_state(self._state_manager)
super(FeatureLayer, self).build(None)
def call(self, features, cols_to_output_tensors=None):
@@ -424,19 +415,20 @@
output_tensors = []
ordered_columns = []
for column in sorted(self._feature_columns, key=lambda x: x.name):
- ordered_columns.append(column)
- if isinstance(column, SharedEmbeddingColumn):
- tensor = column.get_dense_tensor(transformation_cache,
- self._shared_state_manager)
- else:
- tensor = column.get_dense_tensor(transformation_cache,
- self._state_manager)
- num_elements = column.variable_shape.num_elements()
- batch_size = array_ops.shape(tensor)[0]
- tensor = array_ops.reshape(tensor, shape=(batch_size, num_elements))
- output_tensors.append(tensor)
- if cols_to_output_tensors is not None:
- cols_to_output_tensors[column] = tensor
+ with ops.name_scope(column.name):
+ ordered_columns.append(column)
+ if isinstance(column, SharedEmbeddingColumn):
+ tensor = column.get_dense_tensor(transformation_cache,
+ self._shared_state_manager)
+ else:
+ tensor = column.get_dense_tensor(transformation_cache,
+ self._state_manager)
+ num_elements = column.variable_shape.num_elements()
+ batch_size = array_ops.shape(tensor)[0]
+ tensor = array_ops.reshape(tensor, shape=(batch_size, num_elements))
+ output_tensors.append(tensor)
+ if cols_to_output_tensors is not None:
+ cols_to_output_tensors[column] = tensor
_verify_static_batch_size_equality(output_tensors, ordered_columns)
return array_ops.concat(output_tensors, 1)
@@ -448,20 +440,18 @@
return (input_shape[0], total_elements)
-def linear_model(features,
- feature_columns,
- units=1,
- sparse_combiner='sum',
- weight_collections=None,
- trainable=True,
- cols_to_vars=None):
- """Returns a linear prediction `Tensor` based on given `feature_columns`.
+def _strip_leading_slashes(name):
+ return name.rsplit('/', 1)[-1]
- This function generates a weighted sum based on output dimension `units`.
+
+class LinearModel(Layer):
+ """Produces a linear prediction `Tensor` based on given `feature_columns`.
+
+ This layer generates a weighted sum based on output dimension `units`.
Weighted sum refers to logits in classification problems. It refers to the
prediction itself for linear regression problems.
- Note on supported columns: `linear_model` treats categorical columns as
+ Note on supported columns: `LinearModel` treats categorical columns as
`indicator_column`s. To be specific, assume the input as `SparseTensor` looks
like:
@@ -486,308 +476,195 @@
keywords = categorical_column_with_hash_bucket("keywords", 10K)
keywords_price = crossed_column('keywords', price_buckets, ...)
columns = [price_buckets, keywords, keywords_price ...]
+ linear_model = LinearModel(columns)
+
features = tf.parse_example(..., features=make_parse_example_spec(columns))
- prediction = linear_model(features, columns)
+ prediction = linear_model(features)
```
-
- Args:
- features: A mapping from key to tensors. `_FeatureColumn`s look up via these
- keys. For example `numeric_column('price')` will look at 'price' key in
- this dict. Values are `Tensor` or `SparseTensor` depending on
- corresponding `_FeatureColumn`.
- feature_columns: An iterable containing the FeatureColumns to use as inputs
- to your model. All items should be instances of classes derived from
- `_FeatureColumn`s.
- units: An integer, dimensionality of the output space. Default value is 1.
- sparse_combiner: A string specifying how to reduce if a categorical column
- is multivalent. Except `numeric_column`, almost all columns passed to
- `linear_model` are considered as categorical columns. It combines each
- categorical column independently. Currently "mean", "sqrtn" and "sum" are
- supported, with "sum" the default for linear model. "sqrtn" often achieves
- good accuracy, in particular with bag-of-words columns.
- * "sum": do not normalize features in the column
- * "mean": do l1 normalization on features in the column
- * "sqrtn": do l2 normalization on features in the column
- For example, for two features represented as the categorical columns:
-
- ```python
- # Feature 1
-
- shape = [2, 2]
- {
- [0, 0]: "a"
- [0, 1]: "b"
- [1, 0]: "c"
- }
-
- # Feature 2
-
- shape = [2, 3]
- {
- [0, 0]: "d"
- [1, 0]: "e"
- [1, 1]: "f"
- [1, 2]: "g"
- }
- ```
- with `sparse_combiner` as "mean", the linear model outputs conceptly are:
- ```
- y_0 = 1.0 / 2.0 * ( w_a + w_ b) + w_c + b_0
- y_1 = w_d + 1.0 / 3.0 * ( w_e + w_ f + w_g) + b_1
- ```
- where `y_i` is the output, `b_i` is the bias, and `w_x` is the weight
- assigned to the presence of `x` in the input features.
- weight_collections: A list of collection names to which the Variable will be
- added. Note that, variables will also be added to collections
- `tf.GraphKeys.GLOBAL_VARIABLES` and `ops.GraphKeys.MODEL_VARIABLES`.
- trainable: If `True` also add the variable to the graph collection
- `GraphKeys.TRAINABLE_VARIABLES` (see `tf.Variable`).
- cols_to_vars: If not `None`, must be a dictionary that will be filled with a
- mapping from `_FeatureColumn` to associated list of `Variable`s. For
- example, after the call, we might have cols_to_vars = {
- _NumericColumn(
- key='numeric_feature1', shape=(1,):
- [<tf.Variable 'linear_model/price2/weights:0' shape=(1, 1)>],
- 'bias': [<tf.Variable 'linear_model/bias_weights:0' shape=(1,)>],
- _NumericColumn(
- key='numeric_feature2', shape=(2,)):
- [<tf.Variable 'linear_model/price1/weights:0' shape=(2, 1)>]}
- If a column creates no variables, its value will be an empty list. Note
- that cols_to_vars will also contain a string key 'bias' that maps to a
- list of Variables.
-
- Returns:
- A `Tensor` which represents predictions/logits of a linear model. Its shape
- is (batch_size, units) and its dtype is `float32`.
-
- Raises:
- ValueError: if an item in `feature_columns` is neither a `_DenseColumn`
- nor `_CategoricalColumn`.
- """
- with variable_scope.variable_scope(None, 'linear_model') as vs:
- model_name = _strip_leading_slashes(vs.name)
- linear_model_layer = _LinearModel(
- feature_columns=feature_columns,
- units=units,
- sparse_combiner=sparse_combiner,
- weight_collections=weight_collections,
- trainable=trainable,
- name=model_name)
- retval = linear_model_layer(features) # pylint: disable=not-callable
- if cols_to_vars is not None:
- cols_to_vars.update(linear_model_layer.cols_to_vars())
- return retval
-
-
-def _add_to_collections(var, weight_collections):
- """Adds a var to the list of weight_collections provided.
-
- Handles the case for partitioned and non-partitioned variables.
-
- Args:
- var: A variable or Partitioned Variable.
- weight_collections: List of collections to add variable to.
- """
- for weight_collection in weight_collections:
- # The layer self.add_variable call already adds it to GLOBAL_VARIABLES.
- if weight_collection == ops.GraphKeys.GLOBAL_VARIABLES:
- continue
- # TODO(rohanj): Explore adding a _get_variable_list method on `Variable`
- # so that we don't have to do this check.
- if isinstance(var, variables.PartitionedVariable):
- for constituent_var in list(var):
- ops.add_to_collection(weight_collection, constituent_var)
- else:
- ops.add_to_collection(weight_collection, var)
-
-
-class _FCLinearWrapper(base.Layer):
- """Wraps a _FeatureColumn in a layer for use in a linear model.
-
- See `linear_model` above.
- """
-
- def __init__(self,
- feature_column,
- units=1,
- sparse_combiner='sum',
- weight_collections=None,
- trainable=True,
- name=None,
- **kwargs):
- super(_FCLinearWrapper, self).__init__(
- trainable=trainable, name=name, **kwargs)
- self._feature_column = feature_column
- self._units = units
- self._sparse_combiner = sparse_combiner
- self._weight_collections = weight_collections
-
- def build(self, _):
- if isinstance(self._feature_column, fc_old._CategoricalColumn): # pylint: disable=protected-access
- weight = self.add_variable(
- name='weights',
- shape=(self._feature_column._num_buckets, self._units), # pylint: disable=protected-access
- initializer=init_ops.zeros_initializer(),
- trainable=self.trainable)
- else:
- num_elements = self._feature_column._variable_shape.num_elements() # pylint: disable=protected-access
- weight = self.add_variable(
- name='weights',
- shape=[num_elements, self._units],
- initializer=init_ops.zeros_initializer(),
- trainable=self.trainable)
- _add_to_collections(weight, self._weight_collections)
- self._weight_var = weight
- self.built = True
-
- def call(self, builder):
- weighted_sum = fc_old._create_weighted_sum( # pylint: disable=protected-access
- column=self._feature_column,
- builder=builder,
- units=self._units,
- sparse_combiner=self._sparse_combiner,
- weight_collections=self._weight_collections,
- trainable=self.trainable,
- weight_var=self._weight_var)
- return weighted_sum
-
-
-class _BiasLayer(base.Layer):
- """A layer for the bias term.
- """
-
- def __init__(self,
- units=1,
- trainable=True,
- weight_collections=None,
- name=None,
- **kwargs):
- super(_BiasLayer, self).__init__(trainable=trainable, name=name, **kwargs)
- self._units = units
- self._weight_collections = weight_collections
-
- def build(self, _):
- self._bias_variable = self.add_variable(
- 'bias_weights',
- shape=[self._units],
- initializer=init_ops.zeros_initializer(),
- trainable=self.trainable)
- _add_to_collections(self._bias_variable, self._weight_collections)
- self.built = True
-
- def call(self, _):
- return self._bias_variable
-
-
-def _get_expanded_variable_list(var_list):
- returned_list = []
- for variable in var_list:
- if (isinstance(variable, variables.Variable) or
- resource_variable_ops.is_resource_variable(variable)):
- returned_list.append(variable) # Single variable case.
- else: # Must be a PartitionedVariable, so convert into a list.
- returned_list.extend(list(variable))
- return returned_list
-
-
-def _strip_leading_slashes(name):
- return name.rsplit('/', 1)[-1]
-
-
-class _LinearModel(training.Model):
- """Creates a linear model using feature columns.
-
- See `linear_model` for details.
"""
def __init__(self,
feature_columns,
units=1,
sparse_combiner='sum',
- weight_collections=None,
trainable=True,
name=None,
+ shared_state_manager=None,
**kwargs):
- super(_LinearModel, self).__init__(name=name, **kwargs)
- self._feature_columns = fc_old._normalize_feature_columns( # pylint: disable=protected-access
- feature_columns)
- self._weight_collections = list(weight_collections or [])
- if ops.GraphKeys.GLOBAL_VARIABLES not in self._weight_collections:
- self._weight_collections.append(ops.GraphKeys.GLOBAL_VARIABLES)
- if ops.GraphKeys.MODEL_VARIABLES not in self._weight_collections:
- self._weight_collections.append(ops.GraphKeys.MODEL_VARIABLES)
+ """Constructs a LinearModel.
- column_layers = {}
- for column in sorted(self._feature_columns, key=lambda x: x.name):
- with variable_scope.variable_scope(
- None, default_name=column._var_scope_name) as vs: # pylint: disable=protected-access
- # Having the fully expressed variable scope name ends up doubly
- # expressing the outer scope (scope with which this method was called)
- # in the name of the variable that would get created.
- column_name = _strip_leading_slashes(vs.name)
- column_layer = _FCLinearWrapper(column, units, sparse_combiner,
- self._weight_collections, trainable,
- column_name, **kwargs)
- column_layers[column_name] = column_layer
- self._column_layers = self._add_layers(column_layers)
- self._bias_layer = _BiasLayer(
- units=units,
- trainable=trainable,
- weight_collections=self._weight_collections,
- name='bias_layer',
- **kwargs)
- self._cols_to_vars = {}
+ Args:
+ feature_columns: An iterable containing the FeatureColumns to use as
+ inputs to your model. All items should be instances of classes derived
+ from `_FeatureColumn`s.
+ units: An integer, dimensionality of the output space. Default value is 1.
+ sparse_combiner: A string specifying how to reduce if a categorical column
+ is multivalent. Except `numeric_column`, almost all columns passed to
+ `linear_model` are considered as categorical columns. It combines each
+ categorical column independently. Currently "mean", "sqrtn" and "sum"
+ are supported, with "sum" the default for linear model. "sqrtn" often
+ achieves good accuracy, in particular with bag-of-words columns.
+ * "sum": do not normalize features in the column
+ * "mean": do l1 normalization on features in the column
+ * "sqrtn": do l2 normalization on features in the column
+ For example, for two features represented as the categorical columns:
- def cols_to_vars(self):
- """Returns a dict mapping _FeatureColumns to variables.
+ ```python
+ # Feature 1
- See `linear_model` for more information.
- This is not populated till `call` is called i.e. layer is built.
+ shape = [2, 2]
+ {
+ [0, 0]: "a"
+ [0, 1]: "b"
+ [1, 0]: "c"
+ }
+
+ # Feature 2
+
+ shape = [2, 3]
+ {
+ [0, 0]: "d"
+ [1, 0]: "e"
+ [1, 1]: "f"
+ [1, 2]: "g"
+ }
+ ```
+
+ with `sparse_combiner` as "mean", the linear model outputs conceptly are
+ ```
+ y_0 = 1.0 / 2.0 * ( w_a + w_ b) + w_c + b_0
+ y_1 = w_d + 1.0 / 3.0 * ( w_e + w_ f + w_g) + b_1
+ ```
+ where `y_i` is the output, `b_i` is the bias, and `w_x` is the weight
+ assigned to the presence of `x` in the input features.
+ trainable: If `True` also add the variable to the graph collection
+ `GraphKeys.TRAINABLE_VARIABLES` (see `tf.Variable`).
+ name: Name to give to the Linear Model. All variables and ops created will
+ be scoped by this name.
+ shared_state_manager: SharedEmbeddingStateManager that manages the state
+ of SharedEmbeddingColumns. For more info, look at `FeatureLayer`.
+ **kwargs: Keyword arguments to construct a layer.
+
+ Raises:
+ ValueError: if an item in `feature_columns` is neither a `DenseColumn`
+ nor `CategoricalColumn`.
"""
- return self._cols_to_vars
+ super(LinearModel, self).__init__(name=name, trainable=trainable, **kwargs)
+
+ self._feature_columns = _normalize_feature_columns(feature_columns)
+ self._feature_columns = sorted(self._feature_columns, key=lambda x: x.name)
+ for column in self._feature_columns:
+ if not isinstance(column, (DenseColumn, CategoricalColumn)):
+ raise ValueError(
+ 'Items of feature_columns must be either a '
+ 'DenseColumn or CategoricalColumn. Given: {}'.format(column))
+
+ self._units = units
+ self._sparse_combiner = sparse_combiner
+
+ self._state_manager = _StateManagerImpl(self, self.trainable)
+ self._shared_state_manager = shared_state_manager
+ self._bias_variable = None
+
+ def build(self, _):
+ # Create state for shared embedding columns.
+ for column in self._feature_columns:
+ if isinstance(column, SharedEmbeddingColumn):
+ column.create_state(self._shared_state_manager)
+
+ # We need variable scopes for now because we want the variable partitioning
+ # information to percolate down. We also use _pure_variable_scope's here
+ # since we want to open up a name_scope in the `call` method while creating
+ # the ops.
+ with variable_scope._pure_variable_scope(self.name): # pylint: disable=protected-access
+ for column in self._feature_columns:
+ with variable_scope._pure_variable_scope(column.name): # pylint: disable=protected-access
+ # Create the state for each feature column
+ if not isinstance(column, SharedEmbeddingColumn):
+ column.create_state(self._state_manager)
+
+ # Create a weight variable for each column.
+ if isinstance(column, CategoricalColumn):
+ first_dim = column.num_buckets
+ else:
+ first_dim = column.variable_shape.num_elements()
+ self._state_manager.create_variable(
+ column,
+ name='weights',
+ dtype=dtypes.float32,
+ shape=(first_dim, self._units),
+ initializer=init_ops.zeros_initializer(),
+ trainable=self.trainable)
+
+ # Create a bias variable.
+ self._bias_variable = self.add_variable(
+ name='bias_weights',
+ dtype=dtypes.float32,
+ shape=[self._units],
+ initializer=init_ops.zeros_initializer(),
+ trainable=self.trainable,
+ use_resource=True,
+ # TODO(rohanj): Get rid of this hack once we have a mechanism for
+ # specifying a default partitioner for an entire layer. In that case,
+ # the default getter for Layers should work.
+ getter=variable_scope.get_variable)
+
+ super(LinearModel, self).build(None)
def call(self, features):
- with variable_scope.variable_scope(self.name):
- for column in self._feature_columns:
- if not isinstance(
- column,
- (
- fc_old._DenseColumn, # pylint: disable=protected-access
- fc_old._CategoricalColumn)): # pylint: disable=protected-access
- raise ValueError(
- 'Items of feature_columns must be either a '
- '_DenseColumn or _CategoricalColumn. Given: {}'.format(column))
- weighted_sums = []
- ordered_columns = []
- builder = fc_old._LazyBuilder(features) # pylint: disable=protected-access
- for layer in sorted(self._column_layers.values(), key=lambda x: x.name):
- column = layer._feature_column # pylint: disable=protected-access
- ordered_columns.append(column)
- weighted_sum = layer(builder)
- weighted_sums.append(weighted_sum)
- self._cols_to_vars[column] = ops.get_collection(
- ops.GraphKeys.GLOBAL_VARIABLES, scope=layer.scope_name)
+ """Returns a `Tensor` the represents the predictions of a linear model.
- _verify_static_batch_size_equality(weighted_sums, ordered_columns)
+ Args:
+ features: A mapping from key to tensors. `_FeatureColumn`s look up via
+ these keys. For example `numeric_column('price')` will look at 'price'
+ key in this dict. Values are `Tensor` or `SparseTensor` depending on
+ corresponding `_FeatureColumn`.
+
+ Returns:
+ A `Tensor` which represents predictions/logits of a linear model. Its
+ shape is (batch_size, units) and its dtype is `float32`.
+
+ Raises:
+ ValueError: If features are not a dictionary.
+ """
+ if not isinstance(features, dict):
+ raise ValueError('We expected a dictionary here. Instead we got: ',
+ features)
+ with ops.name_scope(self.name):
+ transformation_cache = FeatureTransformationCache(features)
+ weighted_sums = []
+ for column in self._feature_columns:
+ with ops.name_scope(column.name):
+ # All the weights used in the linear model are owned by the state
+ # manager associated with this Linear Model.
+ weight_var = self._state_manager.get_variable(column, 'weights')
+
+ # The embedding weights for the SharedEmbeddingColumn are owned by
+ # the shared_state_manager and so we need to pass that in while
+ # creating the weighted sum. For all other columns, the state is owned
+ # by the Linear Model's state manager.
+ if isinstance(column, SharedEmbeddingColumn):
+ state_manager = self._shared_state_manager
+ else:
+ state_manager = self._state_manager
+ weighted_sum = _create_weighted_sum(
+ column=column,
+ transformation_cache=transformation_cache,
+ state_manager=state_manager,
+ sparse_combiner=self._sparse_combiner,
+ weight_var=weight_var)
+ weighted_sums.append(weighted_sum)
+
+ _verify_static_batch_size_equality(weighted_sums, self._feature_columns)
predictions_no_bias = math_ops.add_n(
weighted_sums, name='weighted_sum_no_bias')
predictions = nn_ops.bias_add(
- predictions_no_bias,
- self._bias_layer( # pylint: disable=not-callable
- builder,
- scope=variable_scope.get_variable_scope()), # pylint: disable=not-callable
- name='weighted_sum')
- bias = self._bias_layer.variables[0]
- self._cols_to_vars['bias'] = _get_expanded_variable_list([bias])
- return predictions
+ predictions_no_bias, self._bias_variable, name='weighted_sum')
+ return predictions
- def _add_layers(self, layers):
- # "Magic" required for keras.Model classes to track all the variables in
- # a list of layers.Layer objects.
- # TODO(ashankar): Figure out API so user code doesn't have to do this.
- for name, layer in layers.items():
- setattr(self, 'layer-%s' % name, layer)
- return layers
+ @property
+ def bias_variable(self):
+ return self._bias_variable
def _transform_features(features, feature_columns, state_manager):
@@ -2045,58 +1922,40 @@
pass
-def _create_weighted_sum(column,
- transformation_cache,
- state_manager,
- units,
- sparse_combiner,
- weight_collections,
- trainable,
- weight_var=None):
+def is_feature_column_v2(feature_columns):
+ """Returns True if all feature columns are V2."""
+ for feature_column in feature_columns:
+ if not isinstance(feature_column, FeatureColumn):
+ return False
+ return True
+
+
+def _create_weighted_sum(column, transformation_cache, state_manager,
+ sparse_combiner, weight_var):
"""Creates a weighted sum for a dense/categorical column for linear_model."""
if isinstance(column, CategoricalColumn):
return _create_categorical_column_weighted_sum(
column=column,
transformation_cache=transformation_cache,
state_manager=state_manager,
- units=units,
sparse_combiner=sparse_combiner,
- weight_collections=weight_collections,
- trainable=trainable,
weight_var=weight_var)
else:
return _create_dense_column_weighted_sum(
column=column,
transformation_cache=transformation_cache,
state_manager=state_manager,
- units=units,
- weight_collections=weight_collections,
- trainable=trainable,
weight_var=weight_var)
-def _create_dense_column_weighted_sum(column,
- transformation_cache,
- state_manager,
- units,
- weight_collections,
- trainable,
- weight_var=None):
+def _create_dense_column_weighted_sum(column, transformation_cache,
+ state_manager, weight_var):
"""Create a weighted sum of a dense column for linear_model."""
tensor = column.get_dense_tensor(transformation_cache, state_manager)
num_elements = column.variable_shape.num_elements()
batch_size = array_ops.shape(tensor)[0]
tensor = array_ops.reshape(tensor, shape=(batch_size, num_elements))
- if weight_var is not None:
- weight = weight_var
- else:
- weight = variable_scope.get_variable(
- name='weights',
- shape=[num_elements, units],
- initializer=init_ops.zeros_initializer(),
- trainable=trainable,
- collections=weight_collections)
- return math_ops.matmul(tensor, weight, name='weighted_sum')
+ return math_ops.matmul(tensor, weight_var, name='weighted_sum')
class CategoricalColumn(FeatureColumn):
@@ -2137,14 +1996,8 @@
pass
-def _create_categorical_column_weighted_sum(column,
- transformation_cache,
- state_manager,
- units,
- sparse_combiner,
- weight_collections,
- trainable,
- weight_var=None):
+def _create_categorical_column_weighted_sum(
+ column, transformation_cache, state_manager, sparse_combiner, weight_var):
# pylint: disable=g-doc-return-or-yield,g-doc-args
"""Create a weighted sum of a categorical column for linear_model.
@@ -2183,17 +2036,8 @@
weight_tensor = sparse_ops.sparse_reshape(
weight_tensor, [array_ops.shape(weight_tensor)[0], -1])
- if weight_var is not None:
- weight = weight_var
- else:
- weight = variable_scope.get_variable(
- name='weights',
- shape=(column.num_buckets, units),
- initializer=init_ops.zeros_initializer(),
- trainable=trainable,
- collections=weight_collections)
return _safe_embedding_lookup_sparse(
- weight,
+ weight_var,
id_tensor,
sparse_weights=weight_tensor,
combiner=sparse_combiner,
@@ -2333,7 +2177,7 @@
# Input_tensor must have rank 1.
if isinstance(input_tensor, sparse_tensor_lib.SparseTensor):
return sparse_ops.sparse_reshape(
- input_tensor, [array_ops.shape(input_tensor)[0], -1])
+ input_tensor, [array_ops.shape(input_tensor)[0], 1])
else:
return array_ops.expand_dims(input_tensor, -1)
@@ -2769,6 +2613,7 @@
dtype=dtype,
trainable=self.trainable and trainable,
initializer=initializer,
+ use_resource=True,
# TODO(rohanj): Get rid of this hack once we have a mechanism for
# specifying a default partitioner for an entire layer. In that case,
# the default getter for Layers should work.
@@ -2782,6 +2627,12 @@
return self._var_dict[name]
+def maybe_create_shared_state_manager(feature_columns):
+ if is_feature_column_v2(feature_columns):
+ return SharedEmbeddingStateManager()
+ return None
+
+
class SharedEmbeddingColumn(
DenseColumn, SequenceDenseColumn,
collections.namedtuple(
@@ -2822,6 +2673,10 @@
def create_state(self, state_manager):
"""Creates the shared embedding lookup variable."""
+ if not isinstance(state_manager, SharedEmbeddingStateManager):
+ raise ValueError('Expected state_manager to be of type '
+ 'SharedEmbeddingStateManager. Obtained type: {}'.format(
+ type(state_manager)))
embedding_shape = (self.categorical_column.num_buckets, self.dimension)
state_manager.create_variable(
name=self.shared_collection_name,
@@ -3433,11 +3288,10 @@
raise ValueError('Missing embedding_weights %s.' % embedding_weights)
dtype = sparse_weights.dtype if sparse_weights is not None else None
- if not isinstance(embedding_weights[0],
- resource_variable_ops.ResourceVariable):
- embedding_weights = [
- ops.convert_to_tensor(w, dtype=dtype) for w in embedding_weights
- ]
+ # TODO(rohanj): Look into removing this convert_to_tensor call.
+ embedding_weights = [
+ ops.convert_to_tensor(w, dtype=dtype) for w in embedding_weights
+ ]
with ops.name_scope(name, 'embedding_lookup',
embedding_weights + [sparse_ids,
diff --git a/tensorflow/python/feature_column/feature_column_v2_test.py b/tensorflow/python/feature_column/feature_column_v2_test.py
index 58168e0..d378714 100644
--- a/tensorflow/python/feature_column/feature_column_v2_test.py
+++ b/tensorflow/python/feature_column/feature_column_v2_test.py
@@ -31,9 +31,7 @@
from tensorflow.python.eager import backprop
from tensorflow.python.eager import context
from tensorflow.python.estimator.inputs import numpy_io
-from tensorflow.python.feature_column import feature_column as fc_old
from tensorflow.python.feature_column import feature_column_v2 as fc
-from tensorflow.python.feature_column.feature_column_v2 import _LinearModel
from tensorflow.python.feature_column.feature_column_v2 import _transform_features
from tensorflow.python.feature_column.feature_column_v2 import FeatureColumn
from tensorflow.python.feature_column.feature_column_v2 import FeatureLayer
@@ -48,7 +46,6 @@
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import lookup_ops
from tensorflow.python.ops import parsing_ops
-from tensorflow.python.ops import partitioned_variables
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import variables as variables_lib
from tensorflow.python.platform import test
@@ -177,6 +174,22 @@
TypeError, '"key" must be either a "str" or "FeatureColumn".'):
transformation_cache.get(NotAFeatureColumn(), None)
+ def test_expand_dim_rank_1_sparse_tensor_empty_batch(self):
+ # empty 1-D sparse tensor:
+ transformation_cache = FeatureTransformationCache(
+ features={
+ 'a':
+ sparse_tensor.SparseTensor(
+ indices=np.reshape(np.array([], dtype=np.int64), (0, 1)),
+ dense_shape=[0],
+ values=np.array([]))
+ })
+ with self.cached_session():
+ spv = transformation_cache.get('a', None).eval()
+ self.assertAllEqual(np.array([0, 1], dtype=np.int64), spv.dense_shape)
+ self.assertAllEqual(
+ np.reshape(np.array([], dtype=np.int64), (0, 2)), spv.indices)
+
class NumericColumnTest(test.TestCase):
@@ -344,26 +357,12 @@
self.assertEqual(a.default_value, ((3., 2.),))
def test_linear_model(self):
- price = fc_old.numeric_column('price')
+ price = fc.numeric_column('price')
with ops.Graph().as_default():
features = {'price': [[1.], [5.]]}
- predictions = fc.linear_model(features, [price])
- bias = get_linear_model_bias()
- price_var = get_linear_model_column_var(price)
- with _initialized_session() as sess:
- self.assertAllClose([0.], bias.eval())
- self.assertAllClose([[0.]], price_var.eval())
- self.assertAllClose([[0.], [0.]], predictions.eval())
- sess.run(price_var.assign([[10.]]))
- self.assertAllClose([[10.], [50.]], predictions.eval())
-
- def test_keras_linear_model(self):
- price = fc_old.numeric_column('price')
- with ops.Graph().as_default():
- features = {'price': [[1.], [5.]]}
- predictions = get_keras_linear_model_predictions(features, [price])
- bias = get_linear_model_bias()
- price_var = get_linear_model_column_var(price)
+ model = fc.LinearModel([price])
+ predictions = model(features)
+ price_var, bias = model.variables
with _initialized_session() as sess:
self.assertAllClose([0.], bias.eval())
self.assertAllClose([[0.]], price_var.eval())
@@ -548,13 +547,13 @@
def test_linear_model_one_input_value(self):
"""Tests linear_model() for input with shape=[1]."""
- price = fc_old.numeric_column('price', shape=[1])
- bucketized_price = fc_old.bucketized_column(price, boundaries=[0, 2, 4, 6])
+ price = fc.numeric_column('price', shape=[1])
+ bucketized_price = fc.bucketized_column(price, boundaries=[0, 2, 4, 6])
with ops.Graph().as_default():
features = {'price': [[-1.], [1.], [5.], [6.]]}
- predictions = fc.linear_model(features, [bucketized_price])
- bias = get_linear_model_bias()
- bucketized_price_var = get_linear_model_column_var(bucketized_price)
+ model = fc.LinearModel([bucketized_price])
+ predictions = model(features)
+ bucketized_price_var, bias = model.variables
with _initialized_session() as sess:
self.assertAllClose([0.], bias.eval())
# One weight variable per bucket, all initialized to zero.
@@ -573,13 +572,13 @@
def test_linear_model_two_input_values(self):
"""Tests linear_model() for input with shape=[2]."""
- price = fc_old.numeric_column('price', shape=[2])
- bucketized_price = fc_old.bucketized_column(price, boundaries=[0, 2, 4, 6])
+ price = fc.numeric_column('price', shape=[2])
+ bucketized_price = fc.bucketized_column(price, boundaries=[0, 2, 4, 6])
with ops.Graph().as_default():
features = {'price': [[-1., 1.], [5., 6.]]}
- predictions = fc.linear_model(features, [bucketized_price])
- bias = get_linear_model_bias()
- bucketized_price_var = get_linear_model_column_var(bucketized_price)
+ model = fc.LinearModel([bucketized_price])
+ predictions = model(features)
+ bucketized_price_var, bias = model.variables
with _initialized_session() as sess:
self.assertAllClose([0.], bias.eval())
# One weight per bucket per input column, all initialized to zero.
@@ -600,62 +599,6 @@
sess.run(bias.assign([1.]))
self.assertAllClose([[81.], [141.]], predictions.eval())
- def test_keras_linear_model_one_input_value(self):
- """Tests _LinearModel for input with shape=[1]."""
- price = fc_old.numeric_column('price', shape=[1])
- bucketized_price = fc_old.bucketized_column(price, boundaries=[0, 2, 4, 6])
- with ops.Graph().as_default():
- features = {'price': [[-1.], [1.], [5.], [6.]]}
- predictions = get_keras_linear_model_predictions(features,
- [bucketized_price])
- bias = get_linear_model_bias()
- bucketized_price_var = get_linear_model_column_var(bucketized_price)
- with _initialized_session() as sess:
- self.assertAllClose([0.], bias.eval())
- # One weight variable per bucket, all initialized to zero.
- self.assertAllClose([[0.], [0.], [0.], [0.], [0.]],
- bucketized_price_var.eval())
- self.assertAllClose([[0.], [0.], [0.], [0.]], predictions.eval())
- sess.run(
- bucketized_price_var.assign([[10.], [20.], [30.], [40.], [50.]]))
- # price -1. is in the 0th bucket, whose weight is 10.
- # price 1. is in the 1st bucket, whose weight is 20.
- # price 5. is in the 3rd bucket, whose weight is 40.
- # price 6. is in the 4th bucket, whose weight is 50.
- self.assertAllClose([[10.], [20.], [40.], [50.]], predictions.eval())
- sess.run(bias.assign([1.]))
- self.assertAllClose([[11.], [21.], [41.], [51.]], predictions.eval())
-
- def test_keras_linear_model_two_input_values(self):
- """Tests _LinearModel for input with shape=[2]."""
- price = fc_old.numeric_column('price', shape=[2])
- bucketized_price = fc_old.bucketized_column(price, boundaries=[0, 2, 4, 6])
- with ops.Graph().as_default():
- features = {'price': [[-1., 1.], [5., 6.]]}
- predictions = get_keras_linear_model_predictions(features,
- [bucketized_price])
- bias = get_linear_model_bias()
- bucketized_price_var = get_linear_model_column_var(bucketized_price)
- with _initialized_session() as sess:
- self.assertAllClose([0.], bias.eval())
- # One weight per bucket per input column, all initialized to zero.
- self.assertAllClose(
- [[0.], [0.], [0.], [0.], [0.], [0.], [0.], [0.], [0.], [0.]],
- bucketized_price_var.eval())
- self.assertAllClose([[0.], [0.]], predictions.eval())
- sess.run(
- bucketized_price_var.assign([[10.], [20.], [30.], [40.], [50.],
- [60.], [70.], [80.], [90.], [100.]]))
- # 1st example:
- # price -1. is in the 0th bucket, whose weight is 10.
- # price 1. is in the 6th bucket, whose weight is 70.
- # 2nd example:
- # price 5. is in the 3rd bucket, whose weight is 40.
- # price 6. is in the 9th bucket, whose weight is 100.
- self.assertAllClose([[80.], [140.]], predictions.eval())
- sess.run(bias.assign([1.]))
- self.assertAllClose([[81.], [141.]], predictions.eval())
-
class HashedCategoricalColumnTest(test.TestCase):
@@ -836,39 +779,18 @@
transformation_cache.get(hashed_sparse, None), id_weight_pair.id_tensor)
def test_linear_model(self):
- wire_column = fc_old.categorical_column_with_hash_bucket('wire', 4)
- self.assertEqual(4, wire_column._num_buckets)
+ wire_column = fc.categorical_column_with_hash_bucket('wire', 4)
+ self.assertEqual(4, wire_column.num_buckets)
with ops.Graph().as_default():
- predictions = fc.linear_model({
- wire_column.name: sparse_tensor.SparseTensorValue(
- indices=((0, 0), (1, 0), (1, 1)),
- values=('marlo', 'skywalker', 'omar'),
- dense_shape=(2, 2))
- }, (wire_column,))
- bias = get_linear_model_bias()
- wire_var = get_linear_model_column_var(wire_column)
- with _initialized_session():
- self.assertAllClose((0.,), bias.eval())
- self.assertAllClose(((0.,), (0.,), (0.,), (0.,)), wire_var.eval())
- self.assertAllClose(((0.,), (0.,)), predictions.eval())
- wire_var.assign(((1.,), (2.,), (3.,), (4.,))).eval()
- # 'marlo' -> 3: wire_var[3] = 4
- # 'skywalker' -> 2, 'omar' -> 2: wire_var[2] + wire_var[2] = 3+3 = 6
- self.assertAllClose(((4.,), (6.,)), predictions.eval())
-
- def test_keras_linear_model(self):
- wire_column = fc_old.categorical_column_with_hash_bucket('wire', 4)
- self.assertEqual(4, wire_column._num_buckets)
- with ops.Graph().as_default():
- predictions = get_keras_linear_model_predictions({
+ model = fc.LinearModel((wire_column,))
+ predictions = model({
wire_column.name:
sparse_tensor.SparseTensorValue(
indices=((0, 0), (1, 0), (1, 1)),
values=('marlo', 'skywalker', 'omar'),
dense_shape=(2, 2))
- }, (wire_column,))
- bias = get_linear_model_bias()
- wire_var = get_linear_model_column_var(wire_column)
+ })
+ wire_var, bias = model.variables
with _initialized_session():
self.assertAllClose((0.,), bias.eval())
self.assertAllClose(((0.,), (0.,), (0.,), (0.,)), wire_var.eval())
@@ -1087,19 +1009,21 @@
Uses data from test_get_sparse_tesnsors_simple.
"""
- a = fc_old.numeric_column('a', dtype=dtypes.int32, shape=(2,))
- b = fc_old.bucketized_column(a, boundaries=(0, 1))
- crossed = fc_old.crossed_column([b, 'c'], hash_bucket_size=5, hash_key=5)
+ a = fc.numeric_column('a', dtype=dtypes.int32, shape=(2,))
+ b = fc.bucketized_column(a, boundaries=(0, 1))
+ crossed = fc.crossed_column([b, 'c'], hash_bucket_size=5, hash_key=5)
with ops.Graph().as_default():
- predictions = fc.linear_model({
- 'a': constant_op.constant(((-1., .5), (.5, 1.))),
- 'c': sparse_tensor.SparseTensor(
- indices=((0, 0), (1, 0), (1, 1)),
- values=['cA', 'cB', 'cC'],
- dense_shape=(2, 2)),
- }, (crossed,))
- bias = get_linear_model_bias()
- crossed_var = get_linear_model_column_var(crossed)
+ model = fc.LinearModel((crossed,))
+ predictions = model({
+ 'a':
+ constant_op.constant(((-1., .5), (.5, 1.))),
+ 'c':
+ sparse_tensor.SparseTensor(
+ indices=((0, 0), (1, 0), (1, 1)),
+ values=['cA', 'cB', 'cC'],
+ dense_shape=(2, 2)),
+ })
+ crossed_var, bias = model.variables
with _initialized_session() as sess:
self.assertAllClose((0.,), bias.eval())
self.assertAllClose(
@@ -1113,7 +1037,7 @@
def test_linear_model_with_weights(self):
- class _TestColumnWithWeights(fc_old._CategoricalColumn):
+ class _TestColumnWithWeights(fc.CategoricalColumn):
"""Produces sparse IDs and sparse weights."""
@property
@@ -1121,7 +1045,7 @@
return 'test_column'
@property
- def _parse_example_spec(self):
+ def parse_example_spec(self):
return {
self.name: parsing_ops.VarLenFeature(dtypes.int32),
'{}_weights'.format(self.name): parsing_ops.VarLenFeature(
@@ -1129,114 +1053,28 @@
}
@property
- def _num_buckets(self):
+ def num_buckets(self):
return 5
- def _transform_feature(self, inputs):
- return (inputs.get(self.name),
- inputs.get('{}_weights'.format(self.name)))
+ def transform_feature(self, transformation_cache, state_manager):
+ return (transformation_cache.get(self.name, state_manager),
+ transformation_cache.get('{}_weights'.format(self.name),
+ state_manager))
- def _get_sparse_tensors(self, inputs, weight_collections=None,
- trainable=None):
+ def get_sparse_tensors(self, transformation_cache, state_manager):
"""Populates both id_tensor and weight_tensor."""
- ids_and_weights = inputs.get(self)
- return fc_old._CategoricalColumn.IdWeightPair(
+ ids_and_weights = transformation_cache.get(self, state_manager)
+ return fc.CategoricalColumn.IdWeightPair(
id_tensor=ids_and_weights[0], weight_tensor=ids_and_weights[1])
t = _TestColumnWithWeights()
- crossed = fc_old.crossed_column([t, 'c'], hash_bucket_size=5, hash_key=5)
+ crossed = fc.crossed_column([t, 'c'], hash_bucket_size=5, hash_key=5)
with ops.Graph().as_default():
with self.assertRaisesRegexp(
ValueError,
'crossed_column does not support weight_tensor.*{}'.format(t.name)):
- fc.linear_model({
- t.name: sparse_tensor.SparseTensor(
- indices=((0, 0), (1, 0), (1, 1)),
- values=[0, 1, 2],
- dense_shape=(2, 2)),
- '{}_weights'.format(t.name): sparse_tensor.SparseTensor(
- indices=((0, 0), (1, 0), (1, 1)),
- values=[1., 10., 2.],
- dense_shape=(2, 2)),
- 'c': sparse_tensor.SparseTensor(
- indices=((0, 0), (1, 0), (1, 1)),
- values=['cA', 'cB', 'cC'],
- dense_shape=(2, 2)),
- }, (crossed,))
-
- def test_keras_linear_model(self):
- """Tests _LinearModel.
-
- Uses data from test_get_sparse_tesnsors_simple.
- """
- a = fc_old.numeric_column('a', dtype=dtypes.int32, shape=(2,))
- b = fc_old.bucketized_column(a, boundaries=(0, 1))
- crossed = fc_old.crossed_column([b, 'c'], hash_bucket_size=5, hash_key=5)
- with ops.Graph().as_default():
- predictions = get_keras_linear_model_predictions({
- 'a':
- constant_op.constant(((-1., .5), (.5, 1.))),
- 'c':
- sparse_tensor.SparseTensor(
- indices=((0, 0), (1, 0), (1, 1)),
- values=['cA', 'cB', 'cC'],
- dense_shape=(2, 2)),
- }, (crossed,))
- bias = get_linear_model_bias()
- crossed_var = get_linear_model_column_var(crossed)
- with _initialized_session() as sess:
- self.assertAllClose((0.,), bias.eval())
- self.assertAllClose(((0.,), (0.,), (0.,), (0.,), (0.,)),
- crossed_var.eval())
- self.assertAllClose(((0.,), (0.,)), predictions.eval())
- sess.run(crossed_var.assign(((1.,), (2.,), (3.,), (4.,), (5.,))))
- # Expected ids after cross = (1, 0, 1, 3, 4, 2)
- self.assertAllClose(((3.,), (14.,)), predictions.eval())
- sess.run(bias.assign((.1,)))
- self.assertAllClose(((3.1,), (14.1,)), predictions.eval())
-
- def test_keras_linear_model_with_weights(self):
-
- class _TestColumnWithWeights(fc_old._CategoricalColumn):
- """Produces sparse IDs and sparse weights."""
-
- @property
- def name(self):
- return 'test_column'
-
- @property
- def _parse_example_spec(self):
- return {
- self.name:
- parsing_ops.VarLenFeature(dtypes.int32),
- '{}_weights'.format(self.name):
- parsing_ops.VarLenFeature(dtypes.float32),
- }
-
- @property
- def _num_buckets(self):
- return 5
-
- def _transform_feature(self, inputs):
- return (inputs.get(self.name),
- inputs.get('{}_weights'.format(self.name)))
-
- def _get_sparse_tensors(self,
- inputs,
- weight_collections=None,
- trainable=None):
- """Populates both id_tensor and weight_tensor."""
- ids_and_weights = inputs.get(self)
- return fc_old._CategoricalColumn.IdWeightPair(
- id_tensor=ids_and_weights[0], weight_tensor=ids_and_weights[1])
-
- t = _TestColumnWithWeights()
- crossed = fc_old.crossed_column([t, 'c'], hash_bucket_size=5, hash_key=5)
- with ops.Graph().as_default():
- with self.assertRaisesRegexp(
- ValueError,
- 'crossed_column does not support weight_tensor.*{}'.format(t.name)):
- get_keras_linear_model_predictions({
+ model = fc.LinearModel((crossed,))
+ model({
t.name:
sparse_tensor.SparseTensor(
indices=((0, 0), (1, 0), (1, 1)),
@@ -1252,37 +1090,7 @@
indices=((0, 0), (1, 0), (1, 1)),
values=['cA', 'cB', 'cC'],
dense_shape=(2, 2)),
- }, (crossed,))
-
-
-def get_linear_model_bias(name='linear_model'):
- with variable_scope.variable_scope(name, reuse=True):
- return variable_scope.get_variable('bias_weights')
-
-
-def get_linear_model_column_var(column, name='linear_model'):
- return ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES,
- name + '/' + column.name)[0]
-
-
-def get_keras_linear_model_predictions(features,
- feature_columns,
- units=1,
- sparse_combiner='sum',
- weight_collections=None,
- trainable=True,
- cols_to_vars=None):
- keras_linear_model = _LinearModel(
- feature_columns,
- units,
- sparse_combiner,
- weight_collections,
- trainable,
- name='linear_model')
- retval = keras_linear_model(features) # pylint: disable=not-callable
- if cols_to_vars is not None:
- cols_to_vars.update(keras_linear_model.cols_to_vars())
- return retval
+ })
class LinearModelTest(test.TestCase):
@@ -1290,56 +1098,50 @@
def test_raises_if_empty_feature_columns(self):
with self.assertRaisesRegexp(ValueError,
'feature_columns must not be empty'):
- fc.linear_model(features={}, feature_columns=[])
+ fc.LinearModel(feature_columns=[])
def test_should_be_feature_column(self):
- with self.assertRaisesRegexp(ValueError, 'must be a _FeatureColumn'):
- fc.linear_model(features={'a': [[0]]}, feature_columns='NotSupported')
+ with self.assertRaisesRegexp(ValueError, 'must be a FeatureColumn'):
+ fc.LinearModel(feature_columns='NotSupported')
def test_should_be_dense_or_categorical_column(self):
- class NotSupportedColumn(fc_old._FeatureColumn):
+ class NotSupportedColumn(fc.FeatureColumn):
@property
def name(self):
return 'NotSupportedColumn'
- def _transform_feature(self, cache):
+ def transform_feature(self, transformation_cache, state_manager):
pass
@property
- def _parse_example_spec(self):
+ def parse_example_spec(self):
pass
with self.assertRaisesRegexp(
- ValueError, 'must be either a _DenseColumn or _CategoricalColumn'):
- fc.linear_model(
- features={'a': [[0]]}, feature_columns=[NotSupportedColumn()])
+ ValueError, 'must be either a DenseColumn or CategoricalColumn'):
+ fc.LinearModel(feature_columns=[NotSupportedColumn()])
def test_does_not_support_dict_columns(self):
with self.assertRaisesRegexp(
ValueError, 'Expected feature_columns to be iterable, found dict.'):
- fc.linear_model(
- features={'a': [[0]]},
- feature_columns={'a': fc_old.numeric_column('a')})
+ fc.LinearModel(feature_columns={'a': fc.numeric_column('a')})
def test_raises_if_duplicate_name(self):
with self.assertRaisesRegexp(
ValueError, 'Duplicate feature column name found for columns'):
- fc.linear_model(
- features={'a': [[0]]},
- feature_columns=[
- fc_old.numeric_column('a'),
- fc_old.numeric_column('a')
- ])
+ fc.LinearModel(
+ feature_columns=[fc.numeric_column('a'),
+ fc.numeric_column('a')])
def test_dense_bias(self):
- price = fc_old.numeric_column('price')
+ price = fc.numeric_column('price')
with ops.Graph().as_default():
features = {'price': [[1.], [5.]]}
- predictions = fc.linear_model(features, [price])
- bias = get_linear_model_bias()
- price_var = get_linear_model_column_var(price)
+ model = fc.LinearModel([price])
+ predictions = model(features)
+ price_var, bias = model.variables
with _initialized_session() as sess:
self.assertAllClose([0.], bias.eval())
sess.run(price_var.assign([[10.]]))
@@ -1347,16 +1149,16 @@
self.assertAllClose([[15.], [55.]], predictions.eval())
def test_sparse_bias(self):
- wire_cast = fc_old.categorical_column_with_hash_bucket('wire_cast', 4)
+ wire_cast = fc.categorical_column_with_hash_bucket('wire_cast', 4)
with ops.Graph().as_default():
wire_tensor = sparse_tensor.SparseTensor(
values=['omar', 'stringer', 'marlo'], # hashed to = [2, 0, 3]
indices=[[0, 0], [1, 0], [1, 1]],
dense_shape=[2, 2])
features = {'wire_cast': wire_tensor}
- predictions = fc.linear_model(features, [wire_cast])
- bias = get_linear_model_bias()
- wire_cast_var = get_linear_model_column_var(wire_cast)
+ model = fc.LinearModel([wire_cast])
+ predictions = model(features)
+ wire_cast_var, bias = model.variables
with _initialized_session() as sess:
self.assertAllClose([0.], bias.eval())
self.assertAllClose([[0.], [0.], [0.], [0.]], wire_cast_var.eval())
@@ -1365,18 +1167,17 @@
self.assertAllClose([[1005.], [10015.]], predictions.eval())
def test_dense_and_sparse_bias(self):
- wire_cast = fc_old.categorical_column_with_hash_bucket('wire_cast', 4)
- price = fc_old.numeric_column('price')
+ wire_cast = fc.categorical_column_with_hash_bucket('wire_cast', 4)
+ price = fc.numeric_column('price')
with ops.Graph().as_default():
wire_tensor = sparse_tensor.SparseTensor(
values=['omar', 'stringer', 'marlo'], # hashed to = [2, 0, 3]
indices=[[0, 0], [1, 0], [1, 1]],
dense_shape=[2, 2])
features = {'wire_cast': wire_tensor, 'price': [[1.], [5.]]}
- predictions = fc.linear_model(features, [wire_cast, price])
- bias = get_linear_model_bias()
- wire_cast_var = get_linear_model_column_var(wire_cast)
- price_var = get_linear_model_column_var(price)
+ model = fc.LinearModel([wire_cast, price])
+ predictions = model(features)
+ price_var, wire_cast_var, bias = model.variables
with _initialized_session() as sess:
sess.run(wire_cast_var.assign([[10.], [100.], [1000.], [10000.]]))
sess.run(bias.assign([5.]))
@@ -1386,38 +1187,36 @@
def test_dense_and_sparse_column(self):
"""When the column is both dense and sparse, uses sparse tensors."""
- class _DenseAndSparseColumn(fc_old._DenseColumn, fc_old._CategoricalColumn):
+ class _DenseAndSparseColumn(fc.DenseColumn, fc.CategoricalColumn):
@property
def name(self):
return 'dense_and_sparse_column'
@property
- def _parse_example_spec(self):
+ def parse_example_spec(self):
return {self.name: parsing_ops.VarLenFeature(self.dtype)}
- def _transform_feature(self, inputs):
- return inputs.get(self.name)
+ def transform_feature(self, transformation_cache, state_manager):
+ return transformation_cache.get(self.name, state_manager)
@property
- def _variable_shape(self):
+ def variable_shape(self):
raise ValueError('Should not use this method.')
- def _get_dense_tensor(self, inputs, weight_collections=None,
- trainable=None):
+ def get_dense_tensor(self, transformation_cache, state_manager):
raise ValueError('Should not use this method.')
@property
- def _num_buckets(self):
+ def num_buckets(self):
return 4
- def _get_sparse_tensors(self, inputs, weight_collections=None,
- trainable=None):
+ def get_sparse_tensors(self, transformation_cache, state_manager):
sp_tensor = sparse_tensor.SparseTensor(
indices=[[0, 0], [1, 0], [1, 1]],
values=[2, 0, 3],
dense_shape=[2, 2])
- return fc_old._CategoricalColumn.IdWeightPair(sp_tensor, None)
+ return fc.CategoricalColumn.IdWeightPair(sp_tensor, None)
dense_and_sparse_column = _DenseAndSparseColumn()
with ops.Graph().as_default():
@@ -1426,10 +1225,9 @@
indices=[[0, 0], [1, 0], [1, 1]],
dense_shape=[2, 2])
features = {dense_and_sparse_column.name: sp_tensor}
- predictions = fc.linear_model(features, [dense_and_sparse_column])
- bias = get_linear_model_bias()
- dense_and_sparse_column_var = get_linear_model_column_var(
- dense_and_sparse_column)
+ model = fc.LinearModel([dense_and_sparse_column])
+ predictions = model(features)
+ dense_and_sparse_column_var, bias = model.variables
with _initialized_session() as sess:
sess.run(dense_and_sparse_column_var.assign(
[[10.], [100.], [1000.], [10000.]]))
@@ -1437,12 +1235,12 @@
self.assertAllClose([[1005.], [10015.]], predictions.eval())
def test_dense_multi_output(self):
- price = fc_old.numeric_column('price')
+ price = fc.numeric_column('price')
with ops.Graph().as_default():
features = {'price': [[1.], [5.]]}
- predictions = fc.linear_model(features, [price], units=3)
- bias = get_linear_model_bias()
- price_var = get_linear_model_column_var(price)
+ model = fc.LinearModel([price], units=3)
+ predictions = model(features)
+ price_var, bias = model.variables
with _initialized_session() as sess:
self.assertAllClose(np.zeros((3,)), bias.eval())
self.assertAllClose(np.zeros((1, 3)), price_var.eval())
@@ -1452,16 +1250,16 @@
predictions.eval())
def test_sparse_multi_output(self):
- wire_cast = fc_old.categorical_column_with_hash_bucket('wire_cast', 4)
+ wire_cast = fc.categorical_column_with_hash_bucket('wire_cast', 4)
with ops.Graph().as_default():
wire_tensor = sparse_tensor.SparseTensor(
values=['omar', 'stringer', 'marlo'], # hashed to = [2, 0, 3]
indices=[[0, 0], [1, 0], [1, 1]],
dense_shape=[2, 2])
features = {'wire_cast': wire_tensor}
- predictions = fc.linear_model(features, [wire_cast], units=3)
- bias = get_linear_model_bias()
- wire_cast_var = get_linear_model_column_var(wire_cast)
+ model = fc.LinearModel([wire_cast], units=3)
+ predictions = model(features)
+ wire_cast_var, bias = model.variables
with _initialized_session() as sess:
self.assertAllClose(np.zeros((3,)), bias.eval())
self.assertAllClose(np.zeros((4, 3)), wire_cast_var.eval())
@@ -1474,18 +1272,19 @@
predictions.eval())
def test_dense_multi_dimension(self):
- price = fc_old.numeric_column('price', shape=2)
+ price = fc.numeric_column('price', shape=2)
with ops.Graph().as_default():
features = {'price': [[1., 2.], [5., 6.]]}
- predictions = fc.linear_model(features, [price])
- price_var = get_linear_model_column_var(price)
+ model = fc.LinearModel([price])
+ predictions = model(features)
+ price_var, _ = model.variables
with _initialized_session() as sess:
self.assertAllClose([[0.], [0.]], price_var.eval())
sess.run(price_var.assign([[10.], [100.]]))
self.assertAllClose([[210.], [650.]], predictions.eval())
def test_sparse_multi_rank(self):
- wire_cast = fc_old.categorical_column_with_hash_bucket('wire_cast', 4)
+ wire_cast = fc.categorical_column_with_hash_bucket('wire_cast', 4)
with ops.Graph().as_default():
wire_tensor = array_ops.sparse_placeholder(dtypes.string)
wire_value = sparse_tensor.SparseTensorValue(
@@ -1493,8 +1292,9 @@
indices=[[0, 0, 0], [0, 1, 0], [1, 0, 0], [1, 0, 1]],
dense_shape=[2, 2, 2])
features = {'wire_cast': wire_tensor}
- predictions = fc.linear_model(features, [wire_cast])
- wire_cast_var = get_linear_model_column_var(wire_cast)
+ model = fc.LinearModel([wire_cast])
+ predictions = model(features)
+ wire_cast_var, _ = model.variables
with _initialized_session() as sess:
self.assertAllClose(np.zeros((4, 1)), wire_cast_var.eval())
self.assertAllClose(
@@ -1506,25 +1306,24 @@
predictions.eval(feed_dict={wire_tensor: wire_value}))
def test_sparse_combiner(self):
- wire_cast = fc_old.categorical_column_with_hash_bucket('wire_cast', 4)
+ wire_cast = fc.categorical_column_with_hash_bucket('wire_cast', 4)
with ops.Graph().as_default():
wire_tensor = sparse_tensor.SparseTensor(
values=['omar', 'stringer', 'marlo'], # hashed to = [2, 0, 3]
indices=[[0, 0], [1, 0], [1, 1]],
dense_shape=[2, 2])
features = {'wire_cast': wire_tensor}
- predictions = fc.linear_model(
- features, [wire_cast], sparse_combiner='mean')
- bias = get_linear_model_bias()
- wire_cast_var = get_linear_model_column_var(wire_cast)
+ model = fc.LinearModel([wire_cast], sparse_combiner='mean')
+ predictions = model(features)
+ wire_cast_var, bias = model.variables
with _initialized_session() as sess:
sess.run(wire_cast_var.assign([[10.], [100.], [1000.], [10000.]]))
sess.run(bias.assign([5.]))
self.assertAllClose([[1005.], [5010.]], predictions.eval())
def test_sparse_combiner_with_negative_weights(self):
- wire_cast = fc_old.categorical_column_with_hash_bucket('wire_cast', 4)
- wire_cast_weights = fc_old.weighted_categorical_column(wire_cast, 'weights')
+ wire_cast = fc.categorical_column_with_hash_bucket('wire_cast', 4)
+ wire_cast_weights = fc.weighted_categorical_column(wire_cast, 'weights')
with ops.Graph().as_default():
wire_tensor = sparse_tensor.SparseTensor(
@@ -1535,22 +1334,21 @@
'wire_cast': wire_tensor,
'weights': constant_op.constant([[1., 1., -1.0]])
}
- predictions = fc.linear_model(
- features, [wire_cast_weights], sparse_combiner='sum')
- bias = get_linear_model_bias()
- wire_cast_var = get_linear_model_column_var(wire_cast)
+ model = fc.LinearModel([wire_cast_weights], sparse_combiner='sum')
+ predictions = model(features)
+ wire_cast_var, bias = model.variables
with _initialized_session() as sess:
sess.run(wire_cast_var.assign([[10.], [100.], [1000.], [10000.]]))
sess.run(bias.assign([5.]))
self.assertAllClose([[1005.], [-9985.]], predictions.eval())
def test_dense_multi_dimension_multi_output(self):
- price = fc_old.numeric_column('price', shape=2)
+ price = fc.numeric_column('price', shape=2)
with ops.Graph().as_default():
features = {'price': [[1., 2.], [5., 6.]]}
- predictions = fc.linear_model(features, [price], units=3)
- bias = get_linear_model_bias()
- price_var = get_linear_model_column_var(price)
+ model = fc.LinearModel([price], units=3)
+ predictions = model(features)
+ price_var, bias = model.variables
with _initialized_session() as sess:
self.assertAllClose(np.zeros((3,)), bias.eval())
self.assertAllClose(np.zeros((2, 3)), price_var.eval())
@@ -1560,21 +1358,22 @@
predictions.eval())
def test_raises_if_shape_mismatch(self):
- price = fc_old.numeric_column('price', shape=2)
+ price = fc.numeric_column('price', shape=2)
with ops.Graph().as_default():
features = {'price': [[1.], [5.]]}
with self.assertRaisesRegexp(
Exception,
r'Cannot reshape a tensor with 2 elements to shape \[2,2\]'):
- fc.linear_model(features, [price])
+ model = fc.LinearModel([price])
+ model(features)
def test_dense_reshaping(self):
- price = fc_old.numeric_column('price', shape=[1, 2])
+ price = fc.numeric_column('price', shape=[1, 2])
with ops.Graph().as_default():
features = {'price': [[[1., 2.]], [[5., 6.]]]}
- predictions = fc.linear_model(features, [price])
- bias = get_linear_model_bias()
- price_var = get_linear_model_column_var(price)
+ model = fc.LinearModel([price])
+ predictions = model(features)
+ price_var, bias = model.variables
with _initialized_session() as sess:
self.assertAllClose([0.], bias.eval())
self.assertAllClose([[0.], [0.]], price_var.eval())
@@ -1583,17 +1382,16 @@
self.assertAllClose([[210.], [650.]], predictions.eval())
def test_dense_multi_column(self):
- price1 = fc_old.numeric_column('price1', shape=2)
- price2 = fc_old.numeric_column('price2')
+ price1 = fc.numeric_column('price1', shape=2)
+ price2 = fc.numeric_column('price2')
with ops.Graph().as_default():
features = {
'price1': [[1., 2.], [5., 6.]],
'price2': [[3.], [4.]]
}
- predictions = fc.linear_model(features, [price1, price2])
- bias = get_linear_model_bias()
- price1_var = get_linear_model_column_var(price1)
- price2_var = get_linear_model_column_var(price2)
+ model = fc.LinearModel([price1, price2])
+ predictions = model(features)
+ price1_var, price2_var, bias = model.variables
with _initialized_session() as sess:
self.assertAllClose([0.], bias.eval())
self.assertAllClose([[0.], [0.]], price1_var.eval())
@@ -1604,115 +1402,55 @@
sess.run(bias.assign([7.]))
self.assertAllClose([[3217.], [4657.]], predictions.eval())
- def test_fills_cols_to_vars(self):
- price1 = fc_old.numeric_column('price1', shape=2)
- price2 = fc_old.numeric_column('price2')
- with ops.Graph().as_default():
- features = {'price1': [[1., 2.], [5., 6.]], 'price2': [[3.], [4.]]}
- cols_to_vars = {}
- fc.linear_model(features, [price1, price2], cols_to_vars=cols_to_vars)
- bias = get_linear_model_bias()
- price1_var = get_linear_model_column_var(price1)
- price2_var = get_linear_model_column_var(price2)
- self.assertAllEqual(cols_to_vars['bias'], [bias])
- self.assertAllEqual(cols_to_vars[price1], [price1_var])
- self.assertAllEqual(cols_to_vars[price2], [price2_var])
-
- def test_fills_cols_to_vars_partitioned_variables(self):
- price1 = fc_old.numeric_column('price1', shape=2)
- price2 = fc_old.numeric_column('price2', shape=3)
- with ops.Graph().as_default():
- features = {
- 'price1': [[1., 2.], [6., 7.]],
- 'price2': [[3., 4., 5.], [8., 9., 10.]]
- }
- cols_to_vars = {}
- with variable_scope.variable_scope(
- 'linear',
- partitioner=partitioned_variables.fixed_size_partitioner(2, axis=0)):
- fc.linear_model(features, [price1, price2], cols_to_vars=cols_to_vars)
- with _initialized_session():
- self.assertEqual([0.], cols_to_vars['bias'][0].eval())
- # Partitioning shards the [2, 1] price1 var into 2 [1, 1] Variables.
- self.assertAllEqual([[0.]], cols_to_vars[price1][0].eval())
- self.assertAllEqual([[0.]], cols_to_vars[price1][1].eval())
- # Partitioning shards the [3, 1] price2 var into a [2, 1] Variable and
- # a [1, 1] Variable.
- self.assertAllEqual([[0.], [0.]], cols_to_vars[price2][0].eval())
- self.assertAllEqual([[0.]], cols_to_vars[price2][1].eval())
-
- def test_dense_collection(self):
- price = fc_old.numeric_column('price')
- with ops.Graph().as_default() as g:
- features = {'price': [[1.], [5.]]}
- fc.linear_model(features, [price], weight_collections=['my-vars'])
- my_vars = g.get_collection('my-vars')
- bias = get_linear_model_bias()
- price_var = get_linear_model_column_var(price)
- self.assertIn(bias, my_vars)
- self.assertIn(price_var, my_vars)
-
- def test_sparse_collection(self):
- wire_cast = fc_old.categorical_column_with_hash_bucket('wire_cast', 4)
- with ops.Graph().as_default() as g:
- wire_tensor = sparse_tensor.SparseTensor(
- values=['omar'], indices=[[0, 0]], dense_shape=[1, 1])
- features = {'wire_cast': wire_tensor}
- fc.linear_model(
- features, [wire_cast], weight_collections=['my-vars'])
- my_vars = g.get_collection('my-vars')
- bias = get_linear_model_bias()
- wire_cast_var = get_linear_model_column_var(wire_cast)
- self.assertIn(bias, my_vars)
- self.assertIn(wire_cast_var, my_vars)
-
def test_dense_trainable_default(self):
- price = fc_old.numeric_column('price')
+ price = fc.numeric_column('price')
with ops.Graph().as_default() as g:
features = {'price': [[1.], [5.]]}
- fc.linear_model(features, [price])
- bias = get_linear_model_bias()
- price_var = get_linear_model_column_var(price)
+ model = fc.LinearModel([price])
+ model(features)
+ price_var, bias = model.variables
trainable_vars = g.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES)
self.assertIn(bias, trainable_vars)
self.assertIn(price_var, trainable_vars)
def test_sparse_trainable_default(self):
- wire_cast = fc_old.categorical_column_with_hash_bucket('wire_cast', 4)
+ wire_cast = fc.categorical_column_with_hash_bucket('wire_cast', 4)
with ops.Graph().as_default() as g:
wire_tensor = sparse_tensor.SparseTensor(
values=['omar'], indices=[[0, 0]], dense_shape=[1, 1])
features = {'wire_cast': wire_tensor}
- fc.linear_model(features, [wire_cast])
+ model = fc.LinearModel([wire_cast])
+ model(features)
trainable_vars = g.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES)
- bias = get_linear_model_bias()
- wire_cast_var = get_linear_model_column_var(wire_cast)
+ wire_cast_var, bias = model.variables
self.assertIn(bias, trainable_vars)
self.assertIn(wire_cast_var, trainable_vars)
def test_dense_trainable_false(self):
- price = fc_old.numeric_column('price')
+ price = fc.numeric_column('price')
with ops.Graph().as_default() as g:
features = {'price': [[1.], [5.]]}
- fc.linear_model(features, [price], trainable=False)
+ model = fc.LinearModel([price], trainable=False)
+ model(features)
trainable_vars = g.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES)
self.assertEqual([], trainable_vars)
def test_sparse_trainable_false(self):
- wire_cast = fc_old.categorical_column_with_hash_bucket('wire_cast', 4)
+ wire_cast = fc.categorical_column_with_hash_bucket('wire_cast', 4)
with ops.Graph().as_default() as g:
wire_tensor = sparse_tensor.SparseTensor(
values=['omar'], indices=[[0, 0]], dense_shape=[1, 1])
features = {'wire_cast': wire_tensor}
- fc.linear_model(features, [wire_cast], trainable=False)
+ model = fc.LinearModel([wire_cast], trainable=False)
+ model(features)
trainable_vars = g.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES)
self.assertEqual([], trainable_vars)
def test_column_order(self):
- price_a = fc_old.numeric_column('price_a')
- price_b = fc_old.numeric_column('price_b')
- wire_cast = fc_old.categorical_column_with_hash_bucket('wire_cast', 4)
- with ops.Graph().as_default() as g:
+ price_a = fc.numeric_column('price_a')
+ price_b = fc.numeric_column('price_b')
+ wire_cast = fc.categorical_column_with_hash_bucket('wire_cast', 4)
+ with ops.Graph().as_default():
features = {
'price_a': [[1.]],
'price_b': [[3.]],
@@ -1720,15 +1458,15 @@
sparse_tensor.SparseTensor(
values=['omar'], indices=[[0, 0]], dense_shape=[1, 1])
}
- fc.linear_model(
- features, [price_a, wire_cast, price_b],
- weight_collections=['my-vars'])
- my_vars = g.get_collection('my-vars')
+ model = fc.LinearModel([price_a, wire_cast, price_b])
+ model(features)
+
+ my_vars = model.variables
self.assertIn('price_a', my_vars[0].name)
self.assertIn('price_b', my_vars[1].name)
self.assertIn('wire_cast', my_vars[2].name)
- with ops.Graph().as_default() as g:
+ with ops.Graph().as_default():
features = {
'price_a': [[1.]],
'price_b': [[3.]],
@@ -1736,17 +1474,45 @@
sparse_tensor.SparseTensor(
values=['omar'], indices=[[0, 0]], dense_shape=[1, 1])
}
- fc.linear_model(
- features, [wire_cast, price_b, price_a],
- weight_collections=['my-vars'])
- my_vars = g.get_collection('my-vars')
+ model = fc.LinearModel([wire_cast, price_b, price_a])
+ model(features)
+
+ my_vars = model.variables
self.assertIn('price_a', my_vars[0].name)
self.assertIn('price_b', my_vars[1].name)
self.assertIn('wire_cast', my_vars[2].name)
+ def test_variable_names(self):
+ price1 = fc.numeric_column('price1')
+ dense_feature = fc.numeric_column('dense_feature')
+ dense_feature_bucketized = fc.bucketized_column(
+ dense_feature, boundaries=[0.])
+ some_sparse_column = fc.categorical_column_with_hash_bucket(
+ 'sparse_feature', hash_bucket_size=5)
+ some_embedding_column = fc.embedding_column(
+ some_sparse_column, dimension=10)
+ all_cols = [price1, dense_feature_bucketized, some_embedding_column]
+
+ with ops.Graph().as_default():
+ model = fc.LinearModel(all_cols)
+ features = {
+ 'price1': [[3.], [4.]],
+ 'dense_feature': [[-1.], [4.]],
+ 'sparse_feature': [['a'], ['x']],
+ }
+ model(features)
+ variable_names = [var.name for var in model.variables]
+ self.assertItemsEqual([
+ 'linear_model/dense_feature_bucketized/weights:0',
+ 'linear_model/price1/weights:0',
+ 'linear_model/sparse_feature_embedding/embedding_weights:0',
+ 'linear_model/sparse_feature_embedding/weights:0',
+ 'linear_model/bias_weights:0',
+ ], variable_names)
+
def test_static_batch_size_mismatch(self):
- price1 = fc_old.numeric_column('price1')
- price2 = fc_old.numeric_column('price2')
+ price1 = fc.numeric_column('price1')
+ price2 = fc.numeric_column('price2')
with ops.Graph().as_default():
features = {
'price1': [[1.], [5.], [7.]], # batchsize = 3
@@ -1755,12 +1521,13 @@
with self.assertRaisesRegexp(
ValueError,
'Batch size \(first dimension\) of each feature must be same.'): # pylint: disable=anomalous-backslash-in-string
- fc.linear_model(features, [price1, price2])
+ model = fc.LinearModel([price1, price2])
+ model(features)
def test_subset_of_static_batch_size_mismatch(self):
- price1 = fc_old.numeric_column('price1')
- price2 = fc_old.numeric_column('price2')
- price3 = fc_old.numeric_column('price3')
+ price1 = fc.numeric_column('price1')
+ price2 = fc.numeric_column('price2')
+ price3 = fc.numeric_column('price3')
with ops.Graph().as_default():
features = {
'price1': array_ops.placeholder(dtype=dtypes.int64), # batchsize = 3
@@ -1770,17 +1537,19 @@
with self.assertRaisesRegexp(
ValueError,
'Batch size \(first dimension\) of each feature must be same.'): # pylint: disable=anomalous-backslash-in-string
- fc.linear_model(features, [price1, price2, price3])
+ model = fc.LinearModel([price1, price2, price3])
+ model(features)
def test_runtime_batch_size_mismatch(self):
- price1 = fc_old.numeric_column('price1')
- price2 = fc_old.numeric_column('price2')
+ price1 = fc.numeric_column('price1')
+ price2 = fc.numeric_column('price2')
with ops.Graph().as_default():
features = {
'price1': array_ops.placeholder(dtype=dtypes.int64), # batchsize = 3
'price2': [[3.], [4.]] # batchsize = 2
}
- predictions = fc.linear_model(features, [price1, price2])
+ model = fc.LinearModel([price1, price2])
+ predictions = model(features)
with _initialized_session() as sess:
with self.assertRaisesRegexp(errors.OpError,
'must have the same size and shape'):
@@ -1788,14 +1557,15 @@
predictions, feed_dict={features['price1']: [[1.], [5.], [7.]]})
def test_runtime_batch_size_matches(self):
- price1 = fc_old.numeric_column('price1')
- price2 = fc_old.numeric_column('price2')
+ price1 = fc.numeric_column('price1')
+ price2 = fc.numeric_column('price2')
with ops.Graph().as_default():
features = {
'price1': array_ops.placeholder(dtype=dtypes.int64), # batchsize = 2
'price2': array_ops.placeholder(dtype=dtypes.int64), # batchsize = 2
}
- predictions = fc.linear_model(features, [price1, price2])
+ model = fc.LinearModel([price1, price2])
+ predictions = model(features)
with _initialized_session() as sess:
sess.run(
predictions,
@@ -1805,14 +1575,14 @@
})
def test_with_numpy_input_fn(self):
- price = fc_old.numeric_column('price')
- price_buckets = fc_old.bucketized_column(
+ price = fc.numeric_column('price')
+ price_buckets = fc.bucketized_column(
price, boundaries=[
0.,
10.,
100.,
])
- body_style = fc_old.categorical_column_with_vocabulary_list(
+ body_style = fc.categorical_column_with_vocabulary_list(
'body-style', vocabulary_list=['hardtop', 'wagon', 'sedan'])
input_fn = numpy_io.numpy_input_fn(
@@ -1823,15 +1593,14 @@
batch_size=2,
shuffle=False)
features = input_fn()
- net = fc.linear_model(features, [price_buckets, body_style])
+ model = fc.LinearModel([price_buckets, body_style])
+ net = model(features)
# self.assertEqual(1 + 3 + 5, net.shape[1])
with _initialized_session() as sess:
coord = coordinator.Coordinator()
threads = queue_runner_impl.start_queue_runners(sess, coord=coord)
- bias = get_linear_model_bias()
- price_buckets_var = get_linear_model_column_var(price_buckets)
- body_style_var = get_linear_model_column_var(body_style)
+ body_style_var, price_buckets_var, bias = model.variables
sess.run(price_buckets_var.assign([[10.], [100.], [1000.], [10000.]]))
sess.run(body_style_var.assign([[-10.], [-100.], [-1000.]]))
@@ -1843,14 +1612,14 @@
coord.join(threads)
def test_with_1d_sparse_tensor(self):
- price = fc_old.numeric_column('price')
- price_buckets = fc_old.bucketized_column(
+ price = fc.numeric_column('price')
+ price_buckets = fc.bucketized_column(
price, boundaries=[
0.,
10.,
100.,
])
- body_style = fc_old.categorical_column_with_vocabulary_list(
+ body_style = fc.categorical_column_with_vocabulary_list(
'body-style', vocabulary_list=['hardtop', 'wagon', 'sedan'])
# Provides 1-dim tensor and dense tensor.
@@ -1864,11 +1633,10 @@
self.assertEqual(1, features['price'].shape.ndims)
self.assertEqual(1, features['body-style'].dense_shape.get_shape()[0])
- net = fc.linear_model(features, [price_buckets, body_style])
+ model = fc.LinearModel([price_buckets, body_style])
+ net = model(features)
with _initialized_session() as sess:
- bias = get_linear_model_bias()
- price_buckets_var = get_linear_model_column_var(price_buckets)
- body_style_var = get_linear_model_column_var(body_style)
+ body_style_var, price_buckets_var, bias = model.variables
sess.run(price_buckets_var.assign([[10.], [100.], [1000.], [10000.]]))
sess.run(body_style_var.assign([[-10.], [-100.], [-1000.]]))
@@ -1877,16 +1645,16 @@
self.assertAllClose([[10 - 1000 + 5.], [1000 - 10 + 5.]], sess.run(net))
def test_with_1d_unknown_shape_sparse_tensor(self):
- price = fc_old.numeric_column('price')
- price_buckets = fc_old.bucketized_column(
+ price = fc.numeric_column('price')
+ price_buckets = fc.bucketized_column(
price, boundaries=[
0.,
10.,
100.,
])
- body_style = fc_old.categorical_column_with_vocabulary_list(
+ body_style = fc.categorical_column_with_vocabulary_list(
'body-style', vocabulary_list=['hardtop', 'wagon', 'sedan'])
- country = fc_old.categorical_column_with_vocabulary_list(
+ country = fc.categorical_column_with_vocabulary_list(
'country', vocabulary_list=['US', 'JP', 'CA'])
# Provides 1-dim tensor and dense tensor.
@@ -1905,10 +1673,9 @@
dense_shape=(2,))
country_data = np.array(['US', 'CA'])
- net = fc.linear_model(features, [price_buckets, body_style, country])
- bias = get_linear_model_bias()
- price_buckets_var = get_linear_model_column_var(price_buckets)
- body_style_var = get_linear_model_column_var(body_style)
+ model = fc.LinearModel([price_buckets, body_style, country])
+ net = model(features)
+ body_style_var, _, price_buckets_var, bias = model.variables
with _initialized_session() as sess:
sess.run(price_buckets_var.assign([[10.], [100.], [1000.], [10000.]]))
sess.run(body_style_var.assign([[-10.], [-100.], [-1000.]]))
@@ -1924,7 +1691,7 @@
}))
def test_with_rank_0_feature(self):
- price = fc_old.numeric_column('price')
+ price = fc.numeric_column('price')
features = {
'price': constant_op.constant(0),
}
@@ -1932,29 +1699,31 @@
# Static rank 0 should fail
with self.assertRaisesRegexp(ValueError, 'Feature .* cannot have rank 0'):
- fc.linear_model(features, [price])
+ model = fc.LinearModel([price])
+ model(features)
# Dynamic rank 0 should fail
features = {
'price': array_ops.placeholder(dtypes.float32),
}
- net = fc.linear_model(features, [price])
+ model = fc.LinearModel([price])
+ net = model(features)
self.assertEqual(1, net.shape[1])
with _initialized_session() as sess:
with self.assertRaisesOpError('Feature .* cannot have rank 0'):
sess.run(net, feed_dict={features['price']: np.array(1)})
def test_multiple_linear_models(self):
- price = fc_old.numeric_column('price')
+ price = fc.numeric_column('price')
with ops.Graph().as_default():
features1 = {'price': [[1.], [5.]]}
features2 = {'price': [[2.], [10.]]}
- predictions1 = fc.linear_model(features1, [price])
- predictions2 = fc.linear_model(features2, [price])
- bias1 = get_linear_model_bias(name='linear_model')
- bias2 = get_linear_model_bias(name='linear_model_1')
- price_var1 = get_linear_model_column_var(price, name='linear_model')
- price_var2 = get_linear_model_column_var(price, name='linear_model_1')
+ model1 = fc.LinearModel([price])
+ model2 = fc.LinearModel([price])
+ predictions1 = model1(features1)
+ predictions2 = model2(features2)
+ price_var1, bias1 = model1.variables
+ price_var2, bias2 = model2.variables
with _initialized_session() as sess:
self.assertAllClose([0.], bias1.eval())
sess.run(price_var1.assign([[10.]]))
@@ -1966,664 +1735,6 @@
self.assertAllClose([[25.], [105.]], predictions2.eval())
-class _LinearModelTest(test.TestCase):
-
- def test_raises_if_empty_feature_columns(self):
- with self.assertRaisesRegexp(ValueError,
- 'feature_columns must not be empty'):
- get_keras_linear_model_predictions(features={}, feature_columns=[])
-
- def test_should_be_feature_column(self):
- with self.assertRaisesRegexp(ValueError, 'must be a _FeatureColumn'):
- get_keras_linear_model_predictions(
- features={'a': [[0]]}, feature_columns='NotSupported')
-
- def test_should_be_dense_or_categorical_column(self):
-
- class NotSupportedColumn(fc_old._FeatureColumn):
-
- @property
- def name(self):
- return 'NotSupportedColumn'
-
- def _transform_feature(self, cache):
- pass
-
- @property
- def _parse_example_spec(self):
- pass
-
- with self.assertRaisesRegexp(
- ValueError, 'must be either a _DenseColumn or _CategoricalColumn'):
- get_keras_linear_model_predictions(
- features={'a': [[0]]}, feature_columns=[NotSupportedColumn()])
-
- def test_does_not_support_dict_columns(self):
- with self.assertRaisesRegexp(
- ValueError, 'Expected feature_columns to be iterable, found dict.'):
- fc.linear_model(
- features={'a': [[0]]},
- feature_columns={'a': fc_old.numeric_column('a')})
-
- def test_raises_if_duplicate_name(self):
- with self.assertRaisesRegexp(
- ValueError, 'Duplicate feature column name found for columns'):
- get_keras_linear_model_predictions(
- features={'a': [[0]]},
- feature_columns=[
- fc_old.numeric_column('a'),
- fc_old.numeric_column('a')
- ])
-
- def test_dense_bias(self):
- price = fc_old.numeric_column('price')
- with ops.Graph().as_default():
- features = {'price': [[1.], [5.]]}
- predictions = get_keras_linear_model_predictions(features, [price])
- bias = get_linear_model_bias()
- price_var = get_linear_model_column_var(price)
- with _initialized_session() as sess:
- self.assertAllClose([0.], bias.eval())
- sess.run(price_var.assign([[10.]]))
- sess.run(bias.assign([5.]))
- self.assertAllClose([[15.], [55.]], predictions.eval())
-
- def test_sparse_bias(self):
- wire_cast = fc_old.categorical_column_with_hash_bucket('wire_cast', 4)
- with ops.Graph().as_default():
- wire_tensor = sparse_tensor.SparseTensor(
- values=['omar', 'stringer', 'marlo'], # hashed to = [2, 0, 3]
- indices=[[0, 0], [1, 0], [1, 1]],
- dense_shape=[2, 2])
- features = {'wire_cast': wire_tensor}
- predictions = get_keras_linear_model_predictions(features, [wire_cast])
- bias = get_linear_model_bias()
- wire_cast_var = get_linear_model_column_var(wire_cast)
- with _initialized_session() as sess:
- self.assertAllClose([0.], bias.eval())
- self.assertAllClose([[0.], [0.], [0.], [0.]], wire_cast_var.eval())
- sess.run(wire_cast_var.assign([[10.], [100.], [1000.], [10000.]]))
- sess.run(bias.assign([5.]))
- self.assertAllClose([[1005.], [10015.]], predictions.eval())
-
- def test_dense_and_sparse_bias(self):
- wire_cast = fc_old.categorical_column_with_hash_bucket('wire_cast', 4)
- price = fc_old.numeric_column('price')
- with ops.Graph().as_default():
- wire_tensor = sparse_tensor.SparseTensor(
- values=['omar', 'stringer', 'marlo'], # hashed to = [2, 0, 3]
- indices=[[0, 0], [1, 0], [1, 1]],
- dense_shape=[2, 2])
- features = {'wire_cast': wire_tensor, 'price': [[1.], [5.]]}
- predictions = get_keras_linear_model_predictions(features,
- [wire_cast, price])
- bias = get_linear_model_bias()
- wire_cast_var = get_linear_model_column_var(wire_cast)
- price_var = get_linear_model_column_var(price)
- with _initialized_session() as sess:
- sess.run(wire_cast_var.assign([[10.], [100.], [1000.], [10000.]]))
- sess.run(bias.assign([5.]))
- sess.run(price_var.assign([[10.]]))
- self.assertAllClose([[1015.], [10065.]], predictions.eval())
-
- def test_dense_and_sparse_column(self):
- """When the column is both dense and sparse, uses sparse tensors."""
-
- class _DenseAndSparseColumn(fc_old._DenseColumn, fc_old._CategoricalColumn):
-
- @property
- def name(self):
- return 'dense_and_sparse_column'
-
- @property
- def _parse_example_spec(self):
- return {self.name: parsing_ops.VarLenFeature(self.dtype)}
-
- def _transform_feature(self, inputs):
- return inputs.get(self.name)
-
- @property
- def _variable_shape(self):
- raise ValueError('Should not use this method.')
-
- def _get_dense_tensor(self,
- inputs,
- weight_collections=None,
- trainable=None):
- raise ValueError('Should not use this method.')
-
- @property
- def _num_buckets(self):
- return 4
-
- def _get_sparse_tensors(self,
- inputs,
- weight_collections=None,
- trainable=None):
- sp_tensor = sparse_tensor.SparseTensor(
- indices=[[0, 0], [1, 0], [1, 1]],
- values=[2, 0, 3],
- dense_shape=[2, 2])
- return fc_old._CategoricalColumn.IdWeightPair(sp_tensor, None)
-
- dense_and_sparse_column = _DenseAndSparseColumn()
- with ops.Graph().as_default():
- sp_tensor = sparse_tensor.SparseTensor(
- values=['omar', 'stringer', 'marlo'],
- indices=[[0, 0], [1, 0], [1, 1]],
- dense_shape=[2, 2])
- features = {dense_and_sparse_column.name: sp_tensor}
- predictions = get_keras_linear_model_predictions(
- features, [dense_and_sparse_column])
- bias = get_linear_model_bias()
- dense_and_sparse_column_var = get_linear_model_column_var(
- dense_and_sparse_column)
- with _initialized_session() as sess:
- sess.run(
- dense_and_sparse_column_var.assign([[10.], [100.], [1000.],
- [10000.]]))
- sess.run(bias.assign([5.]))
- self.assertAllClose([[1005.], [10015.]], predictions.eval())
-
- def test_dense_multi_output(self):
- price = fc_old.numeric_column('price')
- with ops.Graph().as_default():
- features = {'price': [[1.], [5.]]}
- predictions = get_keras_linear_model_predictions(
- features, [price], units=3)
- bias = get_linear_model_bias()
- price_var = get_linear_model_column_var(price)
- with _initialized_session() as sess:
- self.assertAllClose(np.zeros((3,)), bias.eval())
- self.assertAllClose(np.zeros((1, 3)), price_var.eval())
- sess.run(price_var.assign([[10., 100., 1000.]]))
- sess.run(bias.assign([5., 6., 7.]))
- self.assertAllClose([[15., 106., 1007.], [55., 506., 5007.]],
- predictions.eval())
-
- def test_sparse_multi_output(self):
- wire_cast = fc_old.categorical_column_with_hash_bucket('wire_cast', 4)
- with ops.Graph().as_default():
- wire_tensor = sparse_tensor.SparseTensor(
- values=['omar', 'stringer', 'marlo'], # hashed to = [2, 0, 3]
- indices=[[0, 0], [1, 0], [1, 1]],
- dense_shape=[2, 2])
- features = {'wire_cast': wire_tensor}
- predictions = get_keras_linear_model_predictions(
- features, [wire_cast], units=3)
- bias = get_linear_model_bias()
- wire_cast_var = get_linear_model_column_var(wire_cast)
- with _initialized_session() as sess:
- self.assertAllClose(np.zeros((3,)), bias.eval())
- self.assertAllClose(np.zeros((4, 3)), wire_cast_var.eval())
- sess.run(
- wire_cast_var.assign([[10., 11., 12.], [100., 110., 120.],
- [1000., 1100.,
- 1200.], [10000., 11000., 12000.]]))
- sess.run(bias.assign([5., 6., 7.]))
- self.assertAllClose([[1005., 1106., 1207.], [10015., 11017., 12019.]],
- predictions.eval())
-
- def test_dense_multi_dimension(self):
- price = fc_old.numeric_column('price', shape=2)
- with ops.Graph().as_default():
- features = {'price': [[1., 2.], [5., 6.]]}
- predictions = get_keras_linear_model_predictions(features, [price])
- price_var = get_linear_model_column_var(price)
- with _initialized_session() as sess:
- self.assertAllClose([[0.], [0.]], price_var.eval())
- sess.run(price_var.assign([[10.], [100.]]))
- self.assertAllClose([[210.], [650.]], predictions.eval())
-
- def test_sparse_multi_rank(self):
- wire_cast = fc_old.categorical_column_with_hash_bucket('wire_cast', 4)
- with ops.Graph().as_default():
- wire_tensor = array_ops.sparse_placeholder(dtypes.string)
- wire_value = sparse_tensor.SparseTensorValue(
- values=['omar', 'stringer', 'marlo', 'omar'], # hashed = [2, 0, 3, 2]
- indices=[[0, 0, 0], [0, 1, 0], [1, 0, 0], [1, 0, 1]],
- dense_shape=[2, 2, 2])
- features = {'wire_cast': wire_tensor}
- predictions = get_keras_linear_model_predictions(features, [wire_cast])
- wire_cast_var = get_linear_model_column_var(wire_cast)
- with _initialized_session() as sess:
- self.assertAllClose(np.zeros((4, 1)), wire_cast_var.eval())
- self.assertAllClose(
- np.zeros((2, 1)),
- predictions.eval(feed_dict={wire_tensor: wire_value}))
- sess.run(wire_cast_var.assign([[10.], [100.], [1000.], [10000.]]))
- self.assertAllClose(
- [[1010.], [11000.]],
- predictions.eval(feed_dict={wire_tensor: wire_value}))
-
- def test_sparse_combiner(self):
- wire_cast = fc_old.categorical_column_with_hash_bucket('wire_cast', 4)
- with ops.Graph().as_default():
- wire_tensor = sparse_tensor.SparseTensor(
- values=['omar', 'stringer', 'marlo'], # hashed to = [2, 0, 3]
- indices=[[0, 0], [1, 0], [1, 1]],
- dense_shape=[2, 2])
- features = {'wire_cast': wire_tensor}
- predictions = get_keras_linear_model_predictions(
- features, [wire_cast], sparse_combiner='mean')
- bias = get_linear_model_bias()
- wire_cast_var = get_linear_model_column_var(wire_cast)
- with _initialized_session() as sess:
- sess.run(wire_cast_var.assign([[10.], [100.], [1000.], [10000.]]))
- sess.run(bias.assign([5.]))
- self.assertAllClose([[1005.], [5010.]], predictions.eval())
-
- def test_dense_multi_dimension_multi_output(self):
- price = fc_old.numeric_column('price', shape=2)
- with ops.Graph().as_default():
- features = {'price': [[1., 2.], [5., 6.]]}
- predictions = get_keras_linear_model_predictions(
- features, [price], units=3)
- bias = get_linear_model_bias()
- price_var = get_linear_model_column_var(price)
- with _initialized_session() as sess:
- self.assertAllClose(np.zeros((3,)), bias.eval())
- self.assertAllClose(np.zeros((2, 3)), price_var.eval())
- sess.run(price_var.assign([[1., 2., 3.], [10., 100., 1000.]]))
- sess.run(bias.assign([2., 3., 4.]))
- self.assertAllClose([[23., 205., 2007.], [67., 613., 6019.]],
- predictions.eval())
-
- def test_raises_if_shape_mismatch(self):
- price = fc_old.numeric_column('price', shape=2)
- with ops.Graph().as_default():
- features = {'price': [[1.], [5.]]}
- with self.assertRaisesRegexp(
- Exception,
- r'Cannot reshape a tensor with 2 elements to shape \[2,2\]'):
- get_keras_linear_model_predictions(features, [price])
-
- def test_dense_reshaping(self):
- price = fc_old.numeric_column('price', shape=[1, 2])
- with ops.Graph().as_default():
- features = {'price': [[[1., 2.]], [[5., 6.]]]}
- predictions = get_keras_linear_model_predictions(features, [price])
- bias = get_linear_model_bias()
- price_var = get_linear_model_column_var(price)
- with _initialized_session() as sess:
- self.assertAllClose([0.], bias.eval())
- self.assertAllClose([[0.], [0.]], price_var.eval())
- self.assertAllClose([[0.], [0.]], predictions.eval())
- sess.run(price_var.assign([[10.], [100.]]))
- self.assertAllClose([[210.], [650.]], predictions.eval())
-
- def test_dense_multi_column(self):
- price1 = fc_old.numeric_column('price1', shape=2)
- price2 = fc_old.numeric_column('price2')
- with ops.Graph().as_default():
- features = {'price1': [[1., 2.], [5., 6.]], 'price2': [[3.], [4.]]}
- predictions = get_keras_linear_model_predictions(features,
- [price1, price2])
- bias = get_linear_model_bias()
- price1_var = get_linear_model_column_var(price1)
- price2_var = get_linear_model_column_var(price2)
- with _initialized_session() as sess:
- self.assertAllClose([0.], bias.eval())
- self.assertAllClose([[0.], [0.]], price1_var.eval())
- self.assertAllClose([[0.]], price2_var.eval())
- self.assertAllClose([[0.], [0.]], predictions.eval())
- sess.run(price1_var.assign([[10.], [100.]]))
- sess.run(price2_var.assign([[1000.]]))
- sess.run(bias.assign([7.]))
- self.assertAllClose([[3217.], [4657.]], predictions.eval())
-
- def test_fills_cols_to_vars(self):
- price1 = fc_old.numeric_column('price1', shape=2)
- price2 = fc_old.numeric_column('price2')
- with ops.Graph().as_default():
- features = {'price1': [[1., 2.], [5., 6.]], 'price2': [[3.], [4.]]}
- cols_to_vars = {}
- get_keras_linear_model_predictions(
- features, [price1, price2], cols_to_vars=cols_to_vars)
- bias = get_linear_model_bias()
- price1_var = get_linear_model_column_var(price1)
- price2_var = get_linear_model_column_var(price2)
- self.assertAllEqual(cols_to_vars['bias'], [bias])
- self.assertAllEqual(cols_to_vars[price1], [price1_var])
- self.assertAllEqual(cols_to_vars[price2], [price2_var])
-
- def test_fills_cols_to_vars_partitioned_variables(self):
- price1 = fc_old.numeric_column('price1', shape=2)
- price2 = fc_old.numeric_column('price2', shape=3)
- with ops.Graph().as_default():
- features = {
- 'price1': [[1., 2.], [6., 7.]],
- 'price2': [[3., 4., 5.], [8., 9., 10.]]
- }
- cols_to_vars = {}
- with variable_scope.variable_scope(
- 'linear',
- partitioner=partitioned_variables.fixed_size_partitioner(2, axis=0)):
- get_keras_linear_model_predictions(
- features, [price1, price2], cols_to_vars=cols_to_vars)
- with _initialized_session():
- self.assertEqual([0.], cols_to_vars['bias'][0].eval())
- # Partitioning shards the [2, 1] price1 var into 2 [1, 1] Variables.
- self.assertAllEqual([[0.]], cols_to_vars[price1][0].eval())
- self.assertAllEqual([[0.]], cols_to_vars[price1][1].eval())
- # Partitioning shards the [3, 1] price2 var into a [2, 1] Variable and
- # a [1, 1] Variable.
- self.assertAllEqual([[0.], [0.]], cols_to_vars[price2][0].eval())
- self.assertAllEqual([[0.]], cols_to_vars[price2][1].eval())
-
- def test_dense_collection(self):
- price = fc_old.numeric_column('price')
- with ops.Graph().as_default() as g:
- features = {'price': [[1.], [5.]]}
- get_keras_linear_model_predictions(
- features, [price], weight_collections=['my-vars'])
- my_vars = g.get_collection('my-vars')
- bias = get_linear_model_bias()
- price_var = get_linear_model_column_var(price)
- self.assertIn(bias, my_vars)
- self.assertIn(price_var, my_vars)
-
- def test_sparse_collection(self):
- wire_cast = fc_old.categorical_column_with_hash_bucket('wire_cast', 4)
- with ops.Graph().as_default() as g:
- wire_tensor = sparse_tensor.SparseTensor(
- values=['omar'], indices=[[0, 0]], dense_shape=[1, 1])
- features = {'wire_cast': wire_tensor}
- get_keras_linear_model_predictions(
- features, [wire_cast], weight_collections=['my-vars'])
- my_vars = g.get_collection('my-vars')
- bias = get_linear_model_bias()
- wire_cast_var = get_linear_model_column_var(wire_cast)
- self.assertIn(bias, my_vars)
- self.assertIn(wire_cast_var, my_vars)
-
- def test_dense_trainable_default(self):
- price = fc_old.numeric_column('price')
- with ops.Graph().as_default() as g:
- features = {'price': [[1.], [5.]]}
- get_keras_linear_model_predictions(features, [price])
- bias = get_linear_model_bias()
- price_var = get_linear_model_column_var(price)
- trainable_vars = g.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES)
- self.assertIn(bias, trainable_vars)
- self.assertIn(price_var, trainable_vars)
-
- def test_sparse_trainable_default(self):
- wire_cast = fc_old.categorical_column_with_hash_bucket('wire_cast', 4)
- with ops.Graph().as_default() as g:
- wire_tensor = sparse_tensor.SparseTensor(
- values=['omar'], indices=[[0, 0]], dense_shape=[1, 1])
- features = {'wire_cast': wire_tensor}
- get_keras_linear_model_predictions(features, [wire_cast])
- trainable_vars = g.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES)
- bias = get_linear_model_bias()
- wire_cast_var = get_linear_model_column_var(wire_cast)
- self.assertIn(bias, trainable_vars)
- self.assertIn(wire_cast_var, trainable_vars)
-
- def test_dense_trainable_false(self):
- price = fc_old.numeric_column('price')
- with ops.Graph().as_default() as g:
- features = {'price': [[1.], [5.]]}
- get_keras_linear_model_predictions(features, [price], trainable=False)
- trainable_vars = g.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES)
- self.assertEqual([], trainable_vars)
-
- def test_sparse_trainable_false(self):
- wire_cast = fc_old.categorical_column_with_hash_bucket('wire_cast', 4)
- with ops.Graph().as_default() as g:
- wire_tensor = sparse_tensor.SparseTensor(
- values=['omar'], indices=[[0, 0]], dense_shape=[1, 1])
- features = {'wire_cast': wire_tensor}
- get_keras_linear_model_predictions(features, [wire_cast], trainable=False)
- trainable_vars = g.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES)
- self.assertEqual([], trainable_vars)
-
- def test_column_order(self):
- price_a = fc_old.numeric_column('price_a')
- price_b = fc_old.numeric_column('price_b')
- wire_cast = fc_old.categorical_column_with_hash_bucket('wire_cast', 4)
- with ops.Graph().as_default() as g:
- features = {
- 'price_a': [[1.]],
- 'price_b': [[3.]],
- 'wire_cast':
- sparse_tensor.SparseTensor(
- values=['omar'], indices=[[0, 0]], dense_shape=[1, 1])
- }
- get_keras_linear_model_predictions(
- features, [price_a, wire_cast, price_b],
- weight_collections=['my-vars'])
- my_vars = g.get_collection('my-vars')
- self.assertIn('price_a', my_vars[0].name)
- self.assertIn('price_b', my_vars[1].name)
- self.assertIn('wire_cast', my_vars[2].name)
-
- with ops.Graph().as_default() as g:
- features = {
- 'price_a': [[1.]],
- 'price_b': [[3.]],
- 'wire_cast':
- sparse_tensor.SparseTensor(
- values=['omar'], indices=[[0, 0]], dense_shape=[1, 1])
- }
- get_keras_linear_model_predictions(
- features, [wire_cast, price_b, price_a],
- weight_collections=['my-vars'])
- my_vars = g.get_collection('my-vars')
- self.assertIn('price_a', my_vars[0].name)
- self.assertIn('price_b', my_vars[1].name)
- self.assertIn('wire_cast', my_vars[2].name)
-
- def test_static_batch_size_mismatch(self):
- price1 = fc_old.numeric_column('price1')
- price2 = fc_old.numeric_column('price2')
- with ops.Graph().as_default():
- features = {
- 'price1': [[1.], [5.], [7.]], # batchsize = 3
- 'price2': [[3.], [4.]] # batchsize = 2
- }
- with self.assertRaisesRegexp(
- ValueError,
- 'Batch size \(first dimension\) of each feature must be same.'): # pylint: disable=anomalous-backslash-in-string
- get_keras_linear_model_predictions(features, [price1, price2])
-
- def test_subset_of_static_batch_size_mismatch(self):
- price1 = fc_old.numeric_column('price1')
- price2 = fc_old.numeric_column('price2')
- price3 = fc_old.numeric_column('price3')
- with ops.Graph().as_default():
- features = {
- 'price1': array_ops.placeholder(dtype=dtypes.int64), # batchsize = 3
- 'price2': [[3.], [4.]], # batchsize = 2
- 'price3': [[3.], [4.], [5.]] # batchsize = 3
- }
- with self.assertRaisesRegexp(
- ValueError,
- 'Batch size \(first dimension\) of each feature must be same.'): # pylint: disable=anomalous-backslash-in-string
- get_keras_linear_model_predictions(features, [price1, price2, price3])
-
- def test_runtime_batch_size_mismatch(self):
- price1 = fc_old.numeric_column('price1')
- price2 = fc_old.numeric_column('price2')
- with ops.Graph().as_default():
- features = {
- 'price1': array_ops.placeholder(dtype=dtypes.int64), # batchsize = 3
- 'price2': [[3.], [4.]] # batchsize = 2
- }
- predictions = get_keras_linear_model_predictions(features,
- [price1, price2])
- with _initialized_session() as sess:
- with self.assertRaisesRegexp(errors.OpError,
- 'must have the same size and shape'):
- sess.run(
- predictions, feed_dict={features['price1']: [[1.], [5.], [7.]]})
-
- def test_runtime_batch_size_matches(self):
- price1 = fc_old.numeric_column('price1')
- price2 = fc_old.numeric_column('price2')
- with ops.Graph().as_default():
- features = {
- 'price1': array_ops.placeholder(dtype=dtypes.int64), # batchsize = 2
- 'price2': array_ops.placeholder(dtype=dtypes.int64), # batchsize = 2
- }
- predictions = get_keras_linear_model_predictions(features,
- [price1, price2])
- with _initialized_session() as sess:
- sess.run(
- predictions,
- feed_dict={
- features['price1']: [[1.], [5.]],
- features['price2']: [[1.], [5.]],
- })
-
- def test_with_numpy_input_fn(self):
- price = fc_old.numeric_column('price')
- price_buckets = fc_old.bucketized_column(
- price, boundaries=[
- 0.,
- 10.,
- 100.,
- ])
- body_style = fc_old.categorical_column_with_vocabulary_list(
- 'body-style', vocabulary_list=['hardtop', 'wagon', 'sedan'])
-
- input_fn = numpy_io.numpy_input_fn(
- x={
- 'price': np.array([-1., 2., 13., 104.]),
- 'body-style': np.array(['sedan', 'hardtop', 'wagon', 'sedan']),
- },
- batch_size=2,
- shuffle=False)
- features = input_fn()
- net = get_keras_linear_model_predictions(features,
- [price_buckets, body_style])
- # self.assertEqual(1 + 3 + 5, net.shape[1])
- with _initialized_session() as sess:
- coord = coordinator.Coordinator()
- threads = queue_runner_impl.start_queue_runners(sess, coord=coord)
-
- bias = get_linear_model_bias()
- price_buckets_var = get_linear_model_column_var(price_buckets)
- body_style_var = get_linear_model_column_var(body_style)
-
- sess.run(price_buckets_var.assign([[10.], [100.], [1000.], [10000.]]))
- sess.run(body_style_var.assign([[-10.], [-100.], [-1000.]]))
- sess.run(bias.assign([5.]))
-
- self.assertAllClose([[10 - 1000 + 5.], [100 - 10 + 5.]], sess.run(net))
-
- coord.request_stop()
- coord.join(threads)
-
- def test_with_1d_sparse_tensor(self):
- price = fc_old.numeric_column('price')
- price_buckets = fc_old.bucketized_column(
- price, boundaries=[
- 0.,
- 10.,
- 100.,
- ])
- body_style = fc_old.categorical_column_with_vocabulary_list(
- 'body-style', vocabulary_list=['hardtop', 'wagon', 'sedan'])
-
- # Provides 1-dim tensor and dense tensor.
- features = {
- 'price':
- constant_op.constant([
- -1.,
- 12.,
- ]),
- 'body-style':
- sparse_tensor.SparseTensor(
- indices=((0,), (1,)),
- values=('sedan', 'hardtop'),
- dense_shape=(2,)),
- }
- self.assertEqual(1, features['price'].shape.ndims)
- self.assertEqual(1, features['body-style'].dense_shape.get_shape()[0])
-
- net = get_keras_linear_model_predictions(features,
- [price_buckets, body_style])
- with _initialized_session() as sess:
- bias = get_linear_model_bias()
- price_buckets_var = get_linear_model_column_var(price_buckets)
- body_style_var = get_linear_model_column_var(body_style)
-
- sess.run(price_buckets_var.assign([[10.], [100.], [1000.], [10000.]]))
- sess.run(body_style_var.assign([[-10.], [-100.], [-1000.]]))
- sess.run(bias.assign([5.]))
-
- self.assertAllClose([[10 - 1000 + 5.], [1000 - 10 + 5.]], sess.run(net))
-
- def test_with_1d_unknown_shape_sparse_tensor(self):
- price = fc_old.numeric_column('price')
- price_buckets = fc_old.bucketized_column(
- price, boundaries=[
- 0.,
- 10.,
- 100.,
- ])
- body_style = fc_old.categorical_column_with_vocabulary_list(
- 'body-style', vocabulary_list=['hardtop', 'wagon', 'sedan'])
- country = fc_old.categorical_column_with_vocabulary_list(
- 'country', vocabulary_list=['US', 'JP', 'CA'])
-
- # Provides 1-dim tensor and dense tensor.
- features = {
- 'price': array_ops.placeholder(dtypes.float32),
- 'body-style': array_ops.sparse_placeholder(dtypes.string),
- 'country': array_ops.placeholder(dtypes.string),
- }
- self.assertIsNone(features['price'].shape.ndims)
- self.assertIsNone(features['body-style'].get_shape().ndims)
-
- price_data = np.array([-1., 12.])
- body_style_data = sparse_tensor.SparseTensorValue(
- indices=((0,), (1,)), values=('sedan', 'hardtop'), dense_shape=(2,))
- country_data = np.array(['US', 'CA'])
-
- net = get_keras_linear_model_predictions(
- features, [price_buckets, body_style, country])
- bias = get_linear_model_bias()
- price_buckets_var = get_linear_model_column_var(price_buckets)
- body_style_var = get_linear_model_column_var(body_style)
- with _initialized_session() as sess:
- sess.run(price_buckets_var.assign([[10.], [100.], [1000.], [10000.]]))
- sess.run(body_style_var.assign([[-10.], [-100.], [-1000.]]))
- sess.run(bias.assign([5.]))
-
- self.assertAllClose([[10 - 1000 + 5.], [1000 - 10 + 5.]],
- sess.run(
- net,
- feed_dict={
- features['price']: price_data,
- features['body-style']: body_style_data,
- features['country']: country_data
- }))
-
- def test_with_rank_0_feature(self):
- price = fc_old.numeric_column('price')
- features = {
- 'price': constant_op.constant(0),
- }
- self.assertEqual(0, features['price'].shape.ndims)
-
- # Static rank 0 should fail
- with self.assertRaisesRegexp(ValueError, 'Feature .* cannot have rank 0'):
- get_keras_linear_model_predictions(features, [price])
-
- # Dynamic rank 0 should fail
- features = {
- 'price': array_ops.placeholder(dtypes.float32),
- }
- net = get_keras_linear_model_predictions(features, [price])
- self.assertEqual(1, net.shape[1])
- with _initialized_session() as sess:
- with self.assertRaisesOpError('Feature .* cannot have rank 0'):
- sess.run(net, feed_dict={features['price']: np.array(1)})
-
-
class FeatureLayerTest(test.TestCase):
@test_util.run_in_graph_and_eager_modes()
@@ -3723,47 +2834,22 @@
id_weight_pair.id_tensor.eval())
def test_linear_model(self):
- wire_column = fc_old.categorical_column_with_vocabulary_file(
+ wire_column = fc.categorical_column_with_vocabulary_file(
key='wire',
vocabulary_file=self._wire_vocabulary_file_name,
vocabulary_size=self._wire_vocabulary_size,
num_oov_buckets=1)
- self.assertEqual(4, wire_column._num_buckets)
+ self.assertEqual(4, wire_column.num_buckets)
with ops.Graph().as_default():
- predictions = fc.linear_model({
- wire_column.name: sparse_tensor.SparseTensorValue(
- indices=((0, 0), (1, 0), (1, 1)),
- values=('marlo', 'skywalker', 'omar'),
- dense_shape=(2, 2))
- }, (wire_column,))
- bias = get_linear_model_bias()
- wire_var = get_linear_model_column_var(wire_column)
- with _initialized_session():
- self.assertAllClose((0.,), bias.eval())
- self.assertAllClose(((0.,), (0.,), (0.,), (0.,)), wire_var.eval())
- self.assertAllClose(((0.,), (0.,)), predictions.eval())
- wire_var.assign(((1.,), (2.,), (3.,), (4.,))).eval()
- # 'marlo' -> 2: wire_var[2] = 3
- # 'skywalker' -> 3, 'omar' -> 0: wire_var[3] + wire_var[0] = 4+1 = 5
- self.assertAllClose(((3.,), (5.,)), predictions.eval())
-
- def test_keras_linear_model(self):
- wire_column = fc_old.categorical_column_with_vocabulary_file(
- key='wire',
- vocabulary_file=self._wire_vocabulary_file_name,
- vocabulary_size=self._wire_vocabulary_size,
- num_oov_buckets=1)
- self.assertEqual(4, wire_column._num_buckets)
- with ops.Graph().as_default():
- predictions = get_keras_linear_model_predictions({
+ model = fc.LinearModel((wire_column,))
+ predictions = model({
wire_column.name:
sparse_tensor.SparseTensorValue(
indices=((0, 0), (1, 0), (1, 1)),
values=('marlo', 'skywalker', 'omar'),
dense_shape=(2, 2))
- }, (wire_column,))
- bias = get_linear_model_bias()
- wire_var = get_linear_model_column_var(wire_column)
+ })
+ wire_var, bias = model.variables
with _initialized_session():
self.assertAllClose((0.,), bias.eval())
self.assertAllClose(((0.,), (0.,), (0.,), (0.,)), wire_var.eval())
@@ -4124,45 +3210,21 @@
id_weight_pair.id_tensor.eval())
def test_linear_model(self):
- wire_column = fc_old.categorical_column_with_vocabulary_list(
+ wire_column = fc.categorical_column_with_vocabulary_list(
key='aaa',
vocabulary_list=('omar', 'stringer', 'marlo'),
num_oov_buckets=1)
- self.assertEqual(4, wire_column._num_buckets)
+ self.assertEqual(4, wire_column.num_buckets)
with ops.Graph().as_default():
- predictions = fc.linear_model({
- wire_column.name: sparse_tensor.SparseTensorValue(
- indices=((0, 0), (1, 0), (1, 1)),
- values=('marlo', 'skywalker', 'omar'),
- dense_shape=(2, 2))
- }, (wire_column,))
- bias = get_linear_model_bias()
- wire_var = get_linear_model_column_var(wire_column)
- with _initialized_session():
- self.assertAllClose((0.,), bias.eval())
- self.assertAllClose(((0.,), (0.,), (0.,), (0.,)), wire_var.eval())
- self.assertAllClose(((0.,), (0.,)), predictions.eval())
- wire_var.assign(((1.,), (2.,), (3.,), (4.,))).eval()
- # 'marlo' -> 2: wire_var[2] = 3
- # 'skywalker' -> 3, 'omar' -> 0: wire_var[3] + wire_var[0] = 4+1 = 5
- self.assertAllClose(((3.,), (5.,)), predictions.eval())
-
- def test_keras_linear_model(self):
- wire_column = fc_old.categorical_column_with_vocabulary_list(
- key='aaa',
- vocabulary_list=('omar', 'stringer', 'marlo'),
- num_oov_buckets=1)
- self.assertEqual(4, wire_column._num_buckets)
- with ops.Graph().as_default():
- predictions = get_keras_linear_model_predictions({
+ model = fc.LinearModel((wire_column,))
+ predictions = model({
wire_column.name:
sparse_tensor.SparseTensorValue(
indices=((0, 0), (1, 0), (1, 1)),
values=('marlo', 'skywalker', 'omar'),
dense_shape=(2, 2))
- }, (wire_column,))
- bias = get_linear_model_bias()
- wire_var = get_linear_model_column_var(wire_column)
+ })
+ wire_var, bias = model.variables
with _initialized_session():
self.assertAllClose((0.,), bias.eval())
self.assertAllClose(((0.,), (0.,), (0.,), (0.,)), wire_var.eval())
@@ -4382,39 +3444,18 @@
}))
def test_linear_model(self):
- column = fc_old.categorical_column_with_identity(key='aaa', num_buckets=3)
+ column = fc.categorical_column_with_identity(key='aaa', num_buckets=3)
self.assertEqual(3, column.num_buckets)
with ops.Graph().as_default():
- predictions = fc.linear_model({
- column.name: sparse_tensor.SparseTensorValue(
- indices=((0, 0), (1, 0), (1, 1)),
- values=(0, 2, 1),
- dense_shape=(2, 2))
- }, (column,))
- bias = get_linear_model_bias()
- weight_var = get_linear_model_column_var(column)
- with _initialized_session():
- self.assertAllClose((0.,), bias.eval())
- self.assertAllClose(((0.,), (0.,), (0.,)), weight_var.eval())
- self.assertAllClose(((0.,), (0.,)), predictions.eval())
- weight_var.assign(((1.,), (2.,), (3.,))).eval()
- # weight_var[0] = 1
- # weight_var[2] + weight_var[1] = 3+2 = 5
- self.assertAllClose(((1.,), (5.,)), predictions.eval())
-
- def test_keras_linear_model(self):
- column = fc_old.categorical_column_with_identity(key='aaa', num_buckets=3)
- self.assertEqual(3, column.num_buckets)
- with ops.Graph().as_default():
- predictions = get_keras_linear_model_predictions({
+ model = fc.LinearModel((column,))
+ predictions = model({
column.name:
sparse_tensor.SparseTensorValue(
indices=((0, 0), (1, 0), (1, 1)),
values=(0, 2, 1),
dense_shape=(2, 2))
- }, (column,))
- bias = get_linear_model_bias()
- weight_var = get_linear_model_column_var(column)
+ })
+ weight_var, bias = model.variables
with _initialized_session():
self.assertAllClose((0.,), bias.eval())
self.assertAllClose(((0.,), (0.,), (0.,)), weight_var.eval())
@@ -4640,8 +3681,8 @@
self.assertAllEqual([[0., 1., 1.]], indicator_tensor.eval())
def test_linear_model(self):
- animal = fc_old.indicator_column(
- fc_old.categorical_column_with_identity('animal', num_buckets=4))
+ animal = fc.indicator_column(
+ fc.categorical_column_with_identity('animal', num_buckets=4))
with ops.Graph().as_default():
features = {
'animal':
@@ -4649,27 +3690,9 @@
indices=[[0, 0], [0, 1]], values=[1, 2], dense_shape=[1, 2])
}
- predictions = fc.linear_model(features, [animal])
- weight_var = get_linear_model_column_var(animal)
- with _initialized_session():
- # All should be zero-initialized.
- self.assertAllClose([[0.], [0.], [0.], [0.]], weight_var.eval())
- self.assertAllClose([[0.]], predictions.eval())
- weight_var.assign([[1.], [2.], [3.], [4.]]).eval()
- self.assertAllClose([[2. + 3.]], predictions.eval())
-
- def test_keras_linear_model(self):
- animal = fc_old.indicator_column(
- fc_old.categorical_column_with_identity('animal', num_buckets=4))
- with ops.Graph().as_default():
- features = {
- 'animal':
- sparse_tensor.SparseTensor(
- indices=[[0, 0], [0, 1]], values=[1, 2], dense_shape=[1, 2])
- }
-
- predictions = get_keras_linear_model_predictions(features, [animal])
- weight_var = get_linear_model_column_var(animal)
+ model = fc.LinearModel([animal])
+ predictions = model(features)
+ weight_var, _ = model.variables
with _initialized_session():
# All should be zero-initialized.
self.assertAllClose([[0.], [0.], [0.], [0.]], weight_var.eval())
@@ -5121,17 +4144,16 @@
return zeros_embedding_values
# Build columns.
- categorical_column = fc_old.categorical_column_with_identity(
+ categorical_column = fc.categorical_column_with_identity(
key='aaa', num_buckets=vocabulary_size)
- embedding_column = fc_old.embedding_column(
+ embedding_column = fc.embedding_column(
categorical_column,
dimension=embedding_dimension,
initializer=_initializer)
with ops.Graph().as_default():
- predictions = fc.linear_model({
- categorical_column.name: sparse_input
- }, (embedding_column,))
+ model = fc.LinearModel((embedding_column,))
+ predictions = model({categorical_column.name: sparse_input})
expected_var_names = (
'linear_model/bias_weights:0',
'linear_model/aaa_embedding/weights:0',
@@ -5173,82 +4195,6 @@
# = [4*7 + 6*11, 4*2 + 6*3.5, 4*0 + 6*0, 4*3 + 6*5] = [94, 29, 0, 42]
self.assertAllClose(((94.,), (29.,), (0.,), (42.,)), predictions.eval())
- def test_keras_linear_model(self):
- # Inputs.
- batch_size = 4
- vocabulary_size = 3
- sparse_input = sparse_tensor.SparseTensorValue(
- # example 0, ids [2]
- # example 1, ids [0, 1]
- # example 2, ids []
- # example 3, ids [1]
- indices=((0, 0), (1, 0), (1, 4), (3, 0)),
- values=(2, 0, 1, 1),
- dense_shape=(batch_size, 5))
-
- # Embedding variable.
- embedding_dimension = 2
- embedding_shape = (vocabulary_size, embedding_dimension)
- zeros_embedding_values = np.zeros(embedding_shape)
-
- def _initializer(shape, dtype, partition_info):
- self.assertAllEqual(embedding_shape, shape)
- self.assertEqual(dtypes.float32, dtype)
- self.assertIsNone(partition_info)
- return zeros_embedding_values
-
- # Build columns.
- categorical_column = fc_old.categorical_column_with_identity(
- key='aaa', num_buckets=vocabulary_size)
- embedding_column = fc_old.embedding_column(
- categorical_column,
- dimension=embedding_dimension,
- initializer=_initializer)
-
- with ops.Graph().as_default():
- predictions = get_keras_linear_model_predictions({
- categorical_column.name: sparse_input
- }, (embedding_column,))
- expected_var_names = (
- 'linear_model/bias_weights:0',
- 'linear_model/aaa_embedding/weights:0',
- 'linear_model/aaa_embedding/embedding_weights:0',
- )
- self.assertItemsEqual(
- expected_var_names,
- [v.name for v in ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)])
- trainable_vars = {
- v.name: v
- for v in ops.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES)
- }
- self.assertItemsEqual(expected_var_names, trainable_vars.keys())
- bias = trainable_vars['linear_model/bias_weights:0']
- embedding_weights = trainable_vars[
- 'linear_model/aaa_embedding/embedding_weights:0']
- linear_weights = trainable_vars['linear_model/aaa_embedding/weights:0']
- with _initialized_session():
- # Predictions with all zero weights.
- self.assertAllClose(np.zeros((1,)), bias.eval())
- self.assertAllClose(zeros_embedding_values, embedding_weights.eval())
- self.assertAllClose(
- np.zeros((embedding_dimension, 1)), linear_weights.eval())
- self.assertAllClose(np.zeros((batch_size, 1)), predictions.eval())
-
- # Predictions with all non-zero weights.
- embedding_weights.assign((
- (1., 2.), # id 0
- (3., 5.), # id 1
- (7., 11.) # id 2
- )).eval()
- linear_weights.assign(((4.,), (6.,))).eval()
- # example 0, ids [2], embedding[0] = [7, 11]
- # example 1, ids [0, 1], embedding[1] = mean([1, 2] + [3, 5]) = [2, 3.5]
- # example 2, ids [], embedding[2] = [0, 0]
- # example 3, ids [1], embedding[3] = [3, 5]
- # sum(embeddings * linear_weights)
- # = [4*7 + 6*11, 4*2 + 6*3.5, 4*0 + 6*0, 4*3 + 6*5] = [94, 29, 0, 42]
- self.assertAllClose(((94.,), (29.,), (0.,), (42.,)), predictions.eval())
-
def test_feature_layer(self):
# Inputs.
vocabulary_size = 3
@@ -5749,27 +4695,31 @@
return zeros_embedding_values
# Build columns.
- categorical_column_a = fc_old.categorical_column_with_identity(
+ categorical_column_a = fc.categorical_column_with_identity(
key='aaa', num_buckets=vocabulary_size)
- categorical_column_b = fc_old.categorical_column_with_identity(
+ categorical_column_b = fc.categorical_column_with_identity(
key='bbb', num_buckets=vocabulary_size)
- embedding_column_a, embedding_column_b = fc_old.shared_embedding_columns(
+ embedding_column_a, embedding_column_b = fc.shared_embedding_columns_v2(
[categorical_column_a, categorical_column_b],
dimension=embedding_dimension,
initializer=_initializer)
with ops.Graph().as_default():
- predictions = fc.linear_model({
+ model = fc.LinearModel(
+ (embedding_column_a, embedding_column_b),
+ shared_state_manager=fc.SharedEmbeddingStateManager())
+ predictions = model({
categorical_column_a.name: input_a,
- categorical_column_b.name: input_b,
- }, (embedding_column_a, embedding_column_b))
+ categorical_column_b.name: input_b
+ })
+
# Linear weights do not follow the column name. But this is a rare use
# case, and fixing it would add too much complexity to the code.
expected_var_names = (
'linear_model/bias_weights:0',
- 'linear_model/aaa_bbb_shared_embedding/weights:0',
- 'linear_model/aaa_bbb_shared_embedding/embedding_weights:0',
- 'linear_model/aaa_bbb_shared_embedding_1/weights:0',
+ 'linear_model/aaa_shared_embedding/weights:0',
+ 'shared_embedding_state_manager/aaa_bbb_shared_embedding:0',
+ 'linear_model/bbb_shared_embedding/weights:0',
)
self.assertItemsEqual(
expected_var_names,
@@ -5781,102 +4731,11 @@
self.assertItemsEqual(expected_var_names, trainable_vars.keys())
bias = trainable_vars['linear_model/bias_weights:0']
embedding_weights = trainable_vars[
- 'linear_model/aaa_bbb_shared_embedding/embedding_weights:0']
+ 'shared_embedding_state_manager/aaa_bbb_shared_embedding:0']
linear_weights_a = trainable_vars[
- 'linear_model/aaa_bbb_shared_embedding/weights:0']
+ 'linear_model/aaa_shared_embedding/weights:0']
linear_weights_b = trainable_vars[
- 'linear_model/aaa_bbb_shared_embedding_1/weights:0']
- with _initialized_session():
- # Predictions with all zero weights.
- self.assertAllClose(np.zeros((1,)), bias.eval())
- self.assertAllClose(zeros_embedding_values, embedding_weights.eval())
- self.assertAllClose(
- np.zeros((embedding_dimension, 1)), linear_weights_a.eval())
- self.assertAllClose(
- np.zeros((embedding_dimension, 1)), linear_weights_b.eval())
- self.assertAllClose(np.zeros((batch_size, 1)), predictions.eval())
-
- # Predictions with all non-zero weights.
- embedding_weights.assign((
- (1., 2.), # id 0
- (3., 5.), # id 1
- (7., 11.) # id 2
- )).eval()
- linear_weights_a.assign(((4.,), (6.,))).eval()
- # example 0, ids [2], embedding[0] = [7, 11]
- # example 1, ids [0, 1], embedding[1] = mean([1, 2] + [3, 5]) = [2, 3.5]
- # sum(embeddings * linear_weights)
- # = [4*7 + 6*11, 4*2 + 6*3.5] = [94, 29]
- linear_weights_b.assign(((3.,), (5.,))).eval()
- # example 0, ids [0], embedding[0] = [1, 2]
- # example 1, ids [], embedding[1] = 0, 0]
- # sum(embeddings * linear_weights)
- # = [3*1 + 5*2, 3*0 +5*0] = [13, 0]
- self.assertAllClose([[94. + 13.], [29.]], predictions.eval())
-
- def test_keras_linear_model(self):
- # Inputs.
- batch_size = 2
- vocabulary_size = 3
- # -1 values are ignored.
- input_a = np.array([
- [2, -1, -1], # example 0, ids [2]
- [0, 1, -1]
- ]) # example 1, ids [0, 1]
- input_b = np.array([
- [0, -1, -1], # example 0, ids [0]
- [-1, -1, -1]
- ]) # example 1, ids []
-
- # Embedding variable.
- embedding_dimension = 2
- embedding_shape = (vocabulary_size, embedding_dimension)
- zeros_embedding_values = np.zeros(embedding_shape)
-
- def _initializer(shape, dtype, partition_info):
- self.assertAllEqual(embedding_shape, shape)
- self.assertEqual(dtypes.float32, dtype)
- self.assertIsNone(partition_info)
- return zeros_embedding_values
-
- # Build columns.
- categorical_column_a = fc_old.categorical_column_with_identity(
- key='aaa', num_buckets=vocabulary_size)
- categorical_column_b = fc_old.categorical_column_with_identity(
- key='bbb', num_buckets=vocabulary_size)
- embedding_column_a, embedding_column_b = fc_old.shared_embedding_columns(
- [categorical_column_a, categorical_column_b],
- dimension=embedding_dimension,
- initializer=_initializer)
-
- with ops.Graph().as_default():
- predictions = get_keras_linear_model_predictions({
- categorical_column_a.name: input_a,
- categorical_column_b.name: input_b,
- }, (embedding_column_a, embedding_column_b))
- # Linear weights do not follow the column name. But this is a rare use
- # case, and fixing it would add too much complexity to the code.
- expected_var_names = (
- 'linear_model/bias_weights:0',
- 'linear_model/aaa_bbb_shared_embedding/weights:0',
- 'linear_model/aaa_bbb_shared_embedding/embedding_weights:0',
- 'linear_model/aaa_bbb_shared_embedding_1/weights:0',
- )
- self.assertItemsEqual(
- expected_var_names,
- [v.name for v in ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)])
- trainable_vars = {
- v.name: v
- for v in ops.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES)
- }
- self.assertItemsEqual(expected_var_names, trainable_vars.keys())
- bias = trainable_vars['linear_model/bias_weights:0']
- embedding_weights = trainable_vars[
- 'linear_model/aaa_bbb_shared_embedding/embedding_weights:0']
- linear_weights_a = trainable_vars[
- 'linear_model/aaa_bbb_shared_embedding/weights:0']
- linear_weights_b = trainable_vars[
- 'linear_model/aaa_bbb_shared_embedding_1/weights:0']
+ 'linear_model/bbb_shared_embedding/weights:0']
with _initialized_session():
# Predictions with all zero weights.
self.assertAllClose(np.zeros((1,)), bias.eval())
@@ -6275,13 +5134,14 @@
dense_shape=(2, 2)),
weight_tensor.eval())
- def test_keras_linear_model(self):
- column = fc_old.weighted_categorical_column(
- categorical_column=fc_old.categorical_column_with_identity(
+ def test_linear_model(self):
+ column = fc.weighted_categorical_column(
+ categorical_column=fc.categorical_column_with_identity(
key='ids', num_buckets=3),
weight_feature_key='values')
with ops.Graph().as_default():
- predictions = get_keras_linear_model_predictions({
+ model = fc.LinearModel((column,))
+ predictions = model({
'ids':
sparse_tensor.SparseTensorValue(
indices=((0, 0), (1, 0), (1, 1)),
@@ -6292,109 +5152,8 @@
indices=((0, 0), (1, 0), (1, 1)),
values=(.5, 1., .1),
dense_shape=(2, 2))
- }, (column,))
- bias = get_linear_model_bias()
- weight_var = get_linear_model_column_var(column)
- with _initialized_session():
- self.assertAllClose((0.,), bias.eval())
- self.assertAllClose(((0.,), (0.,), (0.,)), weight_var.eval())
- self.assertAllClose(((0.,), (0.,)), predictions.eval())
- weight_var.assign(((1.,), (2.,), (3.,))).eval()
- # weight_var[0] * weights[0, 0] = 1 * .5 = .5
- # weight_var[2] * weights[1, 0] + weight_var[1] * weights[1, 1]
- # = 3*1 + 2*.1 = 3+.2 = 3.2
- self.assertAllClose(((.5,), (3.2,)), predictions.eval())
-
- def test_keras_linear_model_mismatched_shape(self):
- column = fc_old.weighted_categorical_column(
- categorical_column=fc_old.categorical_column_with_identity(
- key='ids', num_buckets=3),
- weight_feature_key='values')
- with ops.Graph().as_default():
- with self.assertRaisesRegexp(ValueError,
- r'Dimensions.*are not compatible'):
- get_keras_linear_model_predictions({
- 'ids':
- sparse_tensor.SparseTensorValue(
- indices=((0, 0), (1, 0), (1, 1)),
- values=(0, 2, 1),
- dense_shape=(2, 2)),
- 'values':
- sparse_tensor.SparseTensorValue(
- indices=((0, 0), (0, 1), (1, 0), (1, 1)),
- values=(.5, 11., 1., .1),
- dense_shape=(2, 2))
- }, (column,))
-
- def test_keras_linear_model_mismatched_dense_values(self):
- column = fc_old.weighted_categorical_column(
- categorical_column=fc_old.categorical_column_with_identity(
- key='ids', num_buckets=3),
- weight_feature_key='values')
- with ops.Graph().as_default():
- predictions = get_keras_linear_model_predictions(
- {
- 'ids':
- sparse_tensor.SparseTensorValue(
- indices=((0, 0), (1, 0), (1, 1)),
- values=(0, 2, 1),
- dense_shape=(2, 2)),
- 'values': ((.5,), (1.,))
- }, (column,),
- sparse_combiner='mean')
- # Disabling the constant folding optimizer here since it changes the
- # error message differently on CPU and GPU.
- config = config_pb2.ConfigProto()
- config.graph_options.rewrite_options.constant_folding = (
- rewriter_config_pb2.RewriterConfig.OFF)
- with _initialized_session(config):
- with self.assertRaisesRegexp(errors.OpError, 'Incompatible shapes'):
- predictions.eval()
-
- def test_keras_linear_model_mismatched_dense_shape(self):
- column = fc_old.weighted_categorical_column(
- categorical_column=fc_old.categorical_column_with_identity(
- key='ids', num_buckets=3),
- weight_feature_key='values')
- with ops.Graph().as_default():
- predictions = get_keras_linear_model_predictions({
- 'ids':
- sparse_tensor.SparseTensorValue(
- indices=((0, 0), (1, 0), (1, 1)),
- values=(0, 2, 1),
- dense_shape=(2, 2)),
- 'values': ((.5,), (1.,), (.1,))
- }, (column,))
- bias = get_linear_model_bias()
- weight_var = get_linear_model_column_var(column)
- with _initialized_session():
- self.assertAllClose((0.,), bias.eval())
- self.assertAllClose(((0.,), (0.,), (0.,)), weight_var.eval())
- self.assertAllClose(((0.,), (0.,)), predictions.eval())
- weight_var.assign(((1.,), (2.,), (3.,))).eval()
- # weight_var[0] * weights[0, 0] = 1 * .5 = .5
- # weight_var[2] * weights[1, 0] + weight_var[1] * weights[1, 1]
- # = 3*1 + 2*.1 = 3+.2 = 3.2
- self.assertAllClose(((.5,), (3.2,)), predictions.eval())
-
- def test_linear_model(self):
- column = fc_old.weighted_categorical_column(
- categorical_column=fc_old.categorical_column_with_identity(
- key='ids', num_buckets=3),
- weight_feature_key='values')
- with ops.Graph().as_default():
- predictions = fc.linear_model({
- 'ids': sparse_tensor.SparseTensorValue(
- indices=((0, 0), (1, 0), (1, 1)),
- values=(0, 2, 1),
- dense_shape=(2, 2)),
- 'values': sparse_tensor.SparseTensorValue(
- indices=((0, 0), (1, 0), (1, 1)),
- values=(.5, 1., .1),
- dense_shape=(2, 2))
- }, (column,))
- bias = get_linear_model_bias()
- weight_var = get_linear_model_column_var(column)
+ })
+ weight_var, bias = model.variables
with _initialized_session():
self.assertAllClose((0.,), bias.eval())
self.assertAllClose(((0.,), (0.,), (0.,)), weight_var.eval())
@@ -6406,40 +5165,42 @@
self.assertAllClose(((.5,), (3.2,)), predictions.eval())
def test_linear_model_mismatched_shape(self):
- column = fc_old.weighted_categorical_column(
- categorical_column=fc_old.categorical_column_with_identity(
+ column = fc.weighted_categorical_column(
+ categorical_column=fc.categorical_column_with_identity(
key='ids', num_buckets=3),
weight_feature_key='values')
with ops.Graph().as_default():
- with self.assertRaisesRegexp(
- ValueError, r'Dimensions.*are not compatible'):
- fc.linear_model({
- 'ids': sparse_tensor.SparseTensorValue(
- indices=((0, 0), (1, 0), (1, 1)),
- values=(0, 2, 1),
- dense_shape=(2, 2)),
- 'values': sparse_tensor.SparseTensorValue(
- indices=((0, 0), (0, 1), (1, 0), (1, 1)),
- values=(.5, 11., 1., .1),
- dense_shape=(2, 2))
- }, (column,))
+ with self.assertRaisesRegexp(ValueError,
+ r'Dimensions.*are not compatible'):
+ model = fc.LinearModel((column,))
+ model({
+ 'ids':
+ sparse_tensor.SparseTensorValue(
+ indices=((0, 0), (1, 0), (1, 1)),
+ values=(0, 2, 1),
+ dense_shape=(2, 2)),
+ 'values':
+ sparse_tensor.SparseTensorValue(
+ indices=((0, 0), (0, 1), (1, 0), (1, 1)),
+ values=(.5, 11., 1., .1),
+ dense_shape=(2, 2))
+ })
def test_linear_model_mismatched_dense_values(self):
- column = fc_old.weighted_categorical_column(
- categorical_column=fc_old.categorical_column_with_identity(
+ column = fc.weighted_categorical_column(
+ categorical_column=fc.categorical_column_with_identity(
key='ids', num_buckets=3),
weight_feature_key='values')
with ops.Graph().as_default():
- predictions = fc.linear_model(
- {
- 'ids':
- sparse_tensor.SparseTensorValue(
- indices=((0, 0), (1, 0), (1, 1)),
- values=(0, 2, 1),
- dense_shape=(2, 2)),
- 'values': ((.5,), (1.,))
- }, (column,),
- sparse_combiner='mean')
+ model = fc.LinearModel((column,), sparse_combiner='mean')
+ predictions = model({
+ 'ids':
+ sparse_tensor.SparseTensorValue(
+ indices=((0, 0), (1, 0), (1, 1)),
+ values=(0, 2, 1),
+ dense_shape=(2, 2)),
+ 'values': ((.5,), (1.,))
+ })
# Disabling the constant folding optimizer here since it changes the
# error message differently on CPU and GPU.
config = config_pb2.ConfigProto()
@@ -6450,20 +5211,21 @@
predictions.eval()
def test_linear_model_mismatched_dense_shape(self):
- column = fc_old.weighted_categorical_column(
- categorical_column=fc_old.categorical_column_with_identity(
+ column = fc.weighted_categorical_column(
+ categorical_column=fc.categorical_column_with_identity(
key='ids', num_buckets=3),
weight_feature_key='values')
with ops.Graph().as_default():
- predictions = fc.linear_model({
- 'ids': sparse_tensor.SparseTensorValue(
- indices=((0, 0), (1, 0), (1, 1)),
- values=(0, 2, 1),
- dense_shape=(2, 2)),
+ model = fc.LinearModel((column,))
+ predictions = model({
+ 'ids':
+ sparse_tensor.SparseTensorValue(
+ indices=((0, 0), (1, 0), (1, 1)),
+ values=(0, 2, 1),
+ dense_shape=(2, 2)),
'values': ((.5,), (1.,), (.1,))
- }, (column,))
- bias = get_linear_model_bias()
- weight_var = get_linear_model_column_var(column)
+ })
+ weight_var, bias = model.variables
with _initialized_session():
self.assertAllClose((0.,), bias.eval())
self.assertAllClose(((0.,), (0.,), (0.,)), weight_var.eval())
diff --git a/tensorflow/python/framework/device.py b/tensorflow/python/framework/device.py
index 06c6530..7f6e0a7 100644
--- a/tensorflow/python/framework/device.py
+++ b/tensorflow/python/framework/device.py
@@ -87,6 +87,7 @@
else:
self.device_type = device_type
self.device_index = device_index
+ self._hash = hash(self.to_string())
def _clear(self):
self._job = None
@@ -234,7 +235,7 @@
return self.to_string() == other.to_string()
def __hash__(self):
- return hash(self.to_string())
+ return self._hash
def check_valid(spec):
@@ -266,6 +267,7 @@
# possible to compare the device function stacks belonging to different
# graphs in a meaningful way.
_cached_device_functions = {}
+_cached_device_specs = {}
_cache_lock = threading.Lock()
@@ -297,7 +299,13 @@
"""
with _cache_lock:
if not isinstance(spec, DeviceSpec):
- spec = DeviceSpec.from_string(spec or "")
+ cached_device_spec = _cached_device_specs.get(spec, None)
+ if cached_device_spec is None:
+ device_spec = DeviceSpec.from_string(spec or "")
+ _cached_device_specs[spec] = device_spec
+ spec = device_spec
+ else:
+ spec = cached_device_spec
cached_function = _cached_device_functions.get(spec, None)
if cached_function is not None:
return cached_function
diff --git a/tensorflow/python/framework/dtypes.py b/tensorflow/python/framework/dtypes.py
index c3f70df..64d3b42 100644
--- a/tensorflow/python/framework/dtypes.py
+++ b/tensorflow/python/framework/dtypes.py
@@ -26,7 +26,7 @@
_np_bfloat16 = pywrap_tensorflow.TF_bfloat16_type()
-@tf_export("DType")
+@tf_export("dtypes.DType", "DType")
class DType(object):
"""Represents the type of the elements in a `Tensor`.
@@ -658,7 +658,7 @@
}
-@tf_export("as_dtype")
+@tf_export("dtypes.as_dtype", "as_dtype")
def as_dtype(type_value):
"""Converts the given `type_value` to a `DType`.
diff --git a/tensorflow/python/framework/errors_impl.py b/tensorflow/python/framework/errors_impl.py
index 5af71f2..8b303fa 100644
--- a/tensorflow/python/framework/errors_impl.py
+++ b/tensorflow/python/framework/errors_impl.py
@@ -25,11 +25,13 @@
from tensorflow.python import pywrap_tensorflow as c_api
from tensorflow.python.framework import c_api_util
from tensorflow.python.util import compat
+from tensorflow.python.util import deprecation
from tensorflow.python.util import tf_inspect
from tensorflow.python.util.tf_export import tf_export
-@tf_export("OpError", "errors.OpError")
+@tf_export("errors.OpError", "OpError")
+@deprecation.deprecated_endpoints("OpError")
class OpError(Exception):
"""A generic error that is raised when TensorFlow execution fails.
@@ -72,7 +74,7 @@
or `Recv` op, there will be no corresponding
`tf.Operation`
object. In that case, this will return `None`, and you should
- instead use the `tf.OpError.node_def` to
+ instead use the `tf.errors.OpError.node_def` to
discover information about the op.
Returns:
diff --git a/tensorflow/python/framework/function_test.py b/tensorflow/python/framework/function_test.py
index f740e5c..87f567d 100644
--- a/tensorflow/python/framework/function_test.py
+++ b/tensorflow/python/framework/function_test.py
@@ -113,7 +113,7 @@
return a
with ops.Graph().as_default():
- var = variables.Variable([18.0])
+ var = variables.VariableV1([18.0])
call = MyIdentityFunc(var._ref()) # pylint: disable=protected-access
self.assertEqual("MyIdentity", call.op.name)
for cfg in _OptimizerOptions():
diff --git a/tensorflow/python/framework/graph_io.py b/tensorflow/python/framework/graph_io.py
index be30b16..47e1344 100644
--- a/tensorflow/python/framework/graph_io.py
+++ b/tensorflow/python/framework/graph_io.py
@@ -27,7 +27,7 @@
from tensorflow.python.util.tf_export import tf_export
-@tf_export('train.write_graph')
+@tf_export('io.write_graph', 'train.write_graph')
def write_graph(graph_or_graph_def, logdir, name, as_text=True):
"""Writes a graph proto to a file.
diff --git a/tensorflow/python/framework/graph_util_test.py b/tensorflow/python/framework/graph_util_test.py
index 2dafb94..563a177 100644
--- a/tensorflow/python/framework/graph_util_test.py
+++ b/tensorflow/python/framework/graph_util_test.py
@@ -104,13 +104,13 @@
def testNestedDeviceFunctions(self):
with ops.Graph().as_default():
- var_0 = variables.Variable(0)
+ var_0 = variables.VariableV1(0)
with ops.device(test_device_func_pin_variable_to_cpu):
- var_1 = variables.Variable(1)
+ var_1 = variables.VariableV1(1)
with ops.device(lambda op: "/device:GPU:0"):
- var_2 = variables.Variable(2)
+ var_2 = variables.VariableV1(2)
with ops.device("/device:GPU:0"): # Implicit merging device function.
- var_3 = variables.Variable(3)
+ var_3 = variables.VariableV1(3)
self.assertDeviceEqual(var_0.device, None)
self.assertDeviceEqual(var_1.device, "/device:CPU:0")
diff --git a/tensorflow/python/framework/importer.py b/tensorflow/python/framework/importer.py
index e48e67c..c659591 100644
--- a/tensorflow/python/framework/importer.py
+++ b/tensorflow/python/framework/importer.py
@@ -329,7 +329,7 @@
node_def.attr[key].CopyFrom(attr_def.default_value)
-@tf_export('import_graph_def')
+@tf_export('graph_util.import_graph_def', 'import_graph_def')
@deprecated_args(None, 'Please file an issue at '
'https://github.com/tensorflow/tensorflow/issues if you depend'
' on this feature.', 'op_dict')
diff --git a/tensorflow/python/framework/random_seed.py b/tensorflow/python/framework/random_seed.py
index 2f95048..6f9f347 100644
--- a/tensorflow/python/framework/random_seed.py
+++ b/tensorflow/python/framework/random_seed.py
@@ -22,6 +22,7 @@
from tensorflow.python.eager import context
from tensorflow.python.framework import ops
+from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
@@ -33,7 +34,8 @@
return seed % _MAXINT32 # Truncate to fit into 32-bit integer
-@tf_export('get_seed')
+@tf_export('random.get_seed', 'get_seed')
+@deprecation.deprecated_endpoints('get_seed')
def get_seed(op_seed):
"""Returns the local seeds an operation should use given an op-specific seed.
@@ -80,7 +82,7 @@
return seeds
-@tf_export('set_random_seed')
+@tf_export('random.set_random_seed', 'set_random_seed')
def set_random_seed(seed):
"""Sets the graph-level random seed.
diff --git a/tensorflow/python/framework/sparse_tensor.py b/tensorflow/python/framework/sparse_tensor.py
index d1bdd9b..440e3a0 100644
--- a/tensorflow/python/framework/sparse_tensor.py
+++ b/tensorflow/python/framework/sparse_tensor.py
@@ -33,7 +33,7 @@
# pylint: enable=protected-access
-@tf_export("SparseTensor")
+@tf_export("sparse.SparseTensor", "SparseTensor")
class SparseTensor(_TensorLike):
"""Represents a sparse tensor.
@@ -245,7 +245,7 @@
SparseTensorValue = collections.namedtuple(
"SparseTensorValue", ["indices", "values", "dense_shape"])
tf_export("SparseTensorValue")(SparseTensorValue)
-pywrap_tensorflow.RegisterSparseTensorValueClass(SparseTensorValue)
+pywrap_tensorflow.RegisterType("SparseTensorValue", SparseTensorValue)
@tf_export("convert_to_tensor_or_sparse_tensor")
diff --git a/tensorflow/python/framework/subscribe_test.py b/tensorflow/python/framework/subscribe_test.py
index 1d594e4..cab4268 100644
--- a/tensorflow/python/framework/subscribe_test.py
+++ b/tensorflow/python/framework/subscribe_test.py
@@ -212,8 +212,8 @@
def testSubscribeVariable(self):
"""Confirm that variables can be subscribed."""
- v1 = variables.Variable(0.0)
- v2 = variables.Variable(4.0)
+ v1 = variables.VariableV1(0.0)
+ v2 = variables.VariableV1(4.0)
add = math_ops.add(v1, v2)
assign_v1 = v1.assign(3.0)
diff --git a/tensorflow/python/framework/test_util.py b/tensorflow/python/framework/test_util.py
index cd0b03b..6673bc5 100644
--- a/tensorflow/python/framework/test_util.py
+++ b/tensorflow/python/framework/test_util.py
@@ -24,8 +24,8 @@
import contextlib
import gc
import itertools
-import os
import math
+import os
import random
import re
import tempfile
@@ -402,11 +402,14 @@
return cls
-def enable_cond_v2(fn):
- """Decorator for enabling CondV2 on a test.
+def enable_control_flow_v2(fn):
+ """Decorator for enabling CondV2 and WhileV2 on a test.
- Note this enables using CondV2 after running the test class's setup/teardown
- methods.
+ Note this enables using CondV2 and WhileV2 after running the test class's
+ setup/teardown methods.
+
+ In addition to this, callers must import the while_v2 module in order to set
+ the _while_v2 module in control_flow_ops.
Args:
fn: the function to be wrapped
@@ -416,21 +419,56 @@
"""
def wrapper(*args, **kwargs):
- prev_value = control_flow_ops.ENABLE_COND_V2
+ enable_cond_v2_old = control_flow_ops.ENABLE_COND_V2
+ enable_while_v2_old = control_flow_ops.ENABLE_WHILE_V2
control_flow_ops.ENABLE_COND_V2 = True
+ control_flow_ops.ENABLE_WHILE_V2 = True
try:
fn(*args, **kwargs)
finally:
- control_flow_ops.ENABLE_COND_V2 = prev_value
+ control_flow_ops.ENABLE_COND_V2 = enable_cond_v2_old
+ control_flow_ops.ENABLE_WHILE_V2 = enable_while_v2_old
return wrapper
-def with_cond_v2(cls):
- """Adds methods that call original methods but with CondV2 enabled.
+def with_control_flow_v2(cls):
+ """Adds methods that call original methods with WhileV2 and CondV2 enabled.
- Note this enables CondV2 in new methods after running the test class's
- setup method.
+ Note this enables CondV2 and WhileV2 in new methods after running the test
+ class's setup method.
+
+ In addition to this, callers must import the while_v2 module in order to set
+ the _while_v2 module in control_flow_ops.
+
+ If a test function has _disable_control_flow_v2 attr set to True (using the
+ @disable_control_flow_v2 decorator), the v2 function is not generated for it.
+
+ Example:
+
+ @test_util.with_control_flow_v2
+ class ControlFlowTest(test.TestCase):
+
+ def testEnabledForV2(self):
+ ...
+
+ @test_util.disable_control_flow_v2("b/xyzabc")
+ def testDisabledForV2(self):
+ ...
+
+ Generated class:
+ class ControlFlowTest(test.TestCase):
+
+ def testEnabledForV2(self):
+ ...
+
+ def testEnabledForV2WithControlFlowV2(self):
+ // Enable V2 flags.
+ testEnabledForV2(self)
+ // Restore V2 flags.
+
+ def testDisabledForV2(self):
+ ...
Args:
cls: class to decorate
@@ -438,15 +476,33 @@
Returns:
cls with new test methods added
"""
- if control_flow_ops.ENABLE_COND_V2:
+ if control_flow_ops.ENABLE_WHILE_V2 and control_flow_ops.ENABLE_COND_V2:
return cls
for name, value in cls.__dict__.copy().items():
- if callable(value) and name.startswith("test"):
- setattr(cls, name + "WithCondV2", enable_cond_v2(value))
+ if (callable(value) and name.startswith("test") and
+ not getattr(value, "_disable_control_flow_v2", False)):
+ setattr(cls, name + "WithControlFlowV2", enable_control_flow_v2(value))
return cls
+def disable_control_flow_v2(unused_msg):
+ """Decorator for a function in a with_control_flow_v2 enabled test class.
+
+ Blocks the function from being run with v2 control flow ops.
+
+ Args:
+ unused_msg: Reason for disabling.
+
+ Returns:
+ The wrapped function with _disable_control_flow_v2 attr set to True.
+ """
+ def wrapper(func):
+ func._disable_control_flow_v2 = True
+ return func
+ return wrapper
+
+
def assert_no_new_pyobjects_executing_eagerly(f):
"""Decorator for asserting that no new Python objects persist after a test.
diff --git a/tensorflow/python/grappler/item_test.py b/tensorflow/python/grappler/item_test.py
index c40de9d..d3d96c6 100644
--- a/tensorflow/python/grappler/item_test.py
+++ b/tensorflow/python/grappler/item_test.py
@@ -110,7 +110,7 @@
def testColocationContraints(self):
with ops.Graph().as_default() as g:
c = constant_op.constant([10])
- v = variables.Variable([3], dtype=dtypes.int32)
+ v = variables.VariableV1([3], dtype=dtypes.int32)
i = gen_array_ops.ref_identity(v)
a = state_ops.assign(i, c)
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
diff --git a/tensorflow/python/grappler/memory_optimizer_test.py b/tensorflow/python/grappler/memory_optimizer_test.py
index b658edf..03b42f6 100644
--- a/tensorflow/python/grappler/memory_optimizer_test.py
+++ b/tensorflow/python/grappler/memory_optimizer_test.py
@@ -39,8 +39,8 @@
def testNoSwapping(self):
"""Make sure the graph is preserved when there is nothing to swap."""
- a = variables.Variable(10, name='a')
- b = variables.Variable(20, name='b')
+ a = variables.VariableV1(10, name='a')
+ b = variables.VariableV1(20, name='b')
c = math_ops.add_n([a, b], name='c')
d = math_ops.add_n([b, c], name='d')
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
@@ -60,8 +60,8 @@
def testSimpleSwap(self):
"""Check that the swap annotations are followed."""
- a = variables.Variable(10, name='a')
- b = variables.Variable(20, name='b')
+ a = variables.VariableV1(10, name='a')
+ b = variables.VariableV1(20, name='b')
c = math_ops.add_n([a, b], name='c')
d = math_ops.add_n([b, c], name='d')
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
@@ -244,7 +244,7 @@
init_op_name=init_op_name,
train_op_name=train_op_name,
loss_op_name=loss_op_name)
- self.assertAllClose(original_loss, memory_optimized_loss, rtol=1e-4)
+ self.assertAllClose(original_loss, memory_optimized_loss, rtol=1e-2)
def _annotated_graph(self):
graph = ops.Graph()
diff --git a/tensorflow/python/grappler/tf_optimizer_test.py b/tensorflow/python/grappler/tf_optimizer_test.py
index 5a9afe7..eca0f67 100644
--- a/tensorflow/python/grappler/tf_optimizer_test.py
+++ b/tensorflow/python/grappler/tf_optimizer_test.py
@@ -57,7 +57,7 @@
def testKeepNodes(self):
g = ops.Graph()
with g.as_default():
- a1 = variables.Variable(
+ a1 = variables.VariableV1(
1.0) # Must be preserved since it's in the collection 'variables'.
a2 = constant_op.constant(0, shape=[50, 50], name='keep')
ops.add_to_collection('a2', a2) # Explicitly add to collection.
diff --git a/tensorflow/python/keras/backend.py b/tensorflow/python/keras/backend.py
index 4589c82..584facc 100644
--- a/tensorflow/python/keras/backend.py
+++ b/tensorflow/python/keras/backend.py
@@ -1511,12 +1511,8 @@
out = math_ops.reduce_sum(
math_ops.multiply(array_ops.transpose(x, [1, 0]), y), axes[1])
else:
- if axes is not None:
- adj_x = None if axes[0] == ndim(x) - 1 else True
- adj_y = True if axes[1] == ndim(y) - 1 else None
- else:
- adj_x = None
- adj_y = None
+ adj_x = None if axes[0] == ndim(x) - 1 else True
+ adj_y = True if axes[1] == ndim(y) - 1 else None
out = math_ops.matmul(x, y, adjoint_a=adj_x, adjoint_b=adj_y)
if diff:
if x_ndim > y_ndim:
diff --git a/tensorflow/python/keras/callbacks_test.py b/tensorflow/python/keras/callbacks_test.py
index b6fae19..467bc4c 100644
--- a/tensorflow/python/keras/callbacks_test.py
+++ b/tensorflow/python/keras/callbacks_test.py
@@ -30,6 +30,7 @@
from tensorflow.core.framework import summary_pb2
from tensorflow.python import keras
+from tensorflow.python.framework import ops
from tensorflow.python.framework import random_seed
from tensorflow.python.framework import test_util
from tensorflow.python.keras import testing_utils
@@ -1222,6 +1223,45 @@
callbacks=cbks,
epochs=1)
+ def test_fit_generator_with_callback(self):
+
+ class TestCallback(keras.callbacks.Callback):
+ def set_model(self, model):
+ # Check the model operations for the optimizer operations that
+ # the _make_train_function adds under a named scope for the
+ # optimizer. This ensurs the full model is populated before the
+ # set_model callback is called.
+ optimizer_name_scope = 'training/' + model.optimizer.__class__.__name__
+ graph_def = ops.get_default_graph().as_graph_def()
+ for node in graph_def.node:
+ if node.name.startswith(optimizer_name_scope):
+ return
+ raise RuntimeError('The optimizer operations are not present in the '
+ 'model graph when the Callback.set_model function '
+ 'is called')
+ np.random.seed(1337)
+
+ def generator():
+ x = np.random.randn(10, 100).astype(np.float32)
+ y = np.random.randn(10, 10).astype(np.float32)
+ while True:
+ yield x, y
+
+ with self.cached_session():
+ model = testing_utils.get_small_sequential_mlp(
+ num_hidden=10, num_classes=10, input_dim=100)
+ model.compile(
+ loss='categorical_crossentropy',
+ optimizer='sgd',
+ metrics=['accuracy'])
+ model.fit_generator(
+ generator(),
+ steps_per_epoch=2,
+ epochs=1,
+ validation_data=generator(),
+ validation_steps=2,
+ callbacks=[TestCallback()],
+ verbose=0)
if __name__ == '__main__':
test.main()
diff --git a/tensorflow/python/keras/engine/base_layer.py b/tensorflow/python/keras/engine/base_layer.py
index cb19a41..a75ce30 100644
--- a/tensorflow/python/keras/engine/base_layer.py
+++ b/tensorflow/python/keras/engine/base_layer.py
@@ -20,6 +20,7 @@
import collections as collections_lib
import enum # pylint: disable=g-bad-import-order
+import functools
import inspect # Necessary supplement to tf_inspect to deal with variadic args.
import numpy as np
@@ -160,9 +161,13 @@
self._trainable_weights = []
self._non_trainable_weights = []
self._updates = []
- # When executing eagerly, _losses is a list of zero-argument lambdas which
- # return tensors. When using graph execution, _losses is a list of ops.
+ # A list of zero-argument lambdas which return Tensors, used for variable
+ # regularizers.
+ self._callable_losses = []
+ # A list of Tensors containing activity regularizers and losses manually
+ # added through `add_loss`. Empty when executing eagerly.
self._losses = []
+ self._in_call = False # Flag for error checking in add_loss
self._dtype = None if dtype is None else dtypes.as_dtype(dtype).name
self._call_fn_args = function_utils.fn_args(self.call)
self._compute_previous_mask = ('mask' in self._call_fn_args or
@@ -359,20 +364,20 @@
def losses(self):
"""Losses which are associated with this `Layer`.
- Note that when executing eagerly, getting this property evaluates
- regularizers. When using graph execution, variable regularization ops have
- already been created and are simply returned here.
+ Variable regularization tensors are created when this property is accessed,
+ so it is eager safe: accessing `losses` under a `tf.GradientTape` will
+ propagate gradients back to the corresponding variables.
Returns:
A list of tensors.
"""
- if context.executing_eagerly():
- # _losses may only contain variable regularization losses when executing
- # eagerly, and they have been saved as lambdas to be executed when
- # requested.
- return [regularizer() for regularizer in self._losses]
- else:
- return self._losses
+ collected_losses = []
+ collected_losses.extend(self._losses)
+ for regularizer in self._callable_losses:
+ loss_tensor = regularizer()
+ if loss_tensor is not None:
+ collected_losses.append(loss_tensor)
+ return collected_losses
@doc_controls.for_subclass_implementers
def add_loss(self, losses, inputs=None):
@@ -393,7 +398,9 @@
from `Layer.call()`).
Arguments:
- losses: Loss tensor, or list/tuple of tensors.
+ losses: Loss tensor, or list/tuple of tensors. Rather than tensors, losses
+ may also be zero-argument callables which create a loss tensor. Only
+ callable losses are supported when executing eagerly.
inputs: If anything other than None is passed, it signals the losses
are conditional on some of the layer's inputs,
and thus they should only be run where these inputs are available.
@@ -403,29 +410,45 @@
(e.g. weight regularization losses).
Raises:
- RuntimeError: If called in Eager mode.
+ RuntimeError: If called in Eager mode with a `Tensor` rather than a
+ callable, or if `inputs` is not None.
"""
- if context.executing_eagerly():
- # TODO(fchollet): it should be possible (and highly desirable) to support
- # `add_loss` in eager mode. This allows great convenience and flexibility
- # in defining custom losses on the fly (e.g. in VAEs).
- # Simply appending the loss value to `self._losses`
- # is the correct behavior.
- # The only caveat is that we need to force the user to only call
- # `add_loss` from inside a model or Layer's `call` method
- # (otherwise the loss computation cannot be backproped through).
- raise RuntimeError('Layer.add_loss not supported in Eager mode.')
-
+ executing_eagerly = context.executing_eagerly()
+ if executing_eagerly:
+ if inputs is not None:
+ raise RuntimeError(
+ 'Activity regularization (via the "inputs" argument to '
+ 'Layer.add_loss) is not supported when executing eagerly. Consider '
+ 'returning activity regularization losses from a Model\'s call() '
+ 'method.')
+ if getattr(self, '_in_call', False):
+ # TODO(psv): Support activity regularization and a way to reset losses.
+ raise RuntimeError(
+ 'Adding losses inside a Layer\'s call() method is not currently '
+ 'supported when executing eagerly. Please file a feature request '
+ 'if you need this limitation lifted.')
losses = generic_utils.to_list(losses)
- losses = [ops.convert_to_tensor(loss, dtype=backend.floatx())
- if not tensor_util.is_tensor(loss) else loss for loss in losses]
- self._losses += losses
- if inputs is None:
- for loss in losses:
- loss._unconditional_loss = True # pylint: disable=protected-access
- else:
- for loss in losses:
- loss._unconditional_loss = False # pylint: disable=protected-access
+
+ def _tag_unconditional(loss):
+ if callable(loss):
+ loss = loss()
+ if loss is None:
+ return None # Will be filtered out when computing the .losses property
+ if not tensor_util.is_tensor(loss):
+ loss = ops.convert_to_tensor(loss, dtype=backend.floatx())
+ loss._unconditional_loss = (inputs is None) # pylint: disable=protected-access
+ return loss
+
+ for loss in losses:
+ if callable(loss):
+ self._callable_losses.append(
+ functools.partial(_tag_unconditional, loss))
+ else:
+ if executing_eagerly:
+ raise RuntimeError(
+ 'Layer.add_loss only supported for zero-argument lambdas when '
+ 'executing eagerly.')
+ self._losses.append(_tag_unconditional(loss))
def get_losses_for(self, inputs):
"""Retrieves losses relevant to a specific set of inputs.
@@ -599,56 +622,20 @@
return variable
def _handle_weight_regularization(self, name, variable, regularizer):
- # `init_graph` should point to the graph in which variable initialization
- # will occur; it should be None if and only if initialization will take
- # place in the eager context.
- init_graph = None
- if not context.executing_eagerly():
- default_graph = ops.get_default_graph()
- if default_graph.building_function:
- with ops.init_scope():
- # Retrieve the variables from the graph into which variables
- # will be lifted; if initialization ops will be lifted into
- # the eager context, then there is nothing to retrieve, since variable
- # collections are not supported when eager execution is enabled.
- if not context.executing_eagerly():
- init_graph = ops.get_default_graph()
- else:
- # Initialization ops will not be lifted out of the default graph.
- init_graph = default_graph
+ """Create lambdas which compute regularization losses."""
- if init_graph is not None: # pylint: disable=protected-access
- # The variable was created and initialized in a graph.
- if regularizer:
- if isinstance(variable, tf_variables.PartitionedVariable):
- for v in variable:
- with ops.colocate_with(v.op):
- with ops.name_scope(name + '/Regularizer'):
- regularization = regularizer(v)
- if regularization is not None:
- self.add_loss(regularization)
- else:
- with ops.colocate_with(variable.op):
- with ops.name_scope(name + '/Regularizer'):
- regularization = regularizer(variable)
- if regularization is not None:
- self.add_loss(regularization)
- elif regularizer: # initialization took place in an eager context
- if isinstance(variable, tf_variables.PartitionedVariable):
- raise RuntimeError(
- 'Partitioned variable regularization is not yet '
- 'supported when executing eagerly. File a feature request'
- 'if this is important to you.')
- # Save a zero-argument lambda which runs the regularizer on the
- # variable, to be executed when `Layer.losses` is requested.
- # This makes losses responsive to variable updates when executing
- # eagerly.
- #
- # TODO(akshayka): Do the same for graphs as well, so that losses
- # collected in a while_loop can be run outside its control flow
- # context and so that losses won't be swallowed up by graph functions
- # (i.e., `.losses()` should always create regularizers).
- self._losses.append(lambda: regularizer(variable))
+ 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)
+ return regularization
+
+ if isinstance(variable, tf_variables.PartitionedVariable):
+ for v in variable:
+ self.add_loss(functools.partial(_loss_for_variable, v))
+ else:
+ self.add_loss(functools.partial(_loss_for_variable, variable))
def _handle_activity_regularization(self, inputs, outputs):
# Apply activity regularization.
@@ -766,7 +753,9 @@
self._assert_input_compatibility(inputs)
if not in_deferred_mode:
+ self._in_call = True
outputs = self.call(inputs, *args, **kwargs)
+ self._in_call = False
if outputs is None:
raise ValueError('A layer\'s `call` method should return a Tensor '
'or a list of Tensors, not None (layer: ' +
@@ -1972,7 +1961,9 @@
if use_resource is None:
use_resource = True
- v = tf_variables.Variable(
+ # TODO(apassos,rohanj) figure out how to remove collections from here so we
+ # can remove the V1.
+ v = tf_variables.VariableV1(
initial_value=init_val,
name=name,
trainable=trainable,
diff --git a/tensorflow/python/keras/engine/training.py b/tensorflow/python/keras/engine/training.py
index ade8a4b..5091cac 100644
--- a/tensorflow/python/keras/engine/training.py
+++ b/tensorflow/python/keras/engine/training.py
@@ -647,12 +647,6 @@
skip_target_indices=skip_target_indices,
sample_weights=self.sample_weights)
- # If using distribution strategy and stateful_metrics, raise an error
- # since we currently don't support stateful metrics.
- if self._distribution_strategy is not None and self.stateful_metric_names:
- raise NotImplementedError('Stateful metrics are not supported with '
- 'DistributionStrategy.')
-
# Prepare gradient updates and state updates.
self.total_loss = total_loss
@@ -857,7 +851,8 @@
# able to clone a Dataset on multiple workers we can remove this lambda.
result = self._distribution_strategy.distribute_dataset(lambda: x)
iterator = result.make_initializable_iterator()
- K.get_session().run(iterator.initializer)
+ with self._distribution_strategy.scope():
+ K.get_session().run(iterator.initializer)
training_utils.validate_iterator_input(x, y, sample_weight,
validation_split)
diff --git a/tensorflow/python/keras/engine/training_distributed.py b/tensorflow/python/keras/engine/training_distributed.py
index 8b434ca..a647045 100644
--- a/tensorflow/python/keras/engine/training_distributed.py
+++ b/tensorflow/python/keras/engine/training_distributed.py
@@ -26,6 +26,7 @@
from tensorflow.python.keras import callbacks as cbks
from tensorflow.python.keras import optimizers
from tensorflow.python.keras.engine import distributed_training_utils
+from tensorflow.python.keras import metrics as metrics_module
from tensorflow.python.keras.utils.generic_utils import Progbar
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import variable_scope
@@ -111,96 +112,99 @@
dataset_targets = distributed_training_utils.flatten_perdevice_values(
current_strategy, targets)
- # Create a train function that is composed of all the parameters above.
- distributed_train_function = K.Function(
- all_inputs, all_outputs,
- updates=all_updates,
- name='distributed_train_function',
- **all_session_args)
+ # Create a train function that is composed of all the parameters above.
+ distributed_train_function = K.Function(
+ all_inputs, all_outputs,
+ updates=all_updates,
+ name='distributed_train_function',
+ **all_session_args)
- # We need to set sample_weights to None since there are sample weight
- # placeholders that are created with default values.
- sample_weights = [None for _ in range(len(model.outputs) *
- current_strategy.num_towers)]
- if model.uses_learning_phase and not isinstance(K.learning_phase(), int):
- ins = dataset_inputs + dataset_targets + sample_weights + [1]
- else:
- ins = dataset_inputs + dataset_targets
+ # We need to set sample_weights to None since there are sample weight
+ # placeholders that are created with default values.
+ sample_weights = [None for _ in range(len(model.outputs) *
+ current_strategy.num_towers)]
+ if model.uses_learning_phase and not isinstance(K.learning_phase(), int):
+ ins = dataset_inputs + dataset_targets + sample_weights + [1]
+ else:
+ ins = dataset_inputs + dataset_targets
- do_validation = False
- if validation_steps:
- do_validation = True
+ do_validation = False
+ if validation_steps:
+ do_validation = True
- # Copy the weights from the original model to each of the replicated models.
- orig_model_weights = model.get_weights()
- with current_strategy.scope():
+ # Copy the weights from the original model to each of the replicated models.
+ orig_model_weights = model.get_weights()
distributed_model = current_strategy.unwrap(model._grouped_model)[0]
distributed_training_utils.set_weights(
current_strategy, distributed_model, orig_model_weights)
- callbacks = cbks.configure_callbacks(
- callbacks,
- model,
- do_validation=do_validation,
- val_inputs=None,
- val_targets=None,
- epochs=epochs,
- steps_per_epoch=steps_per_epoch,
- verbose=verbose)
- out_labels = model.metrics_names or []
- callbacks.on_train_begin()
+ callbacks = cbks.configure_callbacks(
+ callbacks,
+ model,
+ do_validation=do_validation,
+ val_inputs=None,
+ val_targets=None,
+ epochs=epochs,
+ steps_per_epoch=steps_per_epoch,
+ verbose=verbose)
+ out_labels = model.metrics_names or []
+ callbacks.on_train_begin()
- assert steps_per_epoch is not None
+ assert steps_per_epoch is not None
- for epoch in range(initial_epoch, epochs):
- callbacks.on_epoch_begin(epoch)
- epoch_logs = {}
- for step_index in range(steps_per_epoch):
- batch_logs = {'batch': step_index, 'size': 1}
- callbacks.on_batch_begin(step_index, batch_logs)
- try:
- outs = distributed_train_function(ins)
- except errors.OutOfRangeError:
- logging.warning('Your dataset iterator ran out of data; '
- 'interrupting training. Make sure that your dataset '
- 'can generate at least `steps_per_epoch * epochs` '
- 'batches (in this case, %d batches).' %
- steps_per_epoch * epochs)
- break
+ for epoch in range(initial_epoch, epochs):
+ # Reset stateful metrics
+ for m in model.stateful_metric_functions:
+ m.reset_states()
+ callbacks.on_epoch_begin(epoch)
+ epoch_logs = {}
+ for step_index in range(steps_per_epoch):
+ batch_logs = {'batch': step_index, 'size': 1}
+ callbacks.on_batch_begin(step_index, batch_logs)
+ try:
+ outs = distributed_train_function(ins)
+ except errors.OutOfRangeError:
+ logging.warning('Your dataset iterator ran out of data; '
+ 'interrupting training. Make sure that your dataset '
+ 'can generate at least `steps_per_epoch * epochs` '
+ 'batches (in this case, %d batches).' %
+ steps_per_epoch * epochs)
+ break
- if not isinstance(outs, list):
- outs = [outs]
+ if not isinstance(outs, list):
+ outs = [outs]
- outs = _aggregate_metrics_across_towers(
- current_strategy.num_towers, out_labels, outs)
- for l, o in zip(out_labels, outs):
- batch_logs[l] = o
- callbacks.on_batch_end(step_index, batch_logs)
+ outs = _aggregate_metrics_across_towers(current_strategy.num_towers,
+ out_labels,
+ model.stateful_metric_names,
+ outs)
+ for l, o in zip(out_labels, outs):
+ batch_logs[l] = o
+ callbacks.on_batch_end(step_index, batch_logs)
+ if callbacks.model.stop_training:
+ break
+ if do_validation:
+ val_outs = test_loop(
+ model,
+ val_iterator,
+ steps=validation_steps,
+ verbose=0)
+ if not isinstance(val_outs, list):
+ val_outs = [val_outs]
+ # Same labels assumed.
+ for l, o in zip(out_labels, val_outs):
+ epoch_logs['val_' + l] = o
+
+ callbacks.on_epoch_end(epoch, epoch_logs)
if callbacks.model.stop_training:
break
- if do_validation:
- val_outs = test_loop(
- model,
- val_iterator,
- steps=validation_steps,
- verbose=0)
- if not isinstance(val_outs, list):
- val_outs = [val_outs]
- # Same labels assumed.
- for l, o in zip(out_labels, val_outs):
- epoch_logs['val_' + l] = o
+ callbacks.on_train_end()
- callbacks.on_epoch_end(epoch, epoch_logs)
- if callbacks.model.stop_training:
- break
- callbacks.on_train_end()
-
- # Copy the weights back from the replicated model to the original model.
- with current_strategy.scope():
+ # Copy the weights back from the replicated model to the original model.
updated_weights = current_strategy.unwrap(
model._grouped_model)[0].get_weights()
model.set_weights(updated_weights)
- return model.history
+ return model.history
def _experimental_fit_loop(
@@ -422,54 +426,65 @@
dataset_targets = distributed_training_utils.flatten_perdevice_values(
current_strategy, targets)
- distributed_test_function = K.Function(
- all_inputs, all_outputs,
- updates=all_updates,
- name='distributed_test_function',
- **all_session_args)
+ distributed_test_function = K.Function(
+ all_inputs, all_outputs,
+ updates=all_updates,
+ name='distributed_test_function',
+ **all_session_args)
- # We need to set sample_weights to None since there are sample weight
- # placeholders that are created with default values.
- sample_weights = [None for _ in range(len(model.outputs) *
- current_strategy.num_towers)]
- if model.uses_learning_phase and not isinstance(K.learning_phase(), int):
- ins = dataset_inputs + dataset_targets + sample_weights + [0]
- else:
- ins = dataset_inputs + dataset_targets
+ # We need to set sample_weights to None since there are sample weight
+ # placeholders that are created with default values.
+ sample_weights = [None for _ in range(len(model.outputs) *
+ current_strategy.num_towers)]
+ if model.uses_learning_phase and not isinstance(K.learning_phase(), int):
+ ins = dataset_inputs + dataset_targets + sample_weights + [0]
+ else:
+ ins = dataset_inputs + dataset_targets
- outs = []
- if verbose == 1:
- progbar = Progbar(target=steps)
+ for m in model.stateful_metric_functions:
+ m.reset_states()
+ stateful_metric_indices = [
+ i for i, name in enumerate(model.metrics_names)
+ if str(name) in model.stateful_metric_names
+ ]
- # Copy the weights from the original model to each of the replicated models.
- orig_model_weights = model.get_weights()
- with current_strategy.scope():
+ outs = []
+ if verbose == 1:
+ progbar = Progbar(target=steps)
+
+ # Copy the weights from the original model to each of the replicated models.
+ orig_model_weights = model.get_weights()
distributed_model = current_strategy.unwrap(model._grouped_model)[0]
distributed_training_utils.set_weights(
current_strategy, distributed_model, orig_model_weights)
- assert steps is not None
- for step in range(steps):
- batch_outs = distributed_test_function(ins)
- batch_outs = _aggregate_metrics_across_towers(
- current_strategy.num_towers, model.metrics_names, batch_outs)
- if isinstance(batch_outs, list):
- if step == 0:
- outs = [0.] * len(batch_outs)
- for i, batch_out in enumerate(batch_outs):
- outs[i] += batch_out
- else:
- if step == 0:
- outs.append(0.)
- outs[0] += batch_outs
- if verbose >= 1:
- progbar.update(step + 1)
- for i in range(len(outs)):
- outs[i] /= steps
+ assert steps is not None
+ for step in range(steps):
+ batch_outs = distributed_test_function(ins)
+ batch_outs = _aggregate_metrics_across_towers(
+ current_strategy.num_towers, model.metrics_names,
+ model.stateful_metric_names, batch_outs)
+ if isinstance(batch_outs, list):
+ if step == 0:
+ outs = [0.] * len(batch_outs)
+ for i, batch_out in enumerate(batch_outs):
+ if i in stateful_metric_indices:
+ outs[i] = batch_out
+ else:
+ outs[i] += batch_out
+ else:
+ if step == 0:
+ outs.append(0.)
+ outs[0] += batch_outs
+ if verbose >= 1:
+ progbar.update(step + 1)
+ for i in range(len(outs)):
+ if i not in stateful_metric_indices:
+ outs[i] /= steps
- if len(outs) == 1:
- return outs[0]
- return outs
+ if len(outs) == 1:
+ return outs[0]
+ return outs
def _experimental_test_loop(model, iterator, verbose=0, steps=None):
@@ -630,51 +645,50 @@
dataset_inputs = distributed_training_utils.flatten_perdevice_values(
current_strategy, inputs)
- distributed_predict_function = K.Function(
- all_inputs, all_outputs,
- updates=all_updates,
- name='distributed_predict_function',
- **all_session_args)
+ distributed_predict_function = K.Function(
+ all_inputs, all_outputs,
+ updates=all_updates,
+ name='distributed_predict_function',
+ **all_session_args)
- if model.uses_learning_phase and not isinstance(K.learning_phase(), int):
- ins = dataset_inputs + [0]
- else:
- ins = dataset_inputs
+ if model.uses_learning_phase and not isinstance(K.learning_phase(), int):
+ ins = dataset_inputs + [0]
+ else:
+ ins = dataset_inputs
- if verbose == 1:
- progbar = Progbar(target=steps)
+ if verbose == 1:
+ 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():
+ # Copy the weights from the original model to each of the replicated models.
+ orig_model_weights = model.get_weights()
distributed_model = current_strategy.unwrap(model._grouped_model)[0]
distributed_training_utils.set_weights(
current_strategy, distributed_model, orig_model_weights)
- if 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
- # and concatenate them upon returning.
- unconcatenated_outs = []
- for step in range(steps):
- batch_outs = distributed_predict_function(ins)
- if not isinstance(batch_outs, list):
- batch_outs = [batch_outs]
- if step == 0:
- for _ in batch_outs:
- unconcatenated_outs.append([])
- # TODO(anjalisridhar): Should combine the outputs from multiple towers
- # correctly here.
- for i, batch_out in enumerate(batch_outs):
- unconcatenated_outs[i].append(batch_out)
- if verbose >= 1:
- progbar.update(step + 1)
- if len(unconcatenated_outs) == 1:
- return np.concatenate(unconcatenated_outs[0], axis=0)
- return [
- np.concatenate(unconcatenated_outs[i], axis=0)
- for i in range(len(unconcatenated_outs))
- ]
+ if 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 and concatenate them upon returning.
+ unconcatenated_outs = []
+ for step in range(steps):
+ batch_outs = distributed_predict_function(ins)
+ if not isinstance(batch_outs, list):
+ batch_outs = [batch_outs]
+ if step == 0:
+ for _ in batch_outs:
+ unconcatenated_outs.append([])
+ # TODO(anjalisridhar): Should combine the outputs from multiple towers
+ # correctly here.
+ for i, batch_out in enumerate(batch_outs):
+ unconcatenated_outs[i].append(batch_out)
+ if verbose >= 1:
+ progbar.update(step + 1)
+ if len(unconcatenated_outs) == 1:
+ return np.concatenate(unconcatenated_outs[0], axis=0)
+ return [
+ np.concatenate(unconcatenated_outs[i], axis=0)
+ for i in range(len(unconcatenated_outs))
+ ]
def _experimental_predict_loop(model, iterator, verbose=0, steps=None):
@@ -816,10 +830,10 @@
cloned_model.compile(
optimizer,
model.loss,
- metrics=model.metrics,
+ metrics=metrics_module.clone_metrics(model.metrics),
loss_weights=model.loss_weights,
sample_weight_mode=model.sample_weight_mode,
- weighted_metrics=model.weighted_metrics,
+ weighted_metrics=metrics_module.clone_metrics(model.weighted_metrics),
target_tensors=targets)
return cloned_model
@@ -834,8 +848,9 @@
model._make_callback_model()
-def _aggregate_metrics_across_towers(num_devices, out_labels, outs):
- """Aggregate metrics values across all towers.
+def _aggregate_metrics_across_towers(num_devices, out_labels,
+ stateful_metric_names, outs):
+ """Aggregates stateless metrics values across towers.
When using `MirroredStrategy`, the number of towers is equal to the
number of devices over which training is distributed. This may not always be
@@ -844,6 +859,7 @@
Args:
num_devices: Number of devices over which the model is being distributed.
out_labels: The list of metric names passed to `compile`.
+ stateful_metric_names: List of stateful metric names on the model.
outs: The output from all the towers.
Returns:
@@ -858,10 +874,16 @@
# Each label in `out_labels` corresponds to one set of metrics. The
# number of metric values corresponds to the number of devices. We
# currently take the mean of the values.
- for _ in out_labels[1:]:
- m = np.mean(outs[current_index:current_index + num_devices])
- merged_output.append(m)
- current_index += num_devices
+ for metric_name in out_labels[1:]:
+ if metric_name in stateful_metric_names:
+ # For stateful metrics, we get one aggregated result value.
+ merged_output.append(outs[current_index])
+ current_index += 1
+ else:
+ m = np.mean(outs[current_index:current_index + num_devices])
+ merged_output.append(m)
+ current_index += num_devices
+
return merged_output
diff --git a/tensorflow/python/keras/engine/training_eager_test.py b/tensorflow/python/keras/engine/training_eager_test.py
index db7ccb1..1f5176c 100644
--- a/tensorflow/python/keras/engine/training_eager_test.py
+++ b/tensorflow/python/keras/engine/training_eager_test.py
@@ -192,6 +192,20 @@
history = model.fit(iterator, epochs=1, steps_per_epoch=10)
self.assertEqual(np.around(history.history['loss'][-1], decimals=4), 0.6173)
+ def test_no_loss_in_call(self):
+
+ class HasLoss(keras.layers.Layer):
+
+ def call(self, x):
+ self.add_loss(x)
+ return x
+
+ layer = HasLoss()
+ with self.assertRaises(RuntimeError):
+ layer(1.)
+
+ with ops.Graph().as_default():
+ layer(1.)
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 413c1f4..2e07469 100644
--- a/tensorflow/python/keras/engine/training_generator.py
+++ b/tensorflow/python/keras/engine/training_generator.py
@@ -21,6 +21,7 @@
import numpy as np
+from tensorflow.python.eager import context
from tensorflow.python.keras import callbacks as cbks
from tensorflow.python.keras.utils.data_utils import GeneratorEnqueuer
from tensorflow.python.keras.utils.data_utils import OrderedEnqueuer
@@ -48,6 +49,10 @@
epoch = initial_epoch
do_validation = bool(validation_data)
+ if not context.executing_eagerly():
+ model._make_train_function()
+ if do_validation:
+ model._make_test_function()
is_sequence = isinstance(generator, Sequence)
if not is_sequence and use_multiprocessing and workers > 1:
@@ -233,6 +238,9 @@
use_multiprocessing=False,
verbose=0):
"""See docstring for `Model.evaluate_generator`."""
+ if not context.executing_eagerly():
+ model._make_test_function()
+
if hasattr(model, 'metrics'):
for m in model.stateful_metric_functions:
m.reset_states()
@@ -342,6 +350,9 @@
use_multiprocessing=False,
verbose=0):
"""See docstring for `Model.predict_generator`."""
+ if not context.executing_eagerly():
+ model._make_test_function()
+
steps_done = 0
wait_time = 0.01
all_outs = []
diff --git a/tensorflow/python/keras/engine/training_test.py b/tensorflow/python/keras/engine/training_test.py
index 30be413..54ad74c 100644
--- a/tensorflow/python/keras/engine/training_test.py
+++ b/tensorflow/python/keras/engine/training_test.py
@@ -27,6 +27,7 @@
from tensorflow.python import keras
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.eager import context
+from tensorflow.python.eager import function
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import test_util as tf_test_util
@@ -2427,6 +2428,17 @@
scores = model.train_on_batch(x, y, sample_weight=w)
self.assertArrayNear(scores, [0.2, 0.8, 0.8], 0.1)
+ def test_losses_in_defun(self):
+ with context.eager_mode():
+ layer = keras.layers.Dense(1, kernel_regularizer='l1')
+ layer(array_ops.ones([1, 10]))
+
+ @function.defun
+ def get_losses():
+ return layer.losses
+
+ self.assertAllEqual(self.evaluate(layer.losses),
+ self.evaluate(get_losses()))
if __name__ == '__main__':
test.main()
diff --git a/tensorflow/python/keras/layers/core.py b/tensorflow/python/keras/layers/core.py
index 4032202..efa2195 100644
--- a/tensorflow/python/keras/layers/core.py
+++ b/tensorflow/python/keras/layers/core.py
@@ -671,22 +671,34 @@
if mask is not None:
self.supports_masking = True
self.mask = mask
- if output_shape is None:
- self._output_shape = None
- elif isinstance(output_shape, (tuple, list)):
- self._output_shape = tuple(output_shape)
- else:
- if not callable(output_shape):
- raise TypeError('In Lambda, `output_shape` '
- 'must be a list, a tuple, or a function.')
- self._output_shape = output_shape
+ if (output_shape is not None and not isinstance(output_shape,
+ (tuple, list)) and
+ not callable(output_shape)):
+ raise TypeError('In Lambda, `output_shape` '
+ 'must be a list, a tuple, or a function.')
+ # Convert a list representing a single shape into a tuple.
+ if (isinstance(output_shape, list) and isinstance(output_shape[0],
+ (int, type(None)))):
+ output_shape = tuple(output_shape)
+ self._output_shape = output_shape
@tf_utils.shape_type_conversion
def compute_output_shape(self, input_shape):
if self._output_shape is None:
if context.executing_eagerly():
- raise NotImplementedError
- x = K.placeholder(shape=input_shape)
+ # Make use of existing autocomputation for Eager mode but provide
+ # Lambda-specific error message.
+ try:
+ return super(Lambda, self).compute_output_shape(input_shape)
+ except NotImplementedError:
+ raise NotImplementedError('We could not automatically infer '
+ 'the static shape of the Lambda\'s output.'
+ ' Please specify the `output_shape` for'
+ ' this Lambda.')
+ if isinstance(input_shape, list):
+ x = [K.placeholder(shape=shape) for shape in input_shape]
+ else:
+ x = K.placeholder(shape=input_shape)
x = self.call(x)
if isinstance(x, list):
return [tensor_shape.TensorShape(K.int_shape(x_elem)) for x_elem in x]
@@ -697,16 +709,27 @@
num_samples = input_shape[0][0]
else:
num_samples = input_shape[0] if input_shape else None
- return tensor_shape.TensorShape((num_samples,) +
- tuple(self._output_shape))
+ # List here represents multiple outputs.
+ if isinstance(self._output_shape, list):
+ return [
+ tensor_shape.TensorShape((num_samples,) + tuple(single_shape))
+ for single_shape in self._output_shape
+ ]
+ return tensor_shape.TensorShape((num_samples,) + self._output_shape)
else:
shape = self._output_shape(input_shape)
if not isinstance(shape, (list, tuple)):
raise ValueError(
'`output_shape` function must return a tuple or a list of tuples.')
+ # List here can represent multiple outputs or single output.
if isinstance(shape, list):
- if isinstance(shape[0], int) or shape[0] is None:
+ # Convert list representing single output into a tuple.
+ if isinstance(shape[0], (int, type(None))):
shape = tuple(shape)
+ else:
+ return [
+ tensor_shape.TensorShape(single_shape) for single_shape in shape
+ ]
return tensor_shape.TensorShape(shape)
def call(self, inputs, mask=None):
diff --git a/tensorflow/python/keras/layers/core_test.py b/tensorflow/python/keras/layers/core_test.py
index 1df1d57..f0fea1f 100644
--- a/tensorflow/python/keras/layers/core_test.py
+++ b/tensorflow/python/keras/layers/core_test.py
@@ -252,6 +252,51 @@
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))
+ # Dimension(None) != Dimension(None), so check
+ # str representations for equality.
+ self.assertAllEqual(('5', '?', '10'), tuple([str(s) for s in output_shape]))
+
+ @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
diff --git a/tensorflow/python/keras/metrics.py b/tensorflow/python/keras/metrics.py
index e64241e..f4e8419 100644
--- a/tensorflow/python/keras/metrics.py
+++ b/tensorflow/python/keras/metrics.py
@@ -71,6 +71,22 @@
name, x))
+def clone_metric(metric):
+ """Returns a clone of the metric if stateful, otherwise returns it as is."""
+ if isinstance(metric, Metric):
+ return metric.__class__.from_config(metric.get_config())
+ return metric
+
+
+def clone_metrics(metrics):
+ """Clones the given metric list/dict."""
+ if metrics is None:
+ return None
+ if isinstance(metrics, dict):
+ return {key: clone_metric(value) for key, value in metrics.items()}
+ return [clone_metric(metric) for metric in metrics]
+
+
def update_state_wrapper(update_state_fn):
"""Decorator to wrap metric `update_state()` with `add_update()`.
diff --git a/tensorflow/python/keras/models.py b/tensorflow/python/keras/models.py
index 41c5e3c..b04b4df 100644
--- a/tensorflow/python/keras/models.py
+++ b/tensorflow/python/keras/models.py
@@ -20,6 +20,7 @@
from __future__ import print_function
from tensorflow.python.keras import backend as K
+from tensorflow.python.keras import metrics as metrics_module
from tensorflow.python.keras import optimizers
from tensorflow.python.keras.engine import saving
from tensorflow.python.keras.engine import sequential
@@ -290,7 +291,9 @@
if isinstance(value, Layer):
attributes_cache[name] = value
assert value in model._layers
- elif isinstance(value, (list, tuple)) and name not in ('layers', '_layers'):
+ elif isinstance(
+ value, (list, tuple)) and name not in ('layers', '_layers',
+ 'stateful_metric_functions'):
# Handle case: list/tuple of layers (also tracked by the Network API).
if value and all(isinstance(val, Layer) for val in value):
raise ValueError('We do not support the use of list-of-layers '
@@ -466,10 +469,10 @@
clone.compile(
optimizer,
model.loss,
- metrics=model.metrics,
+ metrics=metrics_module.clone_metrics(model.metrics),
loss_weights=model.loss_weights,
sample_weight_mode=model.sample_weight_mode,
- weighted_metrics=model.weighted_metrics,
+ weighted_metrics=metrics_module.clone_metrics(model.weighted_metrics),
target_tensors=target_tensors)
return clone
diff --git a/tensorflow/python/keras/preprocessing/image_test.py b/tensorflow/python/keras/preprocessing/image_test.py
index 362cbc1..4abaadf 100644
--- a/tensorflow/python/keras/preprocessing/image_test.py
+++ b/tensorflow/python/keras/preprocessing/image_test.py
@@ -94,43 +94,6 @@
self.assertEqual(x.shape[1:], images.shape[1:])
break
- def test_image_data_generator_with_validation_split(self):
- if PIL is None:
- return # Skip test if PIL is not available.
-
- for test_images in _generate_test_images():
- img_list = []
- for im in test_images:
- img_list.append(keras.preprocessing.image.img_to_array(im)[None, ...])
-
- images = np.vstack(img_list)
- generator = keras.preprocessing.image.ImageDataGenerator(
- validation_split=0.5)
- seq = generator.flow(
- images,
- np.arange(images.shape[0]),
- shuffle=False,
- batch_size=3,
- subset='validation')
- _, y = seq[0]
- self.assertEqual(list(y), [0, 1, 2])
- seq = generator.flow(
- images,
- np.arange(images.shape[0]),
- shuffle=False,
- batch_size=3,
- subset='training')
- _, y2 = seq[0]
- self.assertEqual(list(y2), [4, 5, 6])
-
- with self.assertRaises(ValueError):
- generator.flow(
- images,
- np.arange(images.shape[0]),
- shuffle=False,
- batch_size=3,
- subset='foo')
-
def test_image_data_generator_with_split_value_error(self):
with self.assertRaises(ValueError):
keras.preprocessing.image.ImageDataGenerator(validation_split=5)
diff --git a/tensorflow/python/kernel_tests/BUILD b/tensorflow/python/kernel_tests/BUILD
index 5183e4d..9490746 100644
--- a/tensorflow/python/kernel_tests/BUILD
+++ b/tensorflow/python/kernel_tests/BUILD
@@ -1097,6 +1097,18 @@
],
)
+tf_py_test(
+ name = "unicode_script_op_test",
+ size = "small",
+ srcs = ["unicode_script_op_test.py"],
+ additional_deps = [
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:constant_op",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:string_ops",
+ ],
+)
+
cuda_py_test(
name = "topk_op_test",
size = "small",
@@ -1468,7 +1480,7 @@
name = "control_flow_ops_py_test",
# TODO(b/70473603): change this back to "small" once the C API is
# permanently enabled
- size = "medium",
+ size = "large",
srcs = ["control_flow_ops_py_test.py"],
additional_deps = [
"//third_party/py/numpy",
@@ -1500,6 +1512,7 @@
"//tensorflow/python:util",
"//tensorflow/python:variable_scope",
"//tensorflow/python:variables",
+ "//tensorflow/python:while_v2",
],
)
@@ -2346,7 +2359,7 @@
cuda_py_test(
name = "transpose_op_test",
- size = "large",
+ size = "medium",
srcs = ["transpose_op_test.py"],
additional_deps = [
"//third_party/py/numpy",
@@ -2354,10 +2367,11 @@
"//tensorflow/python:client_testlib",
"//tensorflow/python:framework_for_generated_wrappers",
],
- shard_count = 2,
+ shard_count = 4,
tags = [
"no_gpu",
"no_oss",
+ "optonly", # times out
],
)
@@ -2476,6 +2490,7 @@
"//tensorflow/python:nn_grad",
"//tensorflow/python:nn_ops",
],
+ shard_count = 2,
tags = [
"optonly", # flaky timeouts unless optimized
],
@@ -2496,7 +2511,7 @@
cuda_py_test(
name = "conv_ops_test",
- size = "large",
+ size = "medium",
srcs = ["conv_ops_test.py"],
additional_deps = [
"//third_party/py/numpy",
@@ -2515,6 +2530,9 @@
"//tensorflow/python:variables",
],
shard_count = 4,
+ tags = [
+ "optonly", # times out
+ ],
)
cuda_py_test(
@@ -2574,7 +2592,7 @@
cuda_py_test(
name = "fft_ops_test",
- size = "large",
+ size = "medium",
srcs = ["fft_ops_test.py"],
additional_deps = [
"//third_party/py/numpy",
@@ -2584,7 +2602,8 @@
"//tensorflow/python:spectral_ops",
"//tensorflow/python:spectral_ops_test_util",
],
- shard_count = 3,
+ shard_count = 4,
+ tags = ["optonly"],
)
cuda_py_test(
@@ -2649,7 +2668,7 @@
cuda_py_test(
name = "scatter_ops_test",
- size = "large", # NOTE: This is not run by default.
+ size = "medium", # NOTE: This is not run by default.
srcs = ["scatter_ops_test.py"],
additional_deps = [
"//third_party/py/numpy",
@@ -2658,11 +2677,13 @@
"//tensorflow/python:state_ops",
"//tensorflow/python:variables",
],
+ shard_count = 2,
+ tags = ["optonly"],
)
cuda_py_test(
name = "slice_op_test",
- size = "large",
+ size = "medium",
srcs = ["slice_op_test.py"],
additional_deps = [
"//third_party/py/numpy",
@@ -3245,8 +3266,7 @@
tags = ["no_gpu"], # TODO(b/111656070)
)
-# TODO(b/116053459): Replace with cuda_py_test.
-tf_py_test(
+cuda_py_test(
name = "while_v2_test",
size = "medium",
srcs = ["while_v2_test.py"],
@@ -3266,5 +3286,4 @@
"//tensorflow/python:while_v2",
],
grpc_enabled = True,
- tags = ["no_gpu"], # TODO(b/116053459)
)
diff --git a/tensorflow/python/kernel_tests/array_ops_test.py b/tensorflow/python/kernel_tests/array_ops_test.py
index 2fe8583..dcc5947 100644
--- a/tensorflow/python/kernel_tests/array_ops_test.py
+++ b/tensorflow/python/kernel_tests/array_ops_test.py
@@ -615,6 +615,14 @@
_ = checker[:, 0]
_ = checker[:, :, 0]
+ def testBothNewAxisAndShrink(self):
+ with self.test_session(use_gpu=True):
+ ones = array_ops.placeholder(shape=[2, 2], dtype=dtypes.int16)
+ self.assertAllEqual(
+ ones[array_ops.newaxis, :, 0].eval(
+ feed_dict={ones: [[1, 1], [1, 1]]}),
+ [[1, 1]])
+
def testTensorIndexing(self):
with self.test_session(use_gpu=True):
raw = [[[[[1, 2, 4, 5], [5, 6, 7, 8], [9, 10, 11, 12]]],
@@ -1001,14 +1009,14 @@
errors.FailedPreconditionError,
"Attempting to use uninitialized value Variable"):
with self.cached_session() as sess:
- v = variables.Variable([1, 2])
+ v = variables.VariableV1([1, 2])
sess.run(v[:].assign([1, 2]))
def testTypeError(self):
init_val = constant_op.constant([1, 2], dtype=dtypes.int32)
too_small_val = constant_op.constant([3, 4], dtype=dtypes.int8)
too_large_val = constant_op.constant([3, 4], dtype=dtypes.int64)
- v = variables.Variable(init_val)
+ v = variables.VariableV1(init_val)
with self.assertRaises(TypeError):
v[:].assign(too_small_val)
with self.assertRaises(TypeError):
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 fc4d2a3..ae61be6 100644
--- a/tensorflow/python/kernel_tests/control_flow_ops_py_test.py
+++ b/tensorflow/python/kernel_tests/control_flow_ops_py_test.py
@@ -23,7 +23,6 @@
import collections
import math
import time
-import unittest
import numpy as np
from six.moves import xrange # pylint: disable=redefined-builtin
@@ -32,6 +31,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 function as eager_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors_impl
@@ -63,6 +63,7 @@
from tensorflow.python.ops import state_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import variables
+from tensorflow.python.ops import while_v2 # pylint: disable=unused-import
# pylint: disable=unused-import
import tensorflow.python.ops.tensor_array_grad
# pylint: enable=unused-import
@@ -125,12 +126,12 @@
return r_s
-@test_util.with_cond_v2
+@test_util.with_control_flow_v2
class ControlFlowTest(test.TestCase):
def testRefIdentity(self):
with self.cached_session():
- v = variables.Variable(7)
+ v = variables.VariableV1(7)
v = control_flow_ops._Identity(v)
op = state_ops.assign(v, 9)
@@ -142,7 +143,7 @@
def testRefEnter(self):
with self.cached_session():
- v = variables.Variable(7)
+ v = variables.VariableV1(7)
enter_v = control_flow_ops._Enter(v, "foo_1", is_constant=True)
nine = constant_op.constant(9)
@@ -155,7 +156,7 @@
def testRefSwitch(self):
with self.cached_session():
- v = variables.Variable(7)
+ v = variables.VariableV1(7)
p = constant_op.constant(True)
v1 = control_flow_ops._SwitchRefOrTensor(v._ref(), p) # pylint: disable=protected-access
@@ -332,10 +333,8 @@
with self.assertRaisesOpError("has inputs from different frames"):
res.eval(feed_dict={data: 1.0})
+ @test_util.disable_control_flow_v2("b/113294340")
def testCondBool(self):
- if control_flow_ops.ENABLE_COND_V2:
- return unittest.skip("b/113296297")
-
values = constant_op.constant(10)
fn1 = lambda: math_ops.add(values, 1)
fn2 = lambda: math_ops.subtract(values, 1)
@@ -366,6 +365,7 @@
"has been marked as not fetchable"):
sess.run(t, feed_dict={x: 3})
+ @test_util.disable_control_flow_v2("Not relevant")
def testFeedable(self):
with self.cached_session() as sess:
c = constant_op.constant(2)
@@ -383,10 +383,8 @@
with self.assertRaisesRegexp(ValueError, "may not be fed"):
sess.run(r, feed_dict={t: 3})
+ @test_util.disable_control_flow_v2("b/113296180 (IndexedSlices)")
def testCondIndexedSlices(self):
- if control_flow_ops.ENABLE_COND_V2:
- return unittest.skip("b/113296180")
-
with self.cached_session():
values = constant_op.constant(10)
indices = constant_op.constant(0)
@@ -401,10 +399,8 @@
self.assertAllEqual(11, val)
self.assertAllEqual(0, ind)
+ @test_util.disable_control_flow_v2("b/113296161 (SparseTensors)")
def testCondSparseTensor(self):
- if control_flow_ops.ENABLE_COND_V2:
- return unittest.skip("b/113296161 (SparseTensors)")
-
with self.cached_session():
values = constant_op.constant([2.0, 4.0], name="values")
indices = constant_op.constant(
@@ -435,10 +431,8 @@
self.assertEqual(1.0, control_flow_ops.cond(rv, case, lambda: t).eval())
+ @test_util.disable_control_flow_v2("b/113293074")
def testCondIndexedSlicesDifferentTypes(self):
- if control_flow_ops.ENABLE_COND_V2:
- return unittest.skip("b/113293074")
-
with self.cached_session():
values = constant_op.constant(10)
i_32 = ops.convert_to_tensor(0, name="one", dtype=dtypes.int32)
@@ -510,10 +504,8 @@
result = r.eval()
self.assertAllEqual(12, result)
+ @test_util.disable_control_flow_v2("b/113324949 (ref vars)")
def testCond_4(self):
- if control_flow_ops.ENABLE_COND_V2:
- return unittest.skip("b/113324949 (ref vars)")
-
with self.cached_session():
v1 = variables.Variable(7)
v2 = variables.Variable(7)
@@ -587,10 +579,8 @@
r = control_flow_ops.cond(constant_op.constant(False), true_fn, false_fn)
self.assertAllEqual([2.0], r.eval())
+ @test_util.disable_control_flow_v2("b/79881896 (control deps)")
def testCondWithControl(self):
- if control_flow_ops.ENABLE_COND_V2:
- return unittest.skip("b/79881896")
-
with self.cached_session():
control_holder = array_ops.placeholder(dtypes.float32, shape=())
a = constant_op.constant(3)
@@ -629,10 +619,9 @@
merged_op = control_flow_ops.merge([assign_v, orig_v])
self.assertAllEqual([1.0], sess.run(merged_op.output))
+ @test_util.disable_control_flow_v2(
+ "b/112477618 (Operation returned from cond)")
def testCondSwitchIdentity(self):
- if control_flow_ops.ENABLE_COND_V2:
- return unittest.skip("b/112477618 (Operation returned from cond)")
-
# Make sure the recv identity is not removed by optimization.
with session.Session(config=opt_cfg()) as sess:
pred = constant_op.constant(True)
@@ -646,10 +635,9 @@
r = control_flow_ops.cond(pred, fn1, fn2)
sess.run(r)
+ @test_util.disable_control_flow_v2(
+ "b/112477618 (Operation returned from cond)")
def testCondRecvIdentity(self):
- if control_flow_ops.ENABLE_COND_V2:
- return unittest.skip("b/112477618 (Operation returned from cond)")
-
# Make sure the switch identity is not removed by optimization.
with session.Session(config=opt_cfg()) as sess:
with ops.device(test.gpu_device_name()):
@@ -665,10 +653,8 @@
r = control_flow_ops.cond(pred, fn1, fn2)
sess.run(r)
+ @test_util.disable_control_flow_v2("b/113346829 (gpu failure)")
def testCondGrad_1(self):
- if control_flow_ops.ENABLE_COND_V2:
- return unittest.skip("b/113346829 (gpu failure)")
-
graph = ops.Graph()
with graph.as_default():
x = constant_op.constant(10.0, name="x")
@@ -694,10 +680,9 @@
self.assertAllEqual(42.0, grad.eval(feed_dict={c: 1}))
self.assertAllEqual(3.0, grad.eval(feed_dict={c: 3}))
+ @test_util.disable_control_flow_v2(
+ "b/110550782 (gradient w.r.t external variable)")
def testCondGrad_3(self):
- if control_flow_ops.ENABLE_COND_V2:
- return unittest.skip("b/110550782 (gradient w.r.t external variable)")
-
with self.cached_session():
c = array_ops.placeholder(dtypes.int32, shape=[])
ox = constant_op.constant(10.0)
@@ -729,10 +714,8 @@
result = gradients_impl.gradients(z, x)[0]
self.assertEqual(1.0, result.eval())
+ @test_util.disable_control_flow_v2("b/113327884")
def testCondGrad_Gather(self):
- if control_flow_ops.ENABLE_COND_V2:
- return unittest.skip("b/113327884")
-
with self.cached_session() as sess:
v1 = variables.Variable([1.0, 42.0])
c = array_ops.placeholder(dtypes.int32, shape=[])
@@ -756,6 +739,7 @@
self.assertAllEqual(dense_gv, [0.0, 2.0])
# Microbenchmark: 256,000 iterations/s.
+ @test_util.disable_control_flow_v2("b/116630618 (Times out)")
def testWhile_1(self):
with self.cached_session():
n = constant_op.constant(0)
@@ -764,6 +748,7 @@
r = control_flow_ops.while_loop(c, b, [n], parallel_iterations=20)
self.assertEqual(10000, r.eval())
+ @test_util.disable_control_flow_v2("b/79881896 (control deps)")
def testWhileExternalControlDependencies(self):
with self.cached_session():
v = variables.Variable(0.0)
@@ -779,6 +764,7 @@
self.assertAllEqual(result.eval(), 2)
self.assertAllEqual(v.eval(), 1.0)
+ @test_util.disable_control_flow_v2("b/79881896 (control deps)")
def testWhileExternalControlDependenciesNoInput(self):
with self.cached_session():
v = variables.Variable(0.0)
@@ -794,9 +780,10 @@
result.eval()
self.assertAllEqual(v.eval(), 1.0)
+ @test_util.disable_control_flow_v2("b/113324949 (RefVariable)")
def testWhileWithRefs_1(self):
with self.cached_session() as sess:
- x = variables.Variable(0)._ref() # pylint: disable=protected-access
+ x = variables.VariableV1(0)._ref() # pylint: disable=protected-access
i = constant_op.constant(0)
c = lambda i, x: math_ops.less(i, 100)
@@ -824,18 +811,22 @@
r = isum(s)
self.assertAllEqual(45, r.eval())
+ @test_util.disable_control_flow_v2("b/115776323 (max_iters)")
def testWhileWithMaximumIterations(self):
with self.cached_session():
s = constant_op.constant([1, 2, 3, 4, 5])
r = isum(s, maximum_iterations=3)
self.assertAllEqual([1 + 3, 2 + 3, 3 + 3, 4 + 3, 5 + 3], r.eval())
+ @test_util.disable_control_flow_v2("b/116339888 (non-tensor loop var)")
def testWhileWithMaximumIterationsAndSingleArgument(self):
with self.cached_session():
r = control_flow_ops.while_loop(
lambda i: i < 3, lambda i: i + 1, [0], maximum_iterations=1)
self.assertEqual(1, r.eval())
+ @test_util.disable_control_flow_v2(
+ "b/116248044 (nested), b/115920078 (gradients)")
def testSingleNestedMaximumIterationsWhileLoopGradientInXLAContext(self):
v = constant_op.constant(1.0)
@@ -861,6 +852,7 @@
# Should execute without issue.
self.assertEqual(3, self.evaluate(loop_execute))
+ @test_util.disable_control_flow_v2("b/116248044 (nested while_loop)")
def testInvalidMaximumIterationsWhileLoopGradientInXLAContext(self):
v = constant_op.constant(1.0)
@@ -904,10 +896,8 @@
r"context '.*' \(currently defined in '.*'\)"):
_ = gradients_impl.gradients(loop_with_maxiter, v)
+ @test_util.disable_control_flow_v2("b/115776323 (max_iters)")
def testInvalidMaximumIterationsFromSiblingContextWhileLoopInXLAContext(self):
- if control_flow_ops.ENABLE_COND_V2:
- return unittest.skip("b/113294340 (enable while_v2)")
-
v = constant_op.constant(1.0)
def create_while_loop():
@@ -939,6 +929,8 @@
r"while loop context '' \(currently defined in 'cond/.+'\)"):
_ = gradients_impl.gradients(loop, v)
+ @test_util.disable_control_flow_v2(
+ "b/116248044 (nesting), b/115776323 (max_iters)")
def testNestedWhileLoopWithMaxItersFromOuterContextInXLAContext(self):
v = constant_op.constant(1.0)
@@ -1048,6 +1040,7 @@
result = r[3].eval()
self.assertAllEqual(42, result)
+ @test_util.disable_control_flow_v2("b/116283162 (shape_invariants)")
def testWhile_5(self):
with self.cached_session():
@@ -1072,6 +1065,7 @@
result = r[2].eval()
self.assertAllEqual(np.array([0, 1, 2, 3, 4, 5, 6]), result)
+ @test_util.disable_control_flow_v2("b/116338794 (buffer_reuse)")
def testBufferForwarding(self):
run_options = config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE)
@@ -1122,6 +1116,7 @@
self._testWhile_Gpu_1(use_gpu=False)
self._testWhile_Gpu_1(use_gpu=True)
+ @test_util.disable_control_flow_v2("b/116283162 (shape_invariants)")
def testWhileShape(self):
with self.cached_session():
i = constant_op.constant(0)
@@ -1139,6 +1134,7 @@
r = r[1] * array_ops.ones([8, 8])
self.assertAllEqual(np.ones((8, 8)), r.eval())
+ @test_util.disable_control_flow_v2("b/116339888 (non-tensor loop var)")
def testWhileWithNonTensorInput_Scalar(self):
with self.cached_session():
n = 0
@@ -1147,6 +1143,7 @@
r = control_flow_ops.while_loop(c, b, [n], parallel_iterations=20)
self.assertEqual(10000, r.eval())
+ @test_util.disable_control_flow_v2("b/116339888 (non-tensor loop var)")
def testWhileWithNonTensorInput_Vector(self):
with self.cached_session():
n = np.array([0]) # Note, [0] would not work here; that is a list
@@ -1155,6 +1152,7 @@
r = control_flow_ops.while_loop(c, b, [n], parallel_iterations=20)
self.assertEqual([10000], r.eval())
+ @test_util.disable_control_flow_v2("b/116283162 (shape_invariants)")
def testWhileShapeInference(self):
with self.cached_session():
i = constant_op.constant(0)
@@ -1169,7 +1167,7 @@
r = control_flow_ops.while_loop(
c, b, [i, m],
[i.get_shape(), tensor_shape.TensorShape([None, 2])])
- self.assertTrue(r[1].get_shape()[0].value is None)
+ self.assertIsNone(r[1].get_shape()[0].value)
self.assertEqual(r[1].get_shape()[1], tensor_shape.Dimension(2))
with self.assertRaisesRegexp(
@@ -1180,6 +1178,7 @@
r"tf.while_loop to specify a less-specific shape."):
r = control_flow_ops.while_loop(c, b, [i, m])
+ @test_util.disable_control_flow_v2("b/116328420 (SparseTensor)")
def testWhileShapeInferenceSparseTensor(self):
with self.cached_session():
values = constant_op.constant([2.0, 4.0], name="values")
@@ -1211,6 +1210,7 @@
c, b, [i, x],
[i.get_shape(), tensor_shape.TensorShape([5])])
+ @test_util.disable_control_flow_v2("b/116282023 (IndexedSlices)")
def testWhileShapeInferenceIndexedSlices(self):
with self.cached_session():
values = constant_op.constant([[2.0, 4.0], [3.0, 5.0]], name="values")
@@ -1265,6 +1265,7 @@
r = control_flow_ops.while_loop(c, b, [n])
self.assertEqual(225, r.eval())
+ @test_util.disable_control_flow_v2("b/116248044 (nested while)")
def testNestedWhile_1(self):
self._testNestedWhile_1(use_gpu=False)
self._testNestedWhile_1(use_gpu=True)
@@ -1297,6 +1298,7 @@
outer_c, outer_b, [s0], parallel_iterations=1)
self.assertEqual(1048576.0, r.eval())
+ @test_util.disable_control_flow_v2("b/116248044 (nested while)")
def testNestedWhile_2(self):
self._testNestedWhile_2(use_gpu=False)
self._testNestedWhile_2(use_gpu=True)
@@ -1350,6 +1352,7 @@
lambda x: x < 10, lambda x: x + array_ops.identity(c), [x0])
self.assertEqual(10, sess.run(r, {b: True}))
+ @test_util.disable_control_flow_v2("b/79881896 (control_deps)")
def testWhileWithControl_5(self):
with self.cached_session() as sess:
b = array_ops.placeholder(dtypes.bool)
@@ -1364,9 +1367,6 @@
self.assertEqual(10, sess.run(r, {b: True}))
def testWhileCondWithControl(self):
- if control_flow_ops.ENABLE_COND_V2:
- return unittest.skip("b/113294377 (unknown shape)")
-
# Ensure that no control edges by an outer control dependency context are
# added to nodes inside cond/while contexts.
with self.cached_session() as sess:
@@ -1380,10 +1380,8 @@
(constant_op.constant(5),))
self.assertEqual(0, sess.run(loop))
+ @test_util.disable_control_flow_v2("b/113324949 (ref vars)")
def testWhileCondWithControl_1(self):
- if control_flow_ops.ENABLE_COND_V2:
- return unittest.skip("b/113324949 (ref vars)")
-
with self.cached_session():
v = variable_scope.get_variable(
"v", [], initializer=init_ops.constant_initializer(2))
@@ -1405,9 +1403,8 @@
self.assertEqual(4, r.eval())
self.assertAllClose(65536.0, v.eval())
+ @test_util.disable_control_flow_v2("b/113324949 (ref vars)")
def testWhileCondExitControl(self):
- if control_flow_ops.ENABLE_COND_V2:
- return unittest.skip("b/113294340 (enable while_v2)")
with self.cached_session():
v = variables.Variable(1)
@@ -1432,8 +1429,6 @@
self.assertEqual(99, v.eval())
def testCondWhile_1(self):
- if control_flow_ops.ENABLE_COND_V2:
- return unittest.skip("b/113294340 (enable while_v2)")
with self.cached_session():
n = ops.convert_to_tensor(0, name="n")
@@ -1445,8 +1440,6 @@
self.assertAllEqual(10, r.eval())
def testCondWhile_2(self):
- if control_flow_ops.ENABLE_COND_V2:
- return unittest.skip("b/113294340 (enable while_v2)")
with self.cached_session():
n = ops.convert_to_tensor(0)
@@ -1458,9 +1451,6 @@
self.assertAllEqual(10, r.eval())
def _testCondWhile_3(self, use_gpu):
- if control_flow_ops.ENABLE_COND_V2:
- return unittest.skip("b/113294340 (enable while_v2)")
-
with self.test_session(use_gpu=use_gpu) as sess:
p = array_ops.placeholder(dtypes.bool)
n = constant_op.constant(0.0)
@@ -1477,18 +1467,17 @@
lambda: control_flow_ops.while_loop(c, b, [n]),
lambda: math_ops.multiply(n, 2.0))
r1 = gradients_impl.gradients(r, [n])
- self.assertEqual(10, sess.run(r, {p: True}))
+ self.assertEqual(10., sess.run(r, {p: True}))
self.assertEqual([1.0], sess.run(r1, {p: True}))
self.assertEqual(0.0, sess.run(r, {p: False}))
self.assertEqual([2.0], sess.run(r1, {p: False}))
+ @test_util.disable_control_flow_v2("b/116743589")
def testCondWhile_3(self):
self._testCondWhile_3(use_gpu=False)
self._testCondWhile_3(use_gpu=True)
def testWhileCond_1(self):
- if control_flow_ops.ENABLE_COND_V2:
- return unittest.skip("b/113294377 (unknown shape)")
with self.cached_session():
i = ops.convert_to_tensor(0, name="i")
@@ -1505,8 +1494,6 @@
self.assertAllEqual(10, r.eval())
def testWhileCond_2(self):
- if control_flow_ops.ENABLE_COND_V2:
- return unittest.skip("b/113294377 (unknown shape)")
with self.cached_session():
n = ops.convert_to_tensor(0, name="n")
@@ -1516,8 +1503,6 @@
self.assertAllEqual(10, r.eval())
def testWhileCond_3(self):
- if control_flow_ops.ENABLE_COND_V2:
- return unittest.skip("b/113294377 (unknown shape)")
with self.cached_session():
n = ops.convert_to_tensor(0)
@@ -1532,6 +1517,7 @@
self.assertAllEqual(10, r.eval())
# NOTE: It is ok to have parallel_iterations > 1
+ @test_util.disable_control_flow_v2("b/113324949 (RefVariable)")
def testWhileUpdateVariable_1(self):
with self.cached_session():
select = variables.Variable([3.0, 4.0, 5.0])
@@ -1554,6 +1540,7 @@
result = select.eval()
self.assertAllClose(np.array([10.0, 10.0, 10.0]), result)
+ @test_util.disable_control_flow_v2("b/113324949 (RefVariable)")
def testWhileUpdateVariable_2(self):
with self.cached_session():
select1 = variables.Variable([3.0, 4.0, 5.0])
@@ -1580,6 +1567,7 @@
result2 = select2.eval()
self.assertAllClose(np.array([10.0, 10.0, 10.0]), result2)
+ @test_util.disable_control_flow_v2("b/113324949 (RefVariable)")
def testWhileUpdateVariable_3(self):
with self.cached_session():
select = variables.Variable([3.0, 4.0, 5.0])
@@ -1601,7 +1589,7 @@
result = r[1].eval()
self.assertAllClose(np.array([10.0, 10.0, 10.0]), result)
- # b/24814703
+ @test_util.disable_control_flow_v2("b/113324949 (RefVariable)")
def testWhileUpdateVariable_4(self):
with self.cached_session():
var_a = variables.Variable(0, name="a")
@@ -1629,7 +1617,7 @@
lpa.eval() # Run the loop
self.assertEqual(10, var_b.eval())
- # b/24736492
+ @test_util.disable_control_flow_v2("b/113324949 (RefVariable)")
def testWhileUpdateVariable_5(self):
with self.cached_session():
# Create some variables.
@@ -1659,7 +1647,7 @@
self.assertEqual(10, var_a.eval())
self.assertEqual(10, var_b.eval())
- # b/24814668
+ @test_util.disable_control_flow_v2("b/113324949 (RefVariable)")
def testWhileUpdateVariable_6(self):
with self.cached_session():
# Create some variables.
@@ -1689,6 +1677,7 @@
self.assertEqual(55, var_b.eval())
self.assertEqual(10, var_a.eval())
+ @test_util.disable_control_flow_v2("b/116742472 (resource accumulator)")
def testWhileQueue_1(self):
with self.cached_session():
q = data_flow_ops.FIFOQueue(-1, dtypes.int32)
@@ -1707,6 +1696,7 @@
for i in xrange(10):
self.assertEqual([i], q.dequeue().eval())
+ @test_util.disable_control_flow_v2("b/116283162 (shape_invariants)")
def testWhileStack_1(self):
with self.cached_session():
s = gen_data_flow_ops.stack_v2(-1, dtypes.int32, stack_name="foo")
@@ -1775,6 +1765,7 @@
with self.session(graph=graph) as sess:
self.assertAllClose(1024.0, sess.run(r))
+ @test_util.disable_control_flow_v2("b/116351701 (colocation)")
def testWhileGrad_ColocateGradients(self):
self._testWhileGrad_ColocateGradients(colocate=False)
self._testWhileGrad_ColocateGradients(colocate=True)
@@ -1790,6 +1781,7 @@
r = gradients_impl.gradients(r, v)[0]
self.assertAllClose(1024.0, r.eval())
+ @test_util.disable_control_flow_v2("b/116283162 (shape_invariants)")
def testWhileGrad_Shape(self):
with self.cached_session():
x = array_ops.placeholder(dtypes.float32, shape=[None])
@@ -1861,8 +1853,6 @@
self._testWhileGrad_Mul(use_gpu=True, p_iters=10)
def _testNestedWhileCondWhileGrad(self, use_gpu):
- if control_flow_ops.ENABLE_COND_V2:
- return unittest.skip("b/113294377 (unknown shape)")
with self.test_session(use_gpu=use_gpu):
v = constant_op.constant(1.0)
@@ -1885,10 +1875,12 @@
r = gradients_impl.gradients(r, v)[0]
self.assertAllClose(512.0, r.eval())
+ @test_util.disable_control_flow_v2("b/116248044 (nested while)")
def testNestedWhileCondWhileGrad(self):
self._testNestedWhileCondWhileGrad(use_gpu=False)
self._testNestedWhileCondWhileGrad(use_gpu=True)
+ @test_util.disable_control_flow_v2("b/116823782")
def testWhileGrad_Variable(self):
with self.cached_session():
a = variables.Variable(3.0)
@@ -1902,8 +1894,6 @@
self.assertAllClose(216.0, r[0].eval())
def testWhileGradInCond(self):
- if control_flow_ops.ENABLE_COND_V2:
- return unittest.skip("b/110550782 (gradient w.r.t external variable)")
with self.cached_session():
n = ops.convert_to_tensor(1.0, name="n")
@@ -1919,6 +1909,7 @@
r = control_flow_ops.cond(math_ops.less(1, 2), fn1, lambda: x)
self.assertAllClose(9.0, r.eval(feed_dict={x: 1.0}))
+ @test_util.disable_control_flow_v2("b/116340060")
def testGradInWhileWrtInitialLoopVal(self):
with self.cached_session():
x = array_ops.placeholder(dtypes.float32, shape=(), name="x")
@@ -1936,6 +1927,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.disable_control_flow_v2("b/116248044 (nested while)")
def testWhileGradInWhile(self):
with self.cached_session():
n = ops.convert_to_tensor(1.0, name="n")
@@ -1952,9 +1944,8 @@
[tensor_shape.unknown_shape()])
self.assertAllClose(9.0, r.eval(feed_dict={x: 1.0}))
+ @test_util.disable_control_flow_v2("b/116248044 (nested while)")
def testCondGradInNestedWhiles(self):
- if control_flow_ops.ENABLE_COND_V2:
- return unittest.skip("b/113346829 (gpu failure)")
def outer_body(i, x):
_, x = control_flow_ops.while_loop(
@@ -1972,6 +1963,7 @@
self.assertEqual(i_val, 3)
self.assertAllClose(x_val, 1.0)
+ @test_util.disable_control_flow_v2("b/116255781 (flat_args)")
def testWhile_NestedInput(self):
with self.cached_session() as sess:
named = collections.namedtuple("named", ("a", "b"))
@@ -1999,6 +1991,7 @@
self.assertEqual([100.0, 1.0, 102.0, 3.0, 4.0 + 100 * 2.0],
sess.run(r_flattened))
+ @test_util.disable_control_flow_v2("b/116255781(flat_args)")
def testWhile_NestedBadArityFails(self):
with self.cached_session():
named = collections.namedtuple("named", ("a", "b"))
@@ -2057,6 +2050,7 @@
r = gradients_impl.gradients([rx], x)
self.assertAllClose(1024.0, r[0].eval())
+ @test_util.disable_control_flow_v2("b/116355153 (back_prop flag)")
def testWhileGrad_NoGradient(self):
with self.cached_session():
v = constant_op.constant(2.0, name="v")
@@ -2067,6 +2061,7 @@
r = gradients_impl.gradients(r, v)
self.assertAllClose(1.0, r[0].eval())
+ @test_util.disable_control_flow_v2("b/113324949 (RefVariable)")
def testWhileGrad_NoDependency(self):
with self.cached_session() as sess:
variable = variables.Variable(array_ops.ones([2, 3]))
@@ -2180,10 +2175,12 @@
r = gradients_impl.gradients(r, v)[0]
self.assertAllClose(8.0, r.eval())
+ @test_util.disable_control_flow_v2("b/116248044 (nested)")
def testNestedWhileGrad_Simple(self):
self._testNestedWhileGrad_Simple(use_gpu=False)
self._testNestedWhileGrad_Simple(use_gpu=True)
+ @test_util.disable_control_flow_v2("b/116248044 (nested)")
def testNestedWhileGrad_SerialInner(self):
with self.cached_session():
v = constant_op.constant(1.0)
@@ -2207,6 +2204,7 @@
r = gradients_impl.gradients(r, v)[0]
self.assertAllClose(256.0, r.eval())
+ @test_util.disable_control_flow_v2("b/116248044 (nested)")
def testNestedWhileGrad_ParallelInner(self):
with self.cached_session():
v = constant_op.constant(1.0)
@@ -2230,6 +2228,8 @@
r = gradients_impl.gradients(r, v)[0]
self.assertAllClose(512.0, r.eval())
+ @test_util.disable_control_flow_v2(
+ "Nested loops and TensorArrays not supported")
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.
@@ -2268,13 +2268,12 @@
r = gradients_impl.gradients(r, v)[0]
self.assertAllClose(1024.0, r.eval())
+ @test_util.disable_control_flow_v2("b/116272044 (cond_in_while)")
def testWhileCondGrad_Simple(self):
- if control_flow_ops.ENABLE_COND_V2:
- return unittest.skip("b/113294377 (unknown shape)")
-
self._testWhileCondGrad_Simple(use_gpu=False)
self._testWhileCondGrad_Simple(use_gpu=True)
+ @test_util.disable_control_flow_v2("b/116272044 (cond_in_while)")
def testWhileCondGrad_UnknownShape(self):
with self.cached_session() as sess:
v = array_ops.placeholder(dtypes.float32)
@@ -2292,6 +2291,7 @@
r = sess.run(r, feed_dict={v: 2.0})
self.assertAllClose(1024.0, r)
+ @test_util.disable_control_flow_v2("b/116283162 (shape_invariants)")
def testWhileGrad_Concat(self):
with self.cached_session() as sess:
x = variable_scope.get_variable("x", initializer=[[1., 2.]])
@@ -2315,9 +2315,10 @@
sess.run(op)
self.assertAllClose([[0.98000002, 1.98000002]], sess.run(x))
+ @test_util.disable_control_flow_v2("b/113324949 (RefVariable)")
def testWhileWithRefsWithGradients_1(self):
with self.cached_session() as sess:
- x = variables.Variable(0.)._ref() # pylint: disable=protected-access
+ x = variables.VariableV1(0.)._ref() # pylint: disable=protected-access
i = constant_op.constant(0)
c = lambda i, x: math_ops.less(i, 10)
@@ -2329,7 +2330,7 @@
r = control_flow_ops.while_loop(c, body, [i, x], parallel_iterations=5)
- grad_ys = [variables.Variable(73)._ref()] # pylint: disable=protected-access
+ 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()
@@ -2343,6 +2344,7 @@
self.assertEqual(0, value_x)
self.assertEqual(73, value_x_grad)
+ @test_util.disable_control_flow_v2("b/116282023 (IndexedSlices)")
def testWhileGrad_IndexedSlices(self):
with self.cached_session():
values = constant_op.constant([2.0, 4.0], name="values")
@@ -2364,6 +2366,7 @@
r = gradients_impl.gradients(r.values, values)[0]
self.assertAllClose(np.array([1024.0, 1024.0]), r.eval())
+ @test_util.disable_control_flow_v2("b/116328420 (SparseTensor)")
def testWhileGrad_SparseTensor(self):
with self.cached_session():
values = constant_op.constant([2.0, 4.0], name="values")
@@ -2386,6 +2389,7 @@
r = gradients_impl.gradients(r.values, values)[0]
self.assertAllClose(np.array([1024.0, 1024.0]), r.eval())
+ @test_util.disable_control_flow_v2("b/115920078 (gradients)")
def testCallGradInLoop(self):
with self.cached_session() as sess:
i0 = constant_op.constant(0)
@@ -2405,6 +2409,8 @@
c, b, [i0, constant_op.constant(0.0)])
self.assertAllClose(600.0, sess.run(output_grad)[1])
+ @test_util.disable_control_flow_v2(
+ "b/116255781 (flat_args), b/115660901 (TensorArray)")
def testWhileAndTensorArray(self):
with self.cached_session() as sess:
param = constant_op.constant(2.0)
@@ -2509,6 +2515,7 @@
all_ops = x.graph.get_operations()
self.assertFalse(any([name in op.name for op in all_ops]))
+ @test_util.disable_control_flow_v2("b/116255781 (flat args)")
def testWhileGradGradFail(self):
theta = variables.Variable(initial_value=1.)
@@ -2538,6 +2545,7 @@
r = gradients_impl.gradients(r, y)[0]
self.assertEqual(388.0, r.eval())
+ @test_util.disable_control_flow_v2("b/113324949 (RefVariable)")
def testWhileGradientWithNontrainablePath1(self):
q = variables.Variable([7., 8.])
@@ -2555,6 +2563,7 @@
sess.run(q.initializer)
self.assertAllClose([0., 0.], sess.run(dy_dq))
+ @test_util.disable_control_flow_v2("b/113324949 (RefVariable)")
def testWhileGradientWithNontrainablePath2(self):
q = variables.Variable([7., 8.])
@@ -2572,6 +2581,7 @@
sess.run(q.initializer)
self.assertAllClose([1., 1.], sess.run(dy_dq))
+ @test_util.disable_control_flow_v2("b/115920078 (gradients)")
def testIssue16504(self):
c = constant_op.constant(np.arange(100), dtype=dtypes.float32)
w = variables.Variable(
@@ -2595,6 +2605,7 @@
grad, = gradients_impl.gradients(w, c)
self.assertIsNotNone(grad)
+ @test_util.disable_control_flow_v2("b/116270461 (resource)")
def testStopGradMultiFlows(self):
with self.cached_session():
@@ -2653,10 +2664,9 @@
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.disable_control_flow_v2(
+ "b/112477618 (Operation returned from cond)")
def testCase(self):
- if control_flow_ops.ENABLE_COND_V2:
- return unittest.skip("b/112477618 (Operation returned from cond)")
-
with self.cached_session():
x = constant_op.constant(1)
y = constant_op.constant(2)
@@ -2708,10 +2718,9 @@
self.assertAllEqual(r6.eval(), 0)
+ @test_util.disable_control_flow_v2(
+ "b/112477618 (Operation returned from cond)")
def testCaseSideEffects(self):
- if control_flow_ops.ENABLE_COND_V2:
- return unittest.skip("b/112477618 (Operation returned from cond)")
-
with self.cached_session() as sess:
v0 = variables.Variable(-1)
v1 = variables.Variable(-1)
@@ -2746,10 +2755,8 @@
self.assertEqual(0, r0.eval())
self.assertAllEqual(sess.run([v0, v1, v2]), [0, -1, -1])
+ @test_util.disable_control_flow_v2("b/113324949 (ref vars)")
def testOneOpCond(self):
- if control_flow_ops.ENABLE_COND_V2:
- return unittest.skip("b/113324949 (ref vars)")
-
with self.cached_session():
v = variables.Variable(0)
c = ops.convert_to_tensor(0)
@@ -2779,7 +2786,7 @@
def testWithOpsDependencies(self):
with self.cached_session() as sess:
- v = variables.Variable(0.0)
+ v = variables.VariableV1(0.0)
c = constant_op.constant(10)
# Fetching v directly will result in an uninitialized error
@@ -2802,7 +2809,7 @@
def testWithTensorDependencies(self):
with self.cached_session():
- v = variables.Variable(0.0)
+ v = variables.VariableV1(0.0)
c1 = constant_op.constant(10)
c2 = constant_op.constant(20)
@@ -2828,7 +2835,7 @@
def testWithIndexedSlicesDependencies(self):
with self.cached_session():
- v = variables.Variable(
+ v = variables.VariableV1(
np.array([[0.0, 1.0], [10.0, 11.0], [20.0, 21.0]]).astype(np.float32))
v_at_1 = ops.IndexedSlices(v, constant_op.constant([1]))
gather_v_at_1 = array_ops.gather(v_at_1.values, v_at_1.indices)
@@ -2851,18 +2858,18 @@
with ops.Graph().as_default():
# device set on tensor => same device on dep.
with ops.device("/job:ps"):
- vd = variables.Variable([0.0])
+ vd = variables.VariableV1([0.0])
with_vd_dep = control_flow_ops.with_dependencies([vd.initializer], vd)
self.assertTrue("/job:ps" in with_vd_dep.device)
# No device set on tensor => no device on dep.
- vnod = variables.Variable([0.0])
+ vnod = variables.VariableV1([0.0])
with_vnod_dep = control_flow_ops.with_dependencies([vnod.initializer],
vnod)
self.assertDeviceEqual(None, with_vnod_dep.device)
# device set on tensor, default device on graph => default device on dep.
- vdef = variables.Variable([0.0], name="vdef")
+ vdef = variables.VariableV1([0.0], name="vdef")
with ops.device("/job:worker/device:GPU:1"):
with_vdef_dep = control_flow_ops.with_dependencies([vdef.initializer],
vdef)
@@ -2872,8 +2879,8 @@
def testGroup(self):
with self.cached_session() as sess:
- v1 = variables.Variable([0.0])
- v2 = variables.Variable([1.0])
+ v1 = variables.VariableV1([0.0])
+ v2 = variables.VariableV1([1.0])
# Group init1 and init2 and run.
init = control_flow_ops.group(v1.initializer, v2.initializer)
@@ -2955,29 +2962,29 @@
p1 = array_ops.placeholder(dtypes.float32)
p2 = array_ops.placeholder(dtypes.float32)
p3 = array_ops.placeholder(dtypes.float32)
- v1 = variables.Variable(p1, validate_shape=False)
- v2 = variables.Variable(p2, validate_shape=False)
- v3 = variables.Variable(p3, validate_shape=False)
+ v1 = variables.VariableV1(p1, validate_shape=False)
+ v2 = variables.VariableV1(p2, validate_shape=False)
+ v3 = variables.VariableV1(p3, validate_shape=False)
self.assertIs(None, v1.get_shape().ndims)
s = control_flow_ops.ref_select(index, [v1, v2, v3])
self.assertIs(None, s.get_shape().ndims)
# All inputs known but different.
- v1 = variables.Variable([[1, 2]])
- v2 = variables.Variable([[2], [1]])
+ v1 = variables.VariableV1([[1, 2]])
+ v2 = variables.VariableV1([[2], [1]])
s = control_flow_ops.ref_select(index, [v1, v2])
self.assertIs(None, s.get_shape().ndims)
# All inputs known and same.
- v1 = variables.Variable([[1, 2]])
- v2 = variables.Variable([[1, 2]])
+ v1 = variables.VariableV1([[1, 2]])
+ v2 = variables.VariableV1([[1, 2]])
s = control_flow_ops.ref_select(index, [v1, v2])
self.assertEqual([1, 2], s.get_shape())
# Possibly the same but not guaranteed.
- v1 = variables.Variable([[1., 2.]])
+ v1 = variables.VariableV1([[1., 2.]])
p2 = array_ops.placeholder(dtypes.float32, shape=[None, 2])
- v2 = variables.Variable(p2, validate_shape=False)
+ v2 = variables.VariableV1(p2, validate_shape=False)
s = control_flow_ops.ref_select(index, [v1, v2])
self.assertEqual(None, s.get_shape())
@@ -3031,9 +3038,11 @@
r = gradients_impl.gradients(r, x)[0]
self.assertEqual(r.eval(), 524288.0)
- self.assertEqual(
- len([op for op in x.graph.get_operations() if op.type == "StackV2"]),
- 1)
+ # while_v2 does not have stacks.
+ if not control_flow_ops.ENABLE_WHILE_V2:
+ self.assertEqual(
+ len([op for op in x.graph.get_operations() if op.type == "StackV2"
+ ]), 1)
class ControlFlowContextCheckTest(test.TestCase):
@@ -3160,11 +3169,11 @@
def testTensors(self):
for v1_first in [True, False]:
with self.cached_session():
- v1 = variables.Variable([1.0])
+ v1 = variables.VariableV1([1.0])
add1 = math_ops.add(
control_flow_ops.with_dependencies([v1.initializer], v1._ref()), # pylint: disable=protected-access
2.0)
- v2 = variables.Variable([10.0])
+ v2 = variables.VariableV1([10.0])
add2 = math_ops.add(
control_flow_ops.with_dependencies([v2.initializer], v2._ref()), # pylint: disable=protected-access
20.0)
@@ -3190,14 +3199,14 @@
def testIndexedSlices(self):
for v1_first in [True, False]:
with self.cached_session():
- v1 = variables.Variable(
+ v1 = variables.VariableV1(
np.array([[0.0, 1.0], [10.0, 11.0], [20.0, 21.0]]).astype(
np.float32))
v1_at_1 = ops.IndexedSlices(
control_flow_ops.with_dependencies([v1.initializer], v1._ref()), # pylint: disable=protected-access
constant_op.constant([1]))
- v2 = variables.Variable(
+ v2 = variables.VariableV1(
np.array([[0.1, 1.1], [10.1, 11.1], [20.1, 21.1]]).astype(
np.float32))
v2_at_1 = ops.IndexedSlices(
@@ -3229,7 +3238,7 @@
def testAcceptTensorsAsControlInputs(self):
with self.cached_session():
- var = variables.Variable(0)
+ var = variables.VariableV1(0)
assign = state_ops.assign(var, 1)
t, = control_flow_ops.tuple(
[constant_op.constant(0)], control_inputs=[assign])
@@ -3393,7 +3402,7 @@
name="unroll_same_device", iters=iters, wall_time=duration)
-@test_util.with_cond_v2
+@test_util.with_control_flow_v2
class EagerTest(test.TestCase):
def testCond(self):
@@ -3406,6 +3415,27 @@
self.assertAllEqual(r.numpy(), 10)
self.assertFalse(isinstance(r, list))
+ def testCondInDefun(self):
+ if "GPU" in [d.device_type for d in device_lib.list_local_devices()]:
+ return unittest.skip("b/113346829 (gpu failure)")
+
+ with context.eager_mode():
+
+ @eager_function.defun
+ def foo(pred):
+ # TODO(b/111124878): this only needs to output one element.
+ fn1 = lambda: (constant_op.constant(10), constant_op.constant(100))
+ fn2 = lambda: (constant_op.constant(20), constant_op.constant(200))
+ return control_flow_ops.cond(constant_op.constant(pred), fn1, fn2)
+
+ r = foo(True)
+ self.assertAllEqual(r[0].numpy(), 10)
+ self.assertNotIsInstance(r, list)
+
+ r = foo(False)
+ self.assertAllEqual(r[0].numpy(), 20)
+ self.assertFalse(isinstance(r, list))
+
def testWhileLoop(self):
with context.eager_mode():
tensor = constant_op.constant([1, 2, 3, 4, 5])
diff --git a/tensorflow/python/kernel_tests/dense_update_ops_test.py b/tensorflow/python/kernel_tests/dense_update_ops_test.py
index 06c3271..120e103 100644
--- a/tensorflow/python/kernel_tests/dense_update_ops_test.py
+++ b/tensorflow/python/kernel_tests/dense_update_ops_test.py
@@ -87,7 +87,7 @@
def testAssignNonStrictShapeChecking(self):
with self.cached_session():
data = array_ops.fill([1024, 1024], 0)
- p = variables.Variable([1])
+ p = variables.VariableV1([1])
a = state_ops.assign(p, data, validate_shape=False)
a.op.run()
self.assertAllEqual(p.eval(), data.eval())
@@ -100,14 +100,14 @@
def testInitRequiredAssignAdd(self):
with self.cached_session():
- p = variables.Variable(array_ops.fill([1024, 1024], 1), dtypes.int32)
+ p = variables.VariableV1(array_ops.fill([1024, 1024], 1), dtypes.int32)
a = state_ops.assign_add(p, array_ops.fill([1024, 1024], 0))
with self.assertRaisesOpError("use uninitialized"):
a.op.run()
def testInitRequiredAssignSub(self):
with self.cached_session():
- p = variables.Variable(array_ops.fill([1024, 1024], 1), dtypes.int32)
+ p = variables.VariableV1(array_ops.fill([1024, 1024], 1), dtypes.int32)
a = state_ops.assign_sub(p, array_ops.fill([1024, 1024], 0))
with self.assertRaisesOpError("use uninitialized"):
a.op.run()
diff --git a/tensorflow/python/kernel_tests/distributions/bernoulli_test.py b/tensorflow/python/kernel_tests/distributions/bernoulli_test.py
index 26d013b..37b35ba 100644
--- a/tensorflow/python/kernel_tests/distributions/bernoulli_test.py
+++ b/tensorflow/python/kernel_tests/distributions/bernoulli_test.py
@@ -118,7 +118,9 @@
self.assertEqual(dist.probs.dtype, dist.stddev().dtype)
self.assertEqual(dist.probs.dtype, dist.entropy().dtype)
self.assertEqual(dist.probs.dtype, dist.prob(0).dtype)
+ self.assertEqual(dist.probs.dtype, dist.prob(0.5).dtype)
self.assertEqual(dist.probs.dtype, dist.log_prob(0).dtype)
+ self.assertEqual(dist.probs.dtype, dist.log_prob(0.5).dtype)
dist64 = make_bernoulli([], dtypes.int64)
self.assertEqual(dist64.dtype, dtypes.int64)
@@ -181,6 +183,16 @@
return
self._testPmf(logits=special.logit(p))
+ @test_util.run_in_graph_and_eager_modes
+ def testPmfWithFloatArgReturnsXEntropy(self):
+ p = [[0.2], [0.4], [0.3], [0.6]]
+ samps = [0, 0.1, 0.8]
+ self.assertAllClose(
+ np.float32(samps) * np.log(np.float32(p)) +
+ (1 - np.float32(samps)) * np.log(1 - np.float32(p)),
+ self.evaluate(
+ bernoulli.Bernoulli(probs=p, validate_args=False).log_prob(samps)))
+
def testBroadcasting(self):
with self.cached_session():
p = array_ops.placeholder(dtypes.float32)
diff --git a/tensorflow/python/kernel_tests/distributions/exponential_test.py b/tensorflow/python/kernel_tests/distributions/exponential_test.py
index 27d1291..367f8bb 100644
--- a/tensorflow/python/kernel_tests/distributions/exponential_test.py
+++ b/tensorflow/python/kernel_tests/distributions/exponential_test.py
@@ -81,6 +81,22 @@
expected_cdf = stats.expon.cdf(x, scale=1 / lam_v)
self.assertAllClose(self.evaluate(cdf), expected_cdf)
+ def testExponentialLogSurvival(self):
+ batch_size = 7
+ lam = constant_op.constant([2.0] * batch_size)
+ lam_v = 2.0
+ x = np.array([2.5, 2.5, 4.0, 0.1, 1.0, 2.0, 10.0], dtype=np.float32)
+
+ exponential = exponential_lib.Exponential(rate=lam)
+
+ log_survival = exponential.log_survival_function(x)
+ self.assertEqual(log_survival.get_shape(), (7,))
+
+ if not stats:
+ return
+ expected_log_survival = stats.expon.logsf(x, scale=1 / lam_v)
+ self.assertAllClose(self.evaluate(log_survival), expected_log_survival)
+
def testExponentialMean(self):
lam_v = np.array([1.0, 4.0, 2.5])
exponential = exponential_lib.Exponential(rate=lam_v)
diff --git a/tensorflow/python/kernel_tests/distributions/normal_test.py b/tensorflow/python/kernel_tests/distributions/normal_test.py
index de73a40..6625a88 100644
--- a/tensorflow/python/kernel_tests/distributions/normal_test.py
+++ b/tensorflow/python/kernel_tests/distributions/normal_test.py
@@ -78,6 +78,14 @@
self.assertEqual(expected, sigma_shape)
@test_util.run_in_graph_and_eager_modes
+ def testSampleLikeArgsGetDistDType(self):
+ dist = normal_lib.Normal(0., 1.)
+ self.assertEqual(dtypes.float32, dist.dtype)
+ for method in ("log_prob", "prob", "log_cdf", "cdf",
+ "log_survival_function", "survival_function", "quantile"):
+ self.assertEqual(dtypes.float32, getattr(dist, method)(1).dtype)
+
+ @test_util.run_in_graph_and_eager_modes
def testParamShapes(self):
sample_shape = [10, 3, 4]
self._testParamShapes(sample_shape, sample_shape)
diff --git a/tensorflow/python/kernel_tests/identity_op_py_test.py b/tensorflow/python/kernel_tests/identity_op_py_test.py
index 37f9f71..88ea10c 100644
--- a/tensorflow/python/kernel_tests/identity_op_py_test.py
+++ b/tensorflow/python/kernel_tests/identity_op_py_test.py
@@ -61,7 +61,7 @@
def testRefIdentityShape(self):
with self.cached_session():
shape = [2, 3]
- tensor = variables.Variable(
+ tensor = variables.VariableV1(
constant_op.constant(
[[1, 2, 3], [6, 5, 4]], dtype=dtypes.int32))
self.assertEquals(shape, tensor.get_shape())
diff --git a/tensorflow/python/kernel_tests/list_ops_test.py b/tensorflow/python/kernel_tests/list_ops_test.py
index 0f56077..ae413ed 100644
--- a/tensorflow/python/kernel_tests/list_ops_test.py
+++ b/tensorflow/python/kernel_tests/list_ops_test.py
@@ -170,6 +170,32 @@
list_ops.tensor_list_pop_back(
l_cpu, element_dtype=dtypes.float32)[1]), 2.0)
+ @test_util.run_in_graph_and_eager_modes
+ def testCPUGPUCopyNested(self):
+ if not context.num_gpus():
+ return
+ t = constant_op.constant([1.0, 2.0])
+ child_l = list_ops.tensor_list_from_tensor(t, element_shape=scalar_shape())
+ l = list_ops.empty_tensor_list(
+ element_shape=constant_op.constant([], dtype=dtypes.int32),
+ element_dtype=dtypes.variant)
+ l = list_ops.tensor_list_push_back(l, child_l)
+ with context.device("gpu:0"):
+ l_gpu = array_ops.identity(l)
+ _, child_l_gpu = list_ops.tensor_list_pop_back(
+ l_gpu, element_dtype=dtypes.variant)
+ self.assertAllEqual(
+ self.evaluate(
+ list_ops.tensor_list_pop_back(
+ child_l_gpu, element_dtype=dtypes.float32)[1]), 2.0)
+ l_cpu = array_ops.identity(l_gpu)
+ _, child_l_cpu = list_ops.tensor_list_pop_back(
+ l_cpu, element_dtype=dtypes.variant)
+ self.assertAllEqual(
+ self.evaluate(
+ list_ops.tensor_list_pop_back(
+ child_l_cpu, element_dtype=dtypes.float32)[1]), 2.0)
+
def testGraphStack(self):
with self.cached_session():
tl = list_ops.empty_tensor_list(
diff --git a/tensorflow/python/kernel_tests/resource_variable_ops_test.py b/tensorflow/python/kernel_tests/resource_variable_ops_test.py
index f90545f..1365d4b 100644
--- a/tensorflow/python/kernel_tests/resource_variable_ops_test.py
+++ b/tensorflow/python/kernel_tests/resource_variable_ops_test.py
@@ -290,7 +290,7 @@
self.assertEqual(self.evaluate(read), [[2]])
def testUseResource(self):
- v = variables.Variable(1.0, use_resource=True)
+ v = variables.VariableV1(1.0, use_resource=True)
self.assertTrue(isinstance(v, resource_variable_ops.ResourceVariable))
def testEagerNoUseResource(self):
diff --git a/tensorflow/python/kernel_tests/rnn_test.py b/tensorflow/python/kernel_tests/rnn_test.py
index 05ad9f6..2f6963f 100644
--- a/tensorflow/python/kernel_tests/rnn_test.py
+++ b/tensorflow/python/kernel_tests/rnn_test.py
@@ -535,6 +535,45 @@
self.assertAllClose(tf_out, k_out)
self.assertAllClose(tf_state, k_state)
+ def testSimpleRNNCellAndBasicRNNCellComparison(self):
+ input_shape = 10
+ output_shape = 5
+ timestep = 4
+ batch = 20
+ (x_train, _), _ = testing_utils.get_test_data(
+ train_samples=batch,
+ test_samples=0,
+ input_shape=(timestep, input_shape),
+ num_classes=output_shape)
+ fix_weights_generator = keras.layers.SimpleRNNCell(output_shape)
+ fix_weights_generator.build((None, input_shape))
+ # The SimpleRNNCell contains 3 weights: kernel, recurrent_kernel, and bias
+ # The BasicRNNCell contains 2 weight: kernel and bias, where kernel is
+ # zipped [kernel, recurrent_kernel] in SimpleRNNCell.
+ keras_weights = fix_weights_generator.get_weights()
+ kernel, recurrent_kernel, bias = keras_weights
+ tf_weights = [np.concatenate((kernel, recurrent_kernel)), bias]
+
+ with self.test_session(graph=ops_lib.Graph()) as sess:
+ inputs = array_ops.placeholder(
+ dtypes.float32, shape=(None, timestep, input_shape))
+ cell = keras.layers.SimpleRNNCell(output_shape)
+ k_out, k_state = rnn.dynamic_rnn(
+ cell, inputs, dtype=dtypes.float32)
+ cell.set_weights(keras_weights)
+ [k_out, k_state] = sess.run([k_out, k_state], {inputs: x_train})
+ with self.test_session(graph=ops_lib.Graph()) as sess:
+ inputs = array_ops.placeholder(
+ dtypes.float32, shape=(None, timestep, input_shape))
+ cell = rnn_cell_impl.BasicRNNCell(output_shape)
+ tf_out, tf_state = rnn.dynamic_rnn(
+ cell, inputs, dtype=dtypes.float32)
+ cell.set_weights(tf_weights)
+ [tf_out, tf_state] = sess.run([tf_out, tf_state], {inputs: x_train})
+
+ self.assertAllClose(tf_out, k_out)
+ self.assertAllClose(tf_state, k_state)
+
def testBasicLSTMCellInterchangeWithLSTMCell(self):
with self.session(graph=ops_lib.Graph()) as sess:
basic_cell = rnn_cell_impl.BasicLSTMCell(1)
diff --git a/tensorflow/python/kernel_tests/scatter_nd_ops_test.py b/tensorflow/python/kernel_tests/scatter_nd_ops_test.py
index 86e063c..4b92309 100644
--- a/tensorflow/python/kernel_tests/scatter_nd_ops_test.py
+++ b/tensorflow/python/kernel_tests/scatter_nd_ops_test.py
@@ -136,7 +136,7 @@
new = ref.copy()
np_scatter(new, indices, updates)
# Scatter via tensorflow
- ref_var = variables.Variable(ref)
+ ref_var = variables.VariableV1(ref)
ref_var.initializer.run()
tf_scatter(ref_var, indices, updates).eval()
@@ -258,7 +258,7 @@
params = np.array([1, 2, 3, 4, 5, 6]).astype(np.float32)
updates = np.array([-3, -4, -5]).astype(np.float32)
with self.test_session(use_gpu=False):
- ref = variables.Variable(params)
+ ref = variables.VariableV1(params)
ref.initializer.run()
# Indices all in range, no problem.
diff --git a/tensorflow/python/kernel_tests/scatter_ops_test.py b/tensorflow/python/kernel_tests/scatter_ops_test.py
index 1a0fa74..527b7da 100644
--- a/tensorflow/python/kernel_tests/scatter_ops_test.py
+++ b/tensorflow/python/kernel_tests/scatter_ops_test.py
@@ -178,7 +178,7 @@
np_scatter = _TF_OPS_TO_NUMPY[tf_scatter]
np_scatter(new, indices, updates)
# Scatter via tensorflow
- ref = variables.Variable(old)
+ ref = variables.VariableV1(old)
ref.initializer.run()
tf_scatter(ref, indices, updates).eval()
self.assertAllClose(ref.eval(), new)
@@ -294,7 +294,7 @@
updates = np.array([-3, -4, -5]).astype(np.float32)
if not test.is_gpu_available():
with self.test_session(use_gpu=False):
- ref = variables.Variable(params)
+ ref = variables.VariableV1(params)
ref.initializer.run()
# Indices all in range, no problem.
diff --git a/tensorflow/python/kernel_tests/softplus_op_test.py b/tensorflow/python/kernel_tests/softplus_op_test.py
index e8dc272..636ed47 100644
--- a/tensorflow/python/kernel_tests/softplus_op_test.py
+++ b/tensorflow/python/kernel_tests/softplus_op_test.py
@@ -126,7 +126,7 @@
with self.assertRaisesRegexp(
TypeError,
"'features' has DataType int32 not in list of allowed values"):
- nn_ops.softplus(constant_op.constant(7)).eval()
+ nn_ops.softplus(constant_op.constant(42)).eval()
if __name__ == "__main__":
diff --git a/tensorflow/python/kernel_tests/unicode_script_op_test.py b/tensorflow/python/kernel_tests/unicode_script_op_test.py
new file mode 100644
index 0000000..927e545
--- /dev/null
+++ b/tensorflow/python/kernel_tests/unicode_script_op_test.py
@@ -0,0 +1,57 @@
+# 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.
+#===============================================================================
+"""Functional tests for UnicodeScript op."""
+
+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.ops import string_ops
+from tensorflow.python.platform import test
+
+
+class UnicodeScriptOpTest(test.TestCase):
+
+ def testValidScripts(self):
+ inputs = [
+ ord("a"),
+ 0x0411, # CYRILLIC CAPITAL LETTER BE
+ 0x82b8, # CJK UNIFIED IDEOGRAPH-82B8
+ ord(",")
+ ]
+ with self.cached_session():
+ input_vector = constant_op.constant(inputs, dtypes.int32)
+ outputs = string_ops.unicode_script(input_vector).eval()
+ self.assertAllEqual(
+ outputs,
+ [
+ 25, # USCRIPT_LATIN (LATN)
+ 8, # USCRIPT_CYRILLIC (CYRL)
+ 17, # USCRIPT_HAN (HANI)
+ 0 # USCRIPT_COMMON (ZYYY)
+ ])
+
+ def testInvalidScript(self):
+ inputs = [-100, 0xffffff]
+ with self.cached_session():
+ input_vector = constant_op.constant(inputs, dtypes.int32)
+ outputs = string_ops.unicode_script(input_vector).eval()
+ self.assertAllEqual(outputs, [-1, -1])
+
+
+if __name__ == "__main__":
+ test.main()
diff --git a/tensorflow/python/kernel_tests/variable_scope_test.py b/tensorflow/python/kernel_tests/variable_scope_test.py
index 401e1ae..33f464f 100644
--- a/tensorflow/python/kernel_tests/variable_scope_test.py
+++ b/tensorflow/python/kernel_tests/variable_scope_test.py
@@ -394,10 +394,10 @@
old = variable_scope._DEFAULT_USE_RESOURCE
try:
variable_scope.enable_resource_variables()
- self.assertTrue(isinstance(variables_lib.Variable(1.0),
+ self.assertTrue(isinstance(variables_lib.VariableV1(1.0),
resource_variable_ops.ResourceVariable))
variable_scope.disable_resource_variables()
- self.assertFalse(isinstance(variables_lib.Variable(1.0),
+ self.assertFalse(isinstance(variables_lib.VariableV1(1.0),
resource_variable_ops.ResourceVariable))
finally:
variable_scope._DEFAULT_USE_RESOURCE = old
diff --git a/tensorflow/python/kernel_tests/variables_test.py b/tensorflow/python/kernel_tests/variables_test.py
index 2e79756..c2b8608 100644
--- a/tensorflow/python/kernel_tests/variables_test.py
+++ b/tensorflow/python/kernel_tests/variables_test.py
@@ -43,14 +43,14 @@
def testInitialization(self):
with self.cached_session():
- var0 = variables.Variable(0.0)
+ var0 = variables.VariableV1(0.0)
self.assertEqual("Variable:0", var0.name)
self.assertEqual("Variable", var0._shared_name)
self.assertEqual([], var0.get_shape())
self.assertEqual([], var0.get_shape())
self.assertEqual([], var0.shape)
- var1 = variables.Variable(1.1)
+ var1 = variables.VariableV1(1.1)
self.assertEqual("Variable_1:0", var1.name)
self.assertEqual("Variable_1", var1._shared_name)
self.assertEqual([], var1.get_shape())
@@ -143,7 +143,7 @@
def testZeroSizeStringAssign(self):
with self.cached_session() as sess:
- array = variables.Variable(
+ array = variables.VariableV1(
initial_value=array_ops.zeros((0,), dtype=dtypes.string),
name="foo",
trainable=False,
@@ -192,7 +192,7 @@
# d get the control dep.
d = constant_op.constant(2.0)
# variables do not.
- var_x = variables.Variable(2.0)
+ var_x = variables.VariableV1(2.0)
self.assertEqual([c.op], d.op.control_inputs)
self.assertEqual([], var_x.initializer.control_inputs)
self.assertEqual([], var_x.value().op.control_inputs)
@@ -280,10 +280,10 @@
def testCollections(self):
with self.cached_session():
- var_x = variables.Variable(2.0)
- var_y = variables.Variable(2.0, trainable=False)
- var_z = variables.Variable(2.0, trainable=True)
- var_t = variables.Variable(
+ var_x = variables.VariableV1(2.0)
+ var_y = variables.VariableV1(2.0, trainable=False)
+ var_z = variables.VariableV1(2.0, trainable=True)
+ var_t = variables.VariableV1(
2.0,
trainable=True,
collections=[
@@ -296,9 +296,9 @@
def testCollectionsWithScope(self):
with self.cached_session():
with ops.name_scope("scope_1"):
- var_x = variables.Variable(2.0)
+ var_x = variables.VariableV1(2.0)
with ops.name_scope("scope_2"):
- var_y = variables.Variable(2.0)
+ var_y = variables.VariableV1(2.0)
self.assertEqual([var_x, var_y], variables.global_variables())
self.assertEqual([var_x], variables.global_variables("scope_1"))
@@ -399,7 +399,7 @@
def testColocation(self):
with ops.device("/job:ps"):
- var = variables.Variable(0, name="v")
+ var = variables.VariableV1(0, name="v")
with ops.device("/job:worker/task:7"):
assign_op = var.assign(1)
self.assertDeviceEqual("/job:ps", assign_op.device)
@@ -522,7 +522,7 @@
self.assertAllClose(np.ones((5, 5), np.float32), var.eval())
def testRepr(self):
- var = variables.Variable(np.zeros((5, 5), np.float32), name="noop")
+ var = variables.VariableV1(np.zeros((5, 5), np.float32), name="noop")
self.assertEqual(
"<tf.Variable 'noop:0' shape=(5, 5) dtype=float32_ref>",
repr(var))
@@ -556,8 +556,8 @@
def testVariableList(self):
with ops.Graph().as_default(), self.cached_session() as sess:
- v = variables.Variable([1, 2], name="v")
- w = variables.Variable([3, 4], name="w")
+ v = variables.VariableV1([1, 2], name="v")
+ w = variables.VariableV1([3, 4], name="w")
uninited = variables.report_uninitialized_variables()
self.assertAllEqual(np.array([b"v", b"w"]), sess.run(uninited))
sess.run(w.initializer)
@@ -593,8 +593,8 @@
def testVariables(self):
with ops.Graph().as_default(), self.cached_session() as sess:
- v = variables.Variable([1, 2])
- w = variables.Variable([3, 4])
+ v = variables.VariableV1([1, 2])
+ w = variables.VariableV1([3, 4])
_ = v, w
inited = variables.assert_variables_initialized()
with self.assertRaisesOpError("Attempting to use uninitialized value"):
@@ -604,8 +604,8 @@
def testVariableList(self):
with ops.Graph().as_default(), self.cached_session() as sess:
- v = variables.Variable([1, 2])
- w = variables.Variable([3, 4])
+ v = variables.VariableV1([1, 2])
+ w = variables.VariableV1([3, 4])
inited = variables.assert_variables_initialized([v])
with self.assertRaisesOpError("Attempting to use uninitialized value"):
inited.op.run()
@@ -696,6 +696,48 @@
variable_list=[v0],
partitions=partitions)
+ def testPartitionedVariableAssignments(self):
+ with ops.Graph().as_default(), self.cached_session() as sess:
+ v0 = variables.Variable(initial_value=[0.0])
+ v1 = variables.Variable(initial_value=[1.0])
+ v0._set_save_slice_info(
+ variables.Variable.SaveSliceInfo(v0.name, [2], [0], [1]))
+ v1._set_save_slice_info(
+ variables.Variable.SaveSliceInfo(v0.name, [2], [1], [1]))
+ partitions = [2]
+
+ # Pass variable_list as [v1, v0] to ensure they are properly
+ # re-sorted to [v0, v1] based on their slice info offsets.
+ partitioned_variable = variables.PartitionedVariable(
+ name="two_vars",
+ shape=[2],
+ dtype=v0.dtype,
+ variable_list=[v0, v1],
+ partitions=partitions)
+
+ deltas_a = constant_op.constant([1.0, 2.0])
+ deltas_b = constant_op.constant([3.0, 4.0])
+ ones = array_ops.ones([2])
+ plus_delta = partitioned_variable.assign_add(deltas_a)
+ minus_delta = partitioned_variable.assign_sub(deltas_b)
+ assign_ones = partitioned_variable.assign(ones)
+ variables.global_variables_initializer().run()
+
+ self.assertEqual([1.0], plus_delta[0].eval())
+ self.assertEqual([1.0], v0.eval())
+ self.assertEqual([3.0], plus_delta[1].eval())
+ self.assertEqual([3.0], v1.eval())
+
+ self.assertEqual([-2.0], minus_delta[0].eval())
+ self.assertEqual([-2.0], v0.eval())
+ self.assertEqual([-1.0], minus_delta[1].eval())
+ self.assertEqual([-1.0], v1.eval())
+
+ self.assertEqual([1.0], assign_ones[0].eval())
+ self.assertEqual([1.0], v0.eval())
+ self.assertEqual([1.0], assign_ones[1].eval())
+ self.assertEqual([1.0], v1.eval())
+
class VariableContainerTest(test.TestCase):
diff --git a/tensorflow/python/layers/base.py b/tensorflow/python/layers/base.py
index 3ba880d..e399ece 100644
--- a/tensorflow/python/layers/base.py
+++ b/tensorflow/python/layers/base.py
@@ -131,10 +131,20 @@
def add_loss(self, losses, inputs=None):
previous_losses_length = len(self._losses)
+ previous_callable_losses_length = len(self._callable_losses)
super(Layer, self).add_loss(losses, inputs=inputs)
- # TODO(fchollet): deprecate collection below.
- new_losses = self._losses[previous_losses_length:]
- _add_elements_to_collection(new_losses, ops.GraphKeys.REGULARIZATION_LOSSES)
+ if not context.executing_eagerly():
+ # TODO(fchollet): deprecate collection below.
+ new_losses = self._losses[previous_losses_length:]
+ new_callable_losses = self._callable_losses[
+ previous_callable_losses_length:]
+ for regularizer in new_callable_losses:
+ loss_tensor = regularizer()
+ if loss_tensor is not None:
+ new_losses.append(loss_tensor)
+ _add_elements_to_collection(
+ new_losses,
+ ops.GraphKeys.REGULARIZATION_LOSSES)
def _name_scope(self):
"""Determines op naming for the Layer."""
diff --git a/tensorflow/python/layers/convolutional_test.py b/tensorflow/python/layers/convolutional_test.py
index d61d3b6..257fa27 100644
--- a/tensorflow/python/layers/convolutional_test.py
+++ b/tensorflow/python/layers/convolutional_test.py
@@ -207,7 +207,8 @@
layer.apply(images)
loss_keys = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)
self.assertEqual(len(loss_keys), 1)
- self.assertListEqual(layer.losses, loss_keys)
+ self.evaluate([v.initializer for v in layer.variables])
+ self.assertListEqual(self.evaluate(layer.losses), self.evaluate(loss_keys))
def testConv2DBiasRegularizer(self):
height, width = 7, 9
@@ -217,7 +218,8 @@
layer.apply(images)
loss_keys = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)
self.assertEqual(len(loss_keys), 1)
- self.assertListEqual(layer.losses, loss_keys)
+ self.evaluate([v.initializer for v in layer.variables])
+ self.assertListEqual(self.evaluate(layer.losses), self.evaluate(loss_keys))
def testConv2DNoBias(self):
height, width = 7, 9
@@ -445,7 +447,8 @@
layer.apply(data)
loss_keys = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)
self.assertEqual(len(loss_keys), 1)
- self.assertEqual(layer.losses, loss_keys)
+ self.evaluate([v.initializer for v in layer.variables])
+ self.assertListEqual(self.evaluate(layer.losses), self.evaluate(loss_keys))
def testSeparableConv1DPointwiseRegularizer(self):
length = 9
@@ -455,7 +458,8 @@
layer.apply(data)
loss_keys = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)
self.assertEqual(len(loss_keys), 1)
- self.assertEqual(layer.losses, loss_keys)
+ self.evaluate([v.initializer for v in layer.variables])
+ self.assertListEqual(self.evaluate(layer.losses), self.evaluate(loss_keys))
def testSeparableConv1DBiasRegularizer(self):
length = 9
@@ -465,7 +469,8 @@
layer.apply(data)
loss_keys = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)
self.assertEqual(len(loss_keys), 1)
- self.assertEqual(layer.losses, loss_keys)
+ self.evaluate([v.initializer for v in layer.variables])
+ self.assertListEqual(self.evaluate(layer.losses), self.evaluate(loss_keys))
def testSeparableConv1DNoBias(self):
length = 9
@@ -682,7 +687,8 @@
layer.apply(images)
loss_keys = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)
self.assertEqual(len(loss_keys), 1)
- self.assertListEqual(layer.losses, loss_keys)
+ self.evaluate([v.initializer for v in layer.variables])
+ self.assertListEqual(self.evaluate(layer.losses), self.evaluate(loss_keys))
def testSeparableConv2DPointwiseRegularizer(self):
height, width = 7, 9
@@ -692,7 +698,8 @@
layer.apply(images)
loss_keys = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)
self.assertEqual(len(loss_keys), 1)
- self.assertListEqual(layer.losses, loss_keys)
+ self.evaluate([v.initializer for v in layer.variables])
+ self.assertListEqual(self.evaluate(layer.losses), self.evaluate(loss_keys))
def testSeparableConv2DBiasRegularizer(self):
height, width = 7, 9
@@ -702,7 +709,8 @@
layer.apply(images)
loss_keys = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)
self.assertEqual(len(loss_keys), 1)
- self.assertListEqual(layer.losses, loss_keys)
+ self.evaluate([v.initializer for v in layer.variables])
+ self.assertListEqual(self.evaluate(layer.losses), self.evaluate(loss_keys))
def testSeparableConv2DNoBias(self):
height, width = 7, 9
@@ -839,7 +847,8 @@
layer.apply(images)
loss_keys = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)
self.assertEqual(len(loss_keys), 1)
- self.assertListEqual(layer.losses, loss_keys)
+ self.evaluate([v.initializer for v in layer.variables])
+ self.assertListEqual(self.evaluate(layer.losses), self.evaluate(loss_keys))
def testConv2DTransposeBiasRegularizer(self):
height, width = 7, 9
@@ -849,7 +858,8 @@
layer.apply(images)
loss_keys = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)
self.assertEqual(len(loss_keys), 1)
- self.assertListEqual(layer.losses, loss_keys)
+ self.evaluate([v.initializer for v in layer.variables])
+ self.assertListEqual(self.evaluate(layer.losses), self.evaluate(loss_keys))
def testConv2DTransposeNoBias(self):
height, width = 7, 9
@@ -1017,7 +1027,8 @@
layer.apply(volumes)
loss_keys = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)
self.assertEqual(len(loss_keys), 1)
- self.assertListEqual(layer.losses, loss_keys)
+ self.evaluate([v.initializer for v in layer.variables])
+ self.assertListEqual(self.evaluate(layer.losses), self.evaluate(loss_keys))
def testConv3DTransposeBiasRegularizer(self):
depth, height, width = 5, 7, 9
@@ -1027,7 +1038,8 @@
layer.apply(volumes)
loss_keys = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)
self.assertEqual(len(loss_keys), 1)
- self.assertListEqual(layer.losses, loss_keys)
+ self.evaluate([v.initializer for v in layer.variables])
+ self.assertListEqual(self.evaluate(layer.losses), self.evaluate(loss_keys))
def testConv3DTransposeNoBias(self):
depth, height, width = 5, 7, 9
diff --git a/tensorflow/python/layers/core_test.py b/tensorflow/python/layers/core_test.py
index 46009a3..d26f3f4 100644
--- a/tensorflow/python/layers/core_test.py
+++ b/tensorflow/python/layers/core_test.py
@@ -197,7 +197,8 @@
_ = dense(inputs)
loss_keys = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)
self.assertEqual(len(loss_keys), 1)
- self.assertListEqual(dense.losses, loss_keys)
+ self.evaluate([v.initializer for v in dense.variables])
+ self.assertAllEqual(self.evaluate(dense.losses), self.evaluate(loss_keys))
def testKernelRegularizerWithReuse(self):
regularizer = lambda x: math_ops.reduce_sum(x) * 1e-3
@@ -218,7 +219,8 @@
_ = dense(inputs)
loss_keys = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)
self.assertEqual(len(loss_keys), 1)
- self.assertListEqual(dense.losses, loss_keys)
+ self.evaluate([v.initializer for v in dense.variables])
+ self.assertAllEqual(self.evaluate(dense.losses), self.evaluate(loss_keys))
def testFunctionalDense(self):
with self.cached_session():
diff --git a/tensorflow/python/lib/io/tf_record.py b/tensorflow/python/lib/io/tf_record.py
index cce71a2..9ab683d 100644
--- a/tensorflow/python/lib/io/tf_record.py
+++ b/tensorflow/python/lib/io/tf_record.py
@@ -22,10 +22,12 @@
from tensorflow.python import pywrap_tensorflow
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
-@tf_export("python_io.TFRecordCompressionType")
+@tf_export("io.TFRecordCompressionType", "python_io.TFRecordCompressionType")
+@deprecation.deprecated_endpoints("python_io.TFRecordCompressionType")
class TFRecordCompressionType(object):
"""The type of compression for the record."""
NONE = 0
@@ -33,7 +35,8 @@
GZIP = 2
-@tf_export("python_io.TFRecordOptions")
+@tf_export("io.TFRecordOptions", "python_io.TFRecordOptions")
+@deprecation.deprecated_endpoints("python_io.TFRecordOptions")
class TFRecordOptions(object):
"""Options used for manipulating TFRecord files."""
compression_type_map = {
@@ -143,7 +146,8 @@
return options
-@tf_export("python_io.tf_record_iterator")
+@tf_export("io.tf_record_iterator", "python_io.tf_record_iterator")
+@deprecation.deprecated_endpoints("python_io.tf_record_iterator")
def tf_record_iterator(path, options=None):
"""An iterator that read the records from a TFRecords file.
@@ -175,7 +179,8 @@
reader.Close()
-@tf_export("python_io.TFRecordWriter")
+@tf_export("io.TFRecordWriter", "python_io.TFRecordWriter")
+@deprecation.deprecated_endpoints("python_io.TFRecordWriter")
class TFRecordWriter(object):
"""A class to write records to a TFRecords file.
diff --git a/tensorflow/python/ops/array_ops.py b/tensorflow/python/ops/array_ops.py
index a7f57e9..9f5149d 100644
--- a/tensorflow/python/ops/array_ops.py
+++ b/tensorflow/python/ops/array_ops.py
@@ -1204,7 +1204,8 @@
return _apply_mask_1d(tensor, mask, axis)
-@tf_export("sparse_mask")
+@tf_export("sparse.mask", "sparse_mask")
+@deprecation.deprecated_endpoints("sparse_mask")
def sparse_mask(a, mask_indices, name=None):
"""Masks elements of `IndexedSlices`.
@@ -1226,7 +1227,7 @@
# `b` will be the subset of `a` slices at its second and third indices, so
# we want to mask its first and last indices (which are at absolute
# indices 12, 45)
- b = tf.sparse_mask(a, [12, 45])
+ b = tf.sparse.mask(a, [12, 45])
b.indices # [26, 37]
tf.shape(b.values) # [2, 10]
@@ -1382,7 +1383,7 @@
[10, 11, 12]]])
# Take the transpose of the matrices in dimension-0
- # (this common operation has a shorthand `matrix_transpose`)
+ # (this common operation has a shorthand `linalg.transpose`)
tf.transpose(x, perm=[0, 2, 1]) # [[[1, 4],
# [2, 5],
# [3, 6]],
@@ -1421,7 +1422,8 @@
# pylint: disable=invalid-name
-@tf_export("matrix_transpose", "linalg.transpose")
+@tf_export("linalg.transpose", "matrix_transpose")
+@deprecation.deprecated_endpoints("matrix_transpose")
def matrix_transpose(a, name="matrix_transpose", conjugate=False):
"""Transposes last two dimensions of tensor `a`.
@@ -1429,19 +1431,19 @@
```python
x = tf.constant([[1, 2, 3], [4, 5, 6]])
- tf.matrix_transpose(x) # [[1, 4],
+ tf.linalg.transpose(x) # [[1, 4],
# [2, 5],
# [3, 6]]
x = tf.constant([[1 + 1j, 2 + 2j, 3 + 3j],
[4 + 4j, 5 + 5j, 6 + 6j]])
- tf.matrix_transpose(x, conjugate=True) # [[1 - 1j, 4 - 4j],
+ tf.linalg.transpose(x, conjugate=True) # [[1 - 1j, 4 - 4j],
# [2 - 2j, 5 - 5j],
# [3 - 3j, 6 - 6j]]
# Matrix with two batch dimensions.
# x.shape is [1, 2, 3, 4]
- # tf.matrix_transpose(x) is shape [1, 2, 4, 3]
+ # tf.linalg.transpose(x) is shape [1, 2, 4, 3]
```
Note that `tf.matmul` provides kwargs allowing for transpose of arguments.
@@ -1452,14 +1454,14 @@
tf.matmul(matrix, b, transpose_b=True)
# Inefficient!
- tf.matmul(matrix, tf.matrix_transpose(b))
+ tf.matmul(matrix, tf.linalg.transpose(b))
```
@compatibility(numpy)
In `numpy` transposes are memory-efficient constant time operations as they
simply return a new view of the same data with adjusted `strides`.
- TensorFlow does not support strides, `matrix_transposes` return a new tensor
+ TensorFlow does not support strides, `linalg.transposes` return a new tensor
with the items permuted.
@end_compatibility
@@ -1467,7 +1469,7 @@
a: A `Tensor` with `rank >= 2`.
name: A name for the operation (optional).
conjugate: Optional bool. Setting it to `True` is mathematically equivalent
- to tf.conj(tf.matrix_transpose(input)).
+ to tf.conj(tf.linalg.transpose(input)).
Returns:
A transposed batch matrix `Tensor`.
@@ -1756,7 +1758,8 @@
return (ops.convert_to_tensor(shape, dtype=dtypes.int64, name=name), rank)
-@tf_export("sparse_placeholder")
+@tf_export("sparse.placeholder", "sparse_placeholder")
+@deprecation.deprecated_endpoints("sparse_placeholder")
def sparse_placeholder(dtype, shape=None, name=None):
"""Inserts a placeholder for a sparse tensor that will be always fed.
@@ -1767,8 +1770,8 @@
For example:
```python
- x = tf.sparse_placeholder(tf.float32)
- y = tf.sparse_reduce_sum(x)
+ x = tf.sparse.placeholder(tf.float32)
+ y = tf.sparse.reduce_sum(x)
with tf.Session() as sess:
print(sess.run(y)) # ERROR: will fail because x was not fed.
@@ -2250,7 +2253,8 @@
return result_paddings, result_crops
-@tf_export("space_to_batch")
+@tf_export("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(
input,
@@ -2264,7 +2268,8 @@
space_to_batch.__doc__ = gen_array_ops.space_to_batch.__doc__
-@tf_export("space_to_depth")
+@tf_export("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)
@@ -2272,7 +2277,8 @@
space_to_depth.__doc__ = gen_array_ops.space_to_depth.__doc__
-@tf_export("depth_to_space")
+@tf_export("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)
@@ -2747,7 +2753,8 @@
@tf_export("quantize_v2")
@deprecation.deprecated(
"2017-10-25",
- "`tf.quantize_v2` is deprecated, please use `tf.quantize` instead.")
+ "`tf.quantize_v2` is deprecated, please use `tf.quantization.quantize` "
+ "instead.") # pylint: disable=missing-docstring
def quantize_v2(input, # pylint: disable=redefined-builtin
min_range,
max_range,
@@ -2769,7 +2776,8 @@
# We want to expose tf.quantize instead of tf.quantize_v2; we can deprecate
# tf.quantize_v2 in next version of TensorFlow.
-@tf_export("quantize")
+@tf_export("quantization.quantize", "quantize")
+@deprecation.deprecated_endpoints("quantize")
def quantize(input, # pylint: disable=redefined-builtin
min_range,
max_range,
diff --git a/tensorflow/python/ops/candidate_sampling_ops.py b/tensorflow/python/ops/candidate_sampling_ops.py
index 9ea1ea9..98dde99 100644
--- a/tensorflow/python/ops/candidate_sampling_ops.py
+++ b/tensorflow/python/ops/candidate_sampling_ops.py
@@ -23,10 +23,12 @@
from tensorflow.python.ops import array_ops # pylint: disable=unused-import
from tensorflow.python.ops import gen_candidate_sampling_ops
from tensorflow.python.ops import math_ops # pylint: disable=unused-import
+from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
-@tf_export('nn.uniform_candidate_sampler')
+@tf_export('random.uniform_candidate_sampler', 'nn.uniform_candidate_sampler')
+@deprecation.deprecated_endpoints('nn.uniform_candidate_sampler')
def uniform_candidate_sampler(true_classes, num_true, num_sampled, unique,
range_max, seed=None, name=None):
"""Samples a set of classes using a uniform base distribution.
@@ -82,7 +84,9 @@
seed2=seed2, name=name)
-@tf_export('nn.log_uniform_candidate_sampler')
+@tf_export('random.log_uniform_candidate_sampler',
+ 'nn.log_uniform_candidate_sampler')
+@deprecation.deprecated_endpoints('nn.log_uniform_candidate_sampler')
def log_uniform_candidate_sampler(true_classes, num_true, num_sampled, unique,
range_max, seed=None, name=None):
"""Samples a set of classes using a log-uniform (Zipfian) base distribution.
diff --git a/tensorflow/python/ops/check_ops.py b/tensorflow/python/ops/check_ops.py
index c3cf6e6..d607f1d 100644
--- a/tensorflow/python/ops/check_ops.py
+++ b/tensorflow/python/ops/check_ops.py
@@ -36,6 +36,7 @@
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.util import compat
+from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
NUMERIC_TYPES = frozenset(
@@ -91,7 +92,8 @@
return 'shape=%s dtype=%s' % (tensor.shape, tensor.dtype.name)
-@tf_export('assert_proper_iterable')
+@tf_export('debugging.assert_proper_iterable', 'assert_proper_iterable')
+@deprecation.deprecated_endpoints('assert_proper_iterable')
def assert_proper_iterable(values):
"""Static assert that values is a "proper" iterable.
@@ -119,7 +121,8 @@
'Expected argument "values" to be iterable. Found: %s' % type(values))
-@tf_export('assert_negative')
+@tf_export('debugging.assert_negative', 'assert_negative')
+@deprecation.deprecated_endpoints('assert_negative')
def assert_negative(x, data=None, summarize=None, message=None, name=None):
"""Assert the condition `x < 0` holds element-wise.
@@ -160,7 +163,8 @@
return assert_less(x, zero, data=data, summarize=summarize)
-@tf_export('assert_positive')
+@tf_export('debugging.assert_positive', 'assert_positive')
+@deprecation.deprecated_endpoints('assert_positive')
def assert_positive(x, data=None, summarize=None, message=None, name=None):
"""Assert the condition `x > 0` holds element-wise.
@@ -200,7 +204,8 @@
return assert_less(zero, x, data=data, summarize=summarize)
-@tf_export('assert_non_negative')
+@tf_export('debugging.assert_non_negative', 'assert_non_negative')
+@deprecation.deprecated_endpoints('assert_non_negative')
def assert_non_negative(x, data=None, summarize=None, message=None, name=None):
"""Assert the condition `x >= 0` holds element-wise.
@@ -242,7 +247,8 @@
return assert_less_equal(zero, x, data=data, summarize=summarize)
-@tf_export('assert_non_positive')
+@tf_export('debugging.assert_non_positive', 'assert_non_positive')
+@deprecation.deprecated_endpoints('assert_non_positive')
def assert_non_positive(x, data=None, summarize=None, message=None, name=None):
"""Assert the condition `x <= 0` holds element-wise.
@@ -284,7 +290,7 @@
return assert_less_equal(x, zero, data=data, summarize=summarize)
-@tf_export('assert_equal')
+@tf_export('debugging.assert_equal', 'assert_equal')
def assert_equal(x, y, data=None, summarize=None, message=None, name=None):
"""Assert the condition `x == y` holds element-wise.
@@ -384,7 +390,8 @@
return control_flow_ops.Assert(condition, data, summarize=summarize)
-@tf_export('assert_none_equal')
+@tf_export('debugging.assert_none_equal', 'assert_none_equal')
+@deprecation.deprecated_endpoints('assert_none_equal')
def assert_none_equal(
x, y, data=None, summarize=None, message=None, name=None):
"""Assert the condition `x != y` holds for all elements.
@@ -435,7 +442,8 @@
return control_flow_ops.Assert(condition, data, summarize=summarize)
-@tf_export('assert_near')
+@tf_export('debugging.assert_near', 'assert_near')
+@deprecation.deprecated_endpoints('assert_near')
def assert_near(
x, y, rtol=None, atol=None, data=None, summarize=None, message=None,
name=None):
@@ -513,7 +521,7 @@
return control_flow_ops.Assert(condition, data, summarize=summarize)
-@tf_export('assert_less')
+@tf_export('debugging.assert_less', 'assert_less')
def assert_less(x, y, data=None, summarize=None, message=None, name=None):
"""Assert the condition `x < y` holds element-wise.
@@ -561,7 +569,8 @@
return control_flow_ops.Assert(condition, data, summarize=summarize)
-@tf_export('assert_less_equal')
+@tf_export('debugging.assert_less_equal', 'assert_less_equal')
+@deprecation.deprecated_endpoints('assert_less_equal')
def assert_less_equal(x, y, data=None, summarize=None, message=None, name=None):
"""Assert the condition `x <= y` holds element-wise.
@@ -609,7 +618,7 @@
return control_flow_ops.Assert(condition, data, summarize=summarize)
-@tf_export('assert_greater')
+@tf_export('debugging.assert_greater', 'assert_greater')
def assert_greater(x, y, data=None, summarize=None, message=None, name=None):
"""Assert the condition `x > y` holds element-wise.
@@ -657,7 +666,8 @@
return control_flow_ops.Assert(condition, data, summarize=summarize)
-@tf_export('assert_greater_equal')
+@tf_export('debugging.assert_greater_equal', 'assert_greater_equal')
+@deprecation.deprecated_endpoints('assert_greater_equal')
def assert_greater_equal(x, y, data=None, summarize=None, message=None,
name=None):
"""Assert the condition `x >= y` holds element-wise.
@@ -755,7 +765,7 @@
return control_flow_ops.Assert(condition, data, summarize=summarize)
-@tf_export('assert_rank')
+@tf_export('debugging.assert_rank', 'assert_rank')
def assert_rank(x, rank, data=None, summarize=None, message=None, name=None):
"""Assert `x` has rank equal to `rank`.
@@ -817,7 +827,8 @@
return assert_op
-@tf_export('assert_rank_at_least')
+@tf_export('debugging.assert_rank_at_least', 'assert_rank_at_least')
+@deprecation.deprecated_endpoints('assert_rank_at_least')
def assert_rank_at_least(
x, rank, data=None, summarize=None, message=None, name=None):
"""Assert `x` has rank equal to `rank` or higher.
@@ -948,7 +959,8 @@
return control_flow_ops.Assert(condition, data, summarize=summarize)
-@tf_export('assert_rank_in')
+@tf_export('debugging.assert_rank_in', 'assert_rank_in')
+@deprecation.deprecated_endpoints('assert_rank_in')
def assert_rank_in(
x, ranks, data=None, summarize=None, message=None, name=None):
"""Assert `x` has rank in `ranks`.
@@ -1010,7 +1022,8 @@
return assert_op
-@tf_export('assert_integer')
+@tf_export('debugging.assert_integer', 'assert_integer')
+@deprecation.deprecated_endpoints('assert_integer')
def assert_integer(x, message=None, name=None):
"""Assert that `x` is of integer dtype.
@@ -1048,7 +1061,8 @@
return control_flow_ops.no_op('statically_determined_was_integer')
-@tf_export('assert_type')
+@tf_export('debugging.assert_type', 'assert_type')
+@deprecation.deprecated_endpoints('assert_type')
def assert_type(tensor, tf_type, message=None, name=None):
"""Statically asserts that the given `Tensor` is of the specified type.
@@ -1095,12 +1109,14 @@
return control_flow_ops.cond(is_shorter_than_two, short_result, diff)
-@tf_export('is_numeric_tensor')
+@tf_export('debugging.is_numeric_tensor', 'is_numeric_tensor')
+@deprecation.deprecated_endpoints('is_numeric_tensor')
def is_numeric_tensor(tensor):
return isinstance(tensor, ops.Tensor) and tensor.dtype in NUMERIC_TYPES
-@tf_export('is_non_decreasing')
+@tf_export('debugging.is_non_decreasing', 'is_non_decreasing')
+@deprecation.deprecated_endpoints('is_non_decreasing')
def is_non_decreasing(x, name=None):
"""Returns `True` if `x` is non-decreasing.
@@ -1127,7 +1143,8 @@
return math_ops.reduce_all(math_ops.less_equal(zero, diff))
-@tf_export('is_strictly_increasing')
+@tf_export('debugging.is_strictly_increasing', 'is_strictly_increasing')
+@deprecation.deprecated_endpoints('is_strictly_increasing')
def is_strictly_increasing(x, name=None):
"""Returns `True` if `x` is strictly increasing.
@@ -1202,7 +1219,8 @@
return expected_type
-@tf_export('assert_same_float_dtype')
+@tf_export('debugging.assert_same_float_dtype', 'assert_same_float_dtype')
+@deprecation.deprecated_endpoints('assert_same_float_dtype')
def assert_same_float_dtype(tensors=None, dtype=None):
"""Validate and return float type based on `tensors` and `dtype`.
@@ -1231,7 +1249,8 @@
return dtype
-@tf_export('assert_scalar')
+@tf_export('debugging.assert_scalar', 'assert_scalar')
+@deprecation.deprecated_endpoints('assert_scalar')
def assert_scalar(tensor, name=None):
with ops.name_scope(name, 'assert_scalar', [tensor]) as name_scope:
tensor = ops.convert_to_tensor(tensor, name=name_scope)
diff --git a/tensorflow/python/ops/clip_ops.py b/tensorflow/python/ops/clip_ops.py
index 2946843..4551606 100644
--- a/tensorflow/python/ops/clip_ops.py
+++ b/tensorflow/python/ops/clip_ops.py
@@ -30,6 +30,7 @@
from tensorflow.python.ops import gen_nn_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import numerics
+from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
@@ -76,8 +77,8 @@
return t_max
# TODO(scottzhu): switch to use new implmentation in 2 weeks.
- # return gen_math_ops.clip_by_value(
- # t, clip_value_min, clip_value_max, name=name)
+ # return gen_math_ops.clip_by_value(
+ # t, clip_value_min, clip_value_max, name=name)
# TODO(scottzhu): switch to use new implmentation in 2 weeks.
@@ -159,7 +160,8 @@
return tclip
-@tf_export("global_norm")
+@tf_export("linalg.global_norm", "global_norm")
+@deprecation.deprecated_endpoints("global_norm")
def global_norm(t_list, name=None):
"""Computes the global norm of multiple tensors.
diff --git a/tensorflow/python/ops/confusion_matrix.py b/tensorflow/python/ops/confusion_matrix.py
index c091541..8259142 100644
--- a/tensorflow/python/ops/confusion_matrix.py
+++ b/tensorflow/python/ops/confusion_matrix.py
@@ -26,6 +26,7 @@
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import sparse_ops
+from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
@@ -89,7 +90,8 @@
return labels, predictions
-@tf_export('confusion_matrix')
+@tf_export('train.confusion_matrix', 'confusion_matrix')
+@deprecation.deprecated_endpoints('confusion_matrix')
def confusion_matrix(labels, predictions, num_classes=None, dtype=dtypes.int32,
name=None, weights=None):
"""Computes the confusion matrix from predictions and labels.
diff --git a/tensorflow/python/ops/control_flow_ops.py b/tensorflow/python/ops/control_flow_ops.py
index 87f8bd8..8ad71fe 100644
--- a/tensorflow/python/ops/control_flow_ops.py
+++ b/tensorflow/python/ops/control_flow_ops.py
@@ -60,8 +60,17 @@
from tensorflow.python.util import tf_should_use
from tensorflow.python.util.tf_export import tf_export
+# The while_v2 module.
+_while_v2 = None
ENABLE_COND_V2 = os.getenv("TF_ENABLE_COND_V2", "0") != "0"
+# Note: Setting this to True is not sufficient to switch to the v2 while_loop.
+# Users must also import the while_v2 module to set the _while_v2 module
+# variable above. We do this to avoid a circular dependency:
+# control_flow_ops -> while_v2 -> gradients_impl -> control_flow_ops
+# A ValueError is raised in tf.while_loop if this is set to True and the
+# `_while_v2` module is not set.
+ENABLE_WHILE_V2 = os.getenv("TF_ENABLE_WHILE_V2", "0") != "0"
# We override the 'tuple' for a control flow op, so we keep python's
@@ -97,7 +106,7 @@
# Assert and Print are special symbols in python, so we must
# use an upper-case version of them.
-@tf_export("Assert")
+@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.
@@ -3211,6 +3220,13 @@
```
"""
+ if ENABLE_WHILE_V2 and not context.executing_eagerly():
+ if not _while_v2:
+ raise ValueError("The while_v2 module is not set. Did you forget to "
+ "import tensorflow.python.ops."
+ "while_v2?")
+ return _while_v2.while_loop(cond, body, loop_vars, name)
+
with ops.name_scope(name, "while", loop_vars):
if not loop_vars:
raise ValueError("No loop variables provided")
diff --git a/tensorflow/python/ops/conv2d_benchmark.py b/tensorflow/python/ops/conv2d_benchmark.py
index 28111c2..f40488a 100644
--- a/tensorflow/python/ops/conv2d_benchmark.py
+++ b/tensorflow/python/ops/conv2d_benchmark.py
@@ -63,9 +63,9 @@
An array of tensors to run()
"""
with ops.device("/%s:0" % device):
- inp = variables.Variable(
+ inp = variables.VariableV1(
random_ops.truncated_normal(input_shape, dtype=dtype))
- filt = variables.Variable(
+ filt = variables.VariableV1(
random_ops.truncated_normal(filter_shape, dtype=dtype))
outputs = []
diff --git a/tensorflow/python/ops/data_flow_ops.py b/tensorflow/python/ops/data_flow_ops.py
index 69c0fcb..97b6f3b 100644
--- a/tensorflow/python/ops/data_flow_ops.py
+++ b/tensorflow/python/ops/data_flow_ops.py
@@ -39,6 +39,7 @@
# go/tf-wildcard-import
# pylint: disable=wildcard-import
from tensorflow.python.ops.gen_data_flow_ops import *
+from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
# pylint: enable=wildcard-import
@@ -112,7 +113,8 @@
# pylint: disable=protected-access
-@tf_export("QueueBase")
+@tf_export("io.QueueBase", "QueueBase")
+@deprecation.deprecated_endpoints("QueueBase")
class QueueBase(object):
"""Base class for queue implementations.
@@ -604,7 +606,8 @@
return shared_name
-@tf_export("RandomShuffleQueue")
+@tf_export("io.RandomShuffleQueue", "RandomShuffleQueue")
+@deprecation.deprecated_endpoints("RandomShuffleQueue")
class RandomShuffleQueue(QueueBase):
"""A queue implementation that dequeues elements in a random order.
@@ -746,7 +749,8 @@
super(FIFOQueue, self).__init__(dtypes, shapes, names, queue_ref)
-@tf_export("PaddingFIFOQueue")
+@tf_export("io.PaddingFIFOQueue", "PaddingFIFOQueue")
+@deprecation.deprecated_endpoints("PaddingFIFOQueue")
class PaddingFIFOQueue(QueueBase):
"""A FIFOQueue that supports batching variable-sized tensors by padding.
@@ -820,7 +824,8 @@
super(PaddingFIFOQueue, self).__init__(dtypes, shapes, names, queue_ref)
-@tf_export("PriorityQueue")
+@tf_export("io.PriorityQueue", "PriorityQueue")
+@deprecation.deprecated_endpoints("PriorityQueue")
class PriorityQueue(QueueBase):
"""A queue implementation that dequeues elements in prioritized order.
@@ -1300,7 +1305,9 @@
return out
-@tf_export("SparseConditionalAccumulator")
+@tf_export("sparse.SparseConditionalAccumulator",
+ "SparseConditionalAccumulator")
+@deprecation.deprecated_endpoints("SparseConditionalAccumulator")
class SparseConditionalAccumulator(ConditionalAccumulatorBase):
"""A conditional accumulator for aggregating sparse gradients.
diff --git a/tensorflow/python/ops/distributions/distribution.py b/tensorflow/python/ops/distributions/distribution.py
index 76d9806..12fd039 100644
--- a/tensorflow/python/ops/distributions/distribution.py
+++ b/tensorflow/python/ops/distributions/distribution.py
@@ -25,6 +25,7 @@
import numpy as np
import six
+from tensorflow.python.eager import context
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
@@ -127,6 +128,18 @@
return old_str + "\n\n" + append_str
+def _convert_to_tensor(value, name=None, preferred_dtype=None):
+ """Converts to tensor avoiding an eager bug that loses float precision."""
+ # TODO(b/116672045): Remove this function.
+ if (context.executing_eagerly() and preferred_dtype is not None and
+ (preferred_dtype.is_integer or preferred_dtype.is_bool)):
+ v = ops.convert_to_tensor(value, name=name)
+ if v.dtype.is_floating:
+ return v
+ return ops.convert_to_tensor(
+ value, name=name, preferred_dtype=preferred_dtype)
+
+
class _DistributionMeta(abc.ABCMeta):
def __new__(mcs, classname, baseclasses, attrs):
@@ -741,7 +754,8 @@
def _call_log_prob(self, value, name, **kwargs):
with self._name_scope(name, values=[value]):
- value = ops.convert_to_tensor(value, name="value")
+ value = _convert_to_tensor(
+ value, name="value", preferred_dtype=self.dtype)
try:
return self._log_prob(value, **kwargs)
except NotImplementedError as original_exception:
@@ -769,7 +783,8 @@
def _call_prob(self, value, name, **kwargs):
with self._name_scope(name, values=[value]):
- value = ops.convert_to_tensor(value, name="value")
+ value = _convert_to_tensor(
+ value, name="value", preferred_dtype=self.dtype)
try:
return self._prob(value, **kwargs)
except NotImplementedError as original_exception:
@@ -797,7 +812,8 @@
def _call_log_cdf(self, value, name, **kwargs):
with self._name_scope(name, values=[value]):
- value = ops.convert_to_tensor(value, name="value")
+ value = _convert_to_tensor(
+ value, name="value", preferred_dtype=self.dtype)
try:
return self._log_cdf(value, **kwargs)
except NotImplementedError as original_exception:
@@ -835,7 +851,8 @@
def _call_cdf(self, value, name, **kwargs):
with self._name_scope(name, values=[value]):
- value = ops.convert_to_tensor(value, name="value")
+ value = _convert_to_tensor(
+ value, name="value", preferred_dtype=self.dtype)
try:
return self._cdf(value, **kwargs)
except NotImplementedError as original_exception:
@@ -870,7 +887,8 @@
def _call_log_survival_function(self, value, name, **kwargs):
with self._name_scope(name, values=[value]):
- value = ops.convert_to_tensor(value, name="value")
+ value = _convert_to_tensor(
+ value, name="value", preferred_dtype=self.dtype)
try:
return self._log_survival_function(value, **kwargs)
except NotImplementedError as original_exception:
@@ -909,7 +927,8 @@
def _call_survival_function(self, value, name, **kwargs):
with self._name_scope(name, values=[value]):
- value = ops.convert_to_tensor(value, name="value")
+ value = _convert_to_tensor(
+ value, name="value", preferred_dtype=self.dtype)
try:
return self._survival_function(value, **kwargs)
except NotImplementedError as original_exception:
@@ -963,7 +982,8 @@
def _call_quantile(self, value, name, **kwargs):
with self._name_scope(name, values=[value]):
- value = ops.convert_to_tensor(value, name="value")
+ value = _convert_to_tensor(
+ value, name="value", preferred_dtype=self.dtype)
return self._quantile(value, **kwargs)
def quantile(self, value, name="quantile"):
diff --git a/tensorflow/python/ops/distributions/exponential.py b/tensorflow/python/ops/distributions/exponential.py
index 4325a14..02129b5 100644
--- a/tensorflow/python/ops/distributions/exponential.py
+++ b/tensorflow/python/ops/distributions/exponential.py
@@ -114,6 +114,9 @@
def rate(self):
return self._rate
+ def _log_survival_function(self, value):
+ return self._log_prob(value) - math_ops.log(self._rate)
+
def _sample_n(self, n, seed=None):
shape = array_ops.concat([[n], array_ops.shape(self._rate)], 0)
# Uniform variates must be sampled from the open-interval `(0, 1)` rather
diff --git a/tensorflow/python/ops/distributions/special_math.py b/tensorflow/python/ops/distributions/special_math.py
index 31b7a36..ccc667c 100644
--- a/tensorflow/python/ops/distributions/special_math.py
+++ b/tensorflow/python/ops/distributions/special_math.py
@@ -12,6 +12,62 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
+
+# Functions "ndtr" and "ndtri" are derived from calculations made in:
+# https://root.cern.ch/doc/v608/SpecFuncCephesInv_8cxx_source.html
+# In the following email exchange, the author gives his consent to redistribute
+# derived works under an Apache 2.0 license.
+#
+# From: Stephen Moshier <steve@moshier.net>
+# Date: Sat, Jun 9, 2018 at 2:36 PM
+# Subject: Re: Licensing cephes under Apache (BSD-like) license.
+# To: rif <rif@google.com>
+#
+#
+#
+# Hello Rif,
+#
+# Yes, Google may distribute Cephes files under the Apache 2 license.
+#
+# If clarification is needed, I do not favor BSD over other free licenses.
+# I would agree that Apache 2 seems to cover the concern you mentioned
+# about sublicensees.
+#
+# Best wishes for good luck with your projects!
+# Steve Moshier
+#
+#
+#
+# On Thu, 31 May 2018, rif wrote:
+#
+# > Hello Steve.
+# > My name is Rif. I work on machine learning software at Google.
+# >
+# > Your cephes software continues to be incredibly useful and widely used. I
+# > was wondering whether it would be permissible for us to use the Cephes code
+# > under the Apache 2.0 license, which is extremely similar in permissions to
+# > the BSD license (Wikipedia comparisons). This would be quite helpful to us
+# > in terms of avoiding multiple licenses on software.
+# >
+# > I'm sorry to bother you with this (I can imagine you're sick of hearing
+# > about this by now), but I want to be absolutely clear we're on the level and
+# > not misusing your important software. In former conversation with Eugene
+# > Brevdo (ebrevdo@google.com), you wrote "If your licensing is similar to BSD,
+# > the formal way that has been handled is simply to add a statement to the
+# > effect that you are incorporating the Cephes software by permission of the
+# > author." I wanted to confirm that (a) we could use the Apache license, (b)
+# > that we don't need to (and probably you don't want to) keep getting
+# > contacted about individual uses, because your intent is generally to allow
+# > this software to be reused under "BSD-like" license, and (c) you're OK
+# > letting incorporators decide whether a license is sufficiently BSD-like?
+# >
+# > Best,
+# >
+# > rif
+# >
+# >
+# >
+
"""Special Math Ops."""
from __future__ import absolute_import
@@ -135,7 +191,7 @@
# Constants used in piece-wise rational approximations. Taken from the cephes
# library:
- # https://github.com/scipy/scipy/blob/master/scipy/special/cephes/ndtri.c
+ # https://root.cern.ch/doc/v608/SpecFuncCephesInv_8cxx_source.html
p0 = list(reversed([-5.99633501014107895267E1,
9.80010754185999661536E1,
-5.66762857469070293439E1,
@@ -305,7 +361,8 @@
else:
raise TypeError("x.dtype=%s is not supported." % x.dtype)
- # The basic idea here was ported from py/scipy/special/cephes/ndtr.c.
+ # The basic idea here was ported from:
+ # https://root.cern.ch/doc/v608/SpecFuncCephesInv_8cxx_source.html
# We copy the main idea, with a few changes
# * For x >> 1, and X ~ Normal(0, 1),
# Log[P[X < x]] = Log[1 - P[X < -x]] approx -P[X < -x],
diff --git a/tensorflow/python/ops/embedding_ops.py b/tensorflow/python/ops/embedding_ops.py
index 60d73a1..6263041 100644
--- a/tensorflow/python/ops/embedding_ops.py
+++ b/tensorflow/python/ops/embedding_ops.py
@@ -550,11 +550,9 @@
raise ValueError('Missing embedding_weights %s.' % embedding_weights)
dtype = sparse_weights.dtype if sparse_weights is not None else None
- if not isinstance(embedding_weights[0],
- resource_variable_ops.ResourceVariable):
- embedding_weights = [
- ops.convert_to_tensor(w, dtype=dtype) for w in embedding_weights
- ]
+ embedding_weights = [
+ ops.convert_to_tensor(w, dtype=dtype) for w in embedding_weights
+ ]
with ops.name_scope(name, 'embedding_lookup',
embedding_weights + [sparse_ids,
diff --git a/tensorflow/python/ops/gradients_test.py b/tensorflow/python/ops/gradients_test.py
index 4f6e5dc..3c9b7a0 100644
--- a/tensorflow/python/ops/gradients_test.py
+++ b/tensorflow/python/ops/gradients_test.py
@@ -273,7 +273,7 @@
def testVariableRefGradient(self):
with ops.Graph().as_default():
init = constant_op.constant(100.0)
- var = variables.Variable(init)
+ var = variables.VariableV1(init)
gradient = gradients.gradients(var._ref(), var)
self.assertIsNotNone(gradient)
diff --git a/tensorflow/python/ops/init_ops.py b/tensorflow/python/ops/init_ops.py
index fff3d9b..65bb77b 100644
--- a/tensorflow/python/ops/init_ops.py
+++ b/tensorflow/python/ops/init_ops.py
@@ -43,6 +43,7 @@
from tensorflow.python.ops import linalg_ops_impl
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import random_ops
+from tensorflow.python.util import deprecation
from tensorflow.python.util.deprecation import deprecated
from tensorflow.python.util.deprecation import deprecated_arg_values
from tensorflow.python.util.tf_export import tf_export
@@ -341,6 +342,7 @@
@tf_export("initializers.uniform_unit_scaling",
"uniform_unit_scaling_initializer")
+@deprecation.deprecated_endpoints("uniform_unit_scaling_initializer")
class UniformUnitScaling(Initializer):
"""Initializer that generates tensors without scaling variance.
@@ -401,6 +403,7 @@
@tf_export("keras.initializers.VarianceScaling",
"initializers.variance_scaling", "variance_scaling_initializer")
+@deprecation.deprecated_endpoints("variance_scaling_initializer")
class VarianceScaling(Initializer):
"""Initializer capable of adapting its scale to the shape of weights tensors.
@@ -494,6 +497,7 @@
@tf_export("keras.initializers.Orthogonal", "initializers.orthogonal",
"orthogonal_initializer", "keras.initializers.orthogonal")
+@deprecation.deprecated_endpoints("orthogonal_initializer")
class Orthogonal(Initializer):
"""Initializer that generates an orthogonal matrix.
@@ -1149,6 +1153,7 @@
@tf_export("glorot_normal_initializer", "keras.initializers.glorot_normal",
"initializers.glorot_normal")
+@deprecation.deprecated_endpoints("glorot_normal_initializer")
class GlorotNormal(VarianceScaling):
"""The Glorot normal initializer, also called Xavier normal initializer.
diff --git a/tensorflow/python/ops/linalg_ops.py b/tensorflow/python/ops/linalg_ops.py
index f4a9356..bf4354f 100644
--- a/tensorflow/python/ops/linalg_ops.py
+++ b/tensorflow/python/ops/linalg_ops.py
@@ -80,6 +80,7 @@
@tf_export('cholesky_solve', 'linalg.cholesky_solve')
+@deprecation.deprecated_endpoints('cholesky_solve')
def cholesky_solve(chol, rhs, name=None):
"""Solves systems of linear eqns `A X = RHS`, given Cholesky factorizations.
@@ -167,7 +168,8 @@
name=name)
-@tf_export('matrix_solve_ls', 'linalg.lstsq')
+@tf_export('linalg.lstsq', 'matrix_solve_ls')
+@deprecation.deprecated_endpoints('matrix_solve_ls')
def matrix_solve_ls(matrix, rhs, l2_regularizer=0.0, fast=True, name=None):
r"""Solves one or more linear least-squares problems.
@@ -220,7 +222,7 @@
squares sense.
Raises:
- NotImplementedError: matrix_solve_ls is currently disabled for complex128
+ NotImplementedError: linalg.lstsq is currently disabled for complex128
and l2_regularizer != 0 due to poor accuracy.
"""
@@ -303,7 +305,8 @@
matrix, rhs, l2_regularizer, fast=fast, name=name)
-@tf_export('self_adjoint_eig', 'linalg.eigh')
+@tf_export('linalg.eigh', 'self_adjoint_eig')
+@deprecation.deprecated_endpoints('self_adjoint_eig')
def self_adjoint_eig(tensor, name=None):
"""Computes the eigen decomposition of a batch of self-adjoint matrices.
@@ -325,12 +328,13 @@
return e, v
-@tf_export('self_adjoint_eigvals', 'linalg.eigvalsh')
+@tf_export('linalg.eigvalsh', 'self_adjoint_eigvals')
+@deprecation.deprecated_endpoints('self_adjoint_eigvals')
def self_adjoint_eigvals(tensor, name=None):
"""Computes the eigenvalues of one or more self-adjoint matrices.
Note: If your program backpropagates through this function, you should replace
- it with a call to tf.self_adjoint_eig (possibly ignoring the second output) to
+ it with a call to tf.linalg.eigvalsh (possibly ignoring the second output) to
avoid computing the eigen decomposition twice. This is because the
eigenvectors are used to compute the gradient w.r.t. the eigenvalues. See
_SelfAdjointEigV2Grad in linalg_grad.py.
@@ -348,6 +352,7 @@
@tf_export('svd', 'linalg.svd')
+@deprecation.deprecated_endpoints('svd')
def svd(tensor, full_matrices=False, compute_uv=True, name=None):
r"""Computes the singular value decompositions of one or more matrices.
diff --git a/tensorflow/python/ops/lookup_ops.py b/tensorflow/python/ops/lookup_ops.py
index 5443699..cffaa98 100644
--- a/tensorflow/python/ops/lookup_ops.py
+++ b/tensorflow/python/ops/lookup_ops.py
@@ -59,7 +59,7 @@
return tables_initializer(name)
-@tf_export("tables_initializer")
+@tf_export("initializers.tables_initializer", "tables_initializer")
def tables_initializer(name="init_all_tables"):
"""Returns an Op that initializes all tables of the default graph.
diff --git a/tensorflow/python/ops/manip_ops.py b/tensorflow/python/ops/manip_ops.py
index 6633565..d9d0728 100644
--- a/tensorflow/python/ops/manip_ops.py
+++ b/tensorflow/python/ops/manip_ops.py
@@ -19,11 +19,13 @@
from __future__ import print_function
from tensorflow.python.ops import gen_manip_ops as _gen_manip_ops
+from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
# pylint: disable=protected-access
-@tf_export('manip.roll')
+@tf_export('roll', 'manip.roll')
+@deprecation.deprecated_endpoints('manip.roll')
def roll(input, shift, axis): # pylint: disable=redefined-builtin
return _gen_manip_ops.roll(input, shift, axis)
diff --git a/tensorflow/python/ops/math_grad.py b/tensorflow/python/ops/math_grad.py
index 8e11c4b..35278d9 100644
--- a/tensorflow/python/ops/math_grad.py
+++ b/tensorflow/python/ops/math_grad.py
@@ -516,6 +516,40 @@
return grad * math_ops.reciprocal(1 + x)
+@ops.RegisterGradient("Xlogy")
+def _XLogyGrad(op, grad):
+ """Returns gradient of xlogy(x, y) with respect to x and y."""
+ x = op.inputs[0]
+ y = op.inputs[1]
+ sx = array_ops.shape(x)
+ sy = array_ops.shape(y)
+ rx, ry = gen_array_ops.broadcast_gradient_args(sx, sy)
+ with ops.control_dependencies([grad]):
+ not_zero_x = math_ops.cast(
+ math_ops.not_equal(x, math_ops.cast(0., dtype=x.dtype)), dtype=x.dtype)
+ partial_x = gen_math_ops.xlogy(not_zero_x, y)
+ partial_y = gen_math_ops.xdivy(x, y)
+ return (array_ops.reshape(math_ops.reduce_sum(partial_x * grad, rx), sx),
+ array_ops.reshape(math_ops.reduce_sum(partial_y * grad, ry), sy))
+
+
+@ops.RegisterGradient("Xdivy")
+def _XDivyGrad(op, grad):
+ """Returns gradient of xdivy(x, y) with respect to x and y."""
+ x = op.inputs[0]
+ y = op.inputs[1]
+ sx = array_ops.shape(x)
+ sy = array_ops.shape(y)
+ rx, ry = gen_array_ops.broadcast_gradient_args(sx, sy)
+ with ops.control_dependencies([grad]):
+ not_zero_x = math_ops.cast(
+ math_ops.not_equal(x, math_ops.cast(0., dtype=x.dtype)), dtype=x.dtype)
+ partial_x = gen_math_ops.xdivy(not_zero_x, y)
+ partial_y = gen_math_ops.xdivy(math_ops.negative(x), y**2)
+ return (array_ops.reshape(math_ops.reduce_sum(partial_x * grad, rx), sx),
+ array_ops.reshape(math_ops.reduce_sum(partial_y * grad, ry), sy))
+
+
@ops.RegisterGradient("Sinh")
def _SinhGrad(op, grad):
"""Returns grad * cosh(x)."""
diff --git a/tensorflow/python/ops/math_grad_test.py b/tensorflow/python/ops/math_grad_test.py
index 7110e09..9cfb050 100644
--- a/tensorflow/python/ops/math_grad_test.py
+++ b/tensorflow/python/ops/math_grad_test.py
@@ -256,5 +256,93 @@
self.assertAllClose(dy.eval(), np.zeros(y.shape.as_list()))
+class XlogyTest(test.TestCase):
+
+ def _xlogy_gradients(self, x, y):
+ xlogy_xgrad = self.evaluate(gradients.gradients(math_ops.xlogy(x, y), x)[0])
+ xlogy_ygrad = self.evaluate(gradients.gradients(math_ops.xlogy(x, y), y)[0])
+ return xlogy_xgrad, xlogy_ygrad
+
+ def testNonZeroValuesGrad(self):
+ for dtype in [dtypes.float16, dtypes.float32, dtypes.float64]:
+ x = constant_op.constant(0.1, dtype=dtype)
+ y = constant_op.constant(3.1, dtype=dtype)
+ xlogy_xgrad, xlogy_ygrad = self._xlogy_gradients(x, y)
+ xlogy_expected_xgrad = self.evaluate(math_ops.log(y))
+ xlogy_expected_ygrad = self.evaluate(x / y)
+ self.assertAllClose(xlogy_expected_xgrad, xlogy_xgrad)
+ self.assertAllClose(xlogy_expected_ygrad, xlogy_ygrad)
+
+ def testZeroXGrad(self):
+ for dtype in [dtypes.float16, dtypes.float32, dtypes.float64]:
+ x = constant_op.constant(0., dtype=dtype)
+ y = constant_op.constant(3.1, dtype=dtype)
+ xlogy_xgrad, xlogy_ygrad = self._xlogy_gradients(x, y)
+ zero = self.evaluate(x)
+ self.assertAllClose(zero, xlogy_xgrad)
+ self.assertAllClose(zero, xlogy_ygrad)
+
+ def testZeroYGrad(self):
+ for dtype in [dtypes.float16, dtypes.float32, dtypes.float64]:
+ x = constant_op.constant(0.1, dtype=dtype)
+ y = constant_op.constant(0., dtype=dtype)
+ xlogy_xgrad, xlogy_ygrad = self._xlogy_gradients(x, y)
+ self.assertAllClose(-np.inf, xlogy_xgrad)
+ self.assertAllClose(np.inf, xlogy_ygrad)
+
+ def testZeroXYGrad(self):
+ for dtype in [dtypes.float16, dtypes.float32, dtypes.float64]:
+ x = constant_op.constant(0., dtype=dtype)
+ y = constant_op.constant(0., dtype=dtype)
+ xlogy_xgrad, xlogy_ygrad = self._xlogy_gradients(x, y)
+ zero = self.evaluate(x)
+ self.assertAllClose(zero, xlogy_xgrad)
+ self.assertAllClose(zero, xlogy_ygrad)
+
+
+class XdivyTest(test.TestCase):
+
+ def _xdivy_gradients(self, x, y):
+ xdivy_xgrad = self.evaluate(gradients.gradients(math_ops.xdivy(x, y), x)[0])
+ xdivy_ygrad = self.evaluate(gradients.gradients(math_ops.xdivy(x, y), y)[0])
+ return xdivy_xgrad, xdivy_ygrad
+
+ def testNonZeroValuesGrad(self):
+ for dtype in [dtypes.float16, dtypes.float32, dtypes.float64]:
+ x = constant_op.constant(0.1, dtype=dtype)
+ y = constant_op.constant(3.1, dtype=dtype)
+ xdivy_xgrad, xdivy_ygrad = self._xdivy_gradients(x, y)
+ xdivy_expected_xgrad = self.evaluate(1 / y)
+ xdivy_expected_ygrad = self.evaluate(-x / y**2)
+ self.assertAllClose(xdivy_expected_xgrad, xdivy_xgrad)
+ self.assertAllClose(xdivy_expected_ygrad, xdivy_ygrad)
+
+ def testZeroXGrad(self):
+ for dtype in [dtypes.float16, dtypes.float32, dtypes.float64]:
+ x = constant_op.constant(0., dtype=dtype)
+ y = constant_op.constant(3.1, dtype=dtype)
+ xdivy_xgrad, xdivy_ygrad = self._xdivy_gradients(x, y)
+ zero = self.evaluate(x)
+ self.assertAllClose(zero, xdivy_xgrad)
+ self.assertAllClose(zero, xdivy_ygrad)
+
+ def testZeroYGrad(self):
+ for dtype in [dtypes.float16, dtypes.float32, dtypes.float64]:
+ x = constant_op.constant(0.1, dtype=dtype)
+ y = constant_op.constant(0., dtype=dtype)
+ xdivy_xgrad, xdivy_ygrad = self._xdivy_gradients(x, y)
+ self.assertAllClose(np.inf, xdivy_xgrad)
+ self.assertAllClose(-np.inf, xdivy_ygrad)
+
+ def testZeroXYGrad(self):
+ for dtype in [dtypes.float16, dtypes.float32, dtypes.float64]:
+ x = constant_op.constant(0., dtype=dtype)
+ y = constant_op.constant(0., dtype=dtype)
+ xdivy_xgrad, xdivy_ygrad = self._xdivy_gradients(x, y)
+ zero = self.evaluate(x)
+ self.assertAllClose(zero, xdivy_xgrad)
+ self.assertAllClose(zero, xdivy_ygrad)
+
+
if __name__ == "__main__":
test.main()
diff --git a/tensorflow/python/ops/math_ops.py b/tensorflow/python/ops/math_ops.py
index f57abf6..83b8b5a 100644
--- a/tensorflow/python/ops/math_ops.py
+++ b/tensorflow/python/ops/math_ops.py
@@ -70,7 +70,7 @@
# pylint: disable=redefined-builtin
-@tf_export("argmax")
+@tf_export("math.argmax", "argmax")
@deprecation.deprecated_args(None, "Use the `axis` argument instead",
"dimension")
@_set_doc(
@@ -88,7 +88,7 @@
return gen_math_ops.arg_max(input, axis, name=name, output_type=output_type)
-@tf_export("argmin")
+@tf_export("math.argmin", "argmin")
@deprecation.deprecated_args(None, "Use the `axis` argument instead",
"dimension")
@_set_doc(
@@ -111,7 +111,7 @@
# pylint: disable=anomalous-backslash-in-string,protected-access
# pylint: disable=g-docstring-has-escape
-@tf_export("abs")
+@tf_export("math.abs", "abs")
def abs(x, name=None): # pylint: disable=redefined-builtin
r"""Computes the absolute value of a tensor.
@@ -186,7 +186,7 @@
return _div_python2(self.x, y, self.name)
-@tf_export("divide")
+@tf_export("math.divide", "divide")
def divide(x, y, name=None):
"""Computes Python style division of `x` by `y`."""
@@ -198,7 +198,7 @@
return x / y
-@tf_export("multiply")
+@tf_export("math.multiply", "multiply")
def multiply(x, y, name=None):
return gen_math_ops.mul(x, y, name)
@@ -218,7 +218,7 @@
gen_math_ops.mul.__doc__ + ("" if _mul.__doc__ is None else _mul.__doc__))
-@tf_export("subtract")
+@tf_export("math.subtract", "subtract")
def subtract(x, y, name=None):
return gen_math_ops.sub(x, y, name)
@@ -239,7 +239,7 @@
# pylint: disable=g-docstring-has-escape
-@tf_export("negative")
+@tf_export("math.negative", "negative")
def negative(x, name=None):
"""Computes numerical negative value element-wise.
@@ -288,7 +288,7 @@
# pylint: enable=g-docstring-has-escape
-@tf_export("sign")
+@tf_export("math.sign", "sign")
def sign(x, name=None):
"""Returns an element-wise indication of the sign of a number.
@@ -319,7 +319,7 @@
return gen_math_ops.sign(x, name=name)
-@tf_export("square")
+@tf_export("math.square", "square")
def square(x, name=None):
r"""Computes square of x element-wise.
@@ -342,7 +342,7 @@
return gen_math_ops.square(x, name=name)
-@tf_export("sqrt")
+@tf_export("math.sqrt", "sqrt")
def sqrt(x, name=None):
r"""Computes square root of x element-wise.
@@ -365,7 +365,8 @@
return gen_math_ops.sqrt(x, name=name)
-@tf_export("erf")
+@tf_export("math.erf", "erf")
+@deprecation.deprecated_endpoints("erf")
def erf(x, name=None):
"""Computes the Gauss error function of `x` element-wise.
@@ -386,7 +387,7 @@
return gen_math_ops.erf(x, name=name)
-@tf_export("scalar_mul")
+@tf_export("math.scalar_mul", "scalar_mul")
def scalar_mul(scalar, x):
"""Multiplies a scalar times a `Tensor` or `IndexedSlices` object.
@@ -416,7 +417,7 @@
raise ValueError("Only scalar multiply works, got shape %s" % shape)
-@tf_export("pow")
+@tf_export("math.pow", "pow")
def pow(x, y, name=None): # pylint: disable=redefined-builtin
r"""Computes the power of one value to another.
@@ -444,7 +445,7 @@
# pylint: disable=redefined-builtin,redefined-outer-name
-@tf_export("complex")
+@tf_export("dtypes.complex", "complex")
def complex(real, imag, name=None):
r"""Converts two real numbers to a complex number.
@@ -486,7 +487,8 @@
return gen_math_ops._complex(real, imag, Tout=Tout, name=name)
-@tf_export("real")
+@tf_export("math.real", "real")
+@deprecation.deprecated_endpoints("real")
def real(input, name=None):
r"""Returns the real part of a complex (or real) tensor.
@@ -517,7 +519,8 @@
return input
-@tf_export("imag")
+@tf_export("math.imag", "imag")
+@deprecation.deprecated_endpoints("imag")
def imag(input, name=None):
r"""Returns the imaginary part of a complex (or real) tensor.
@@ -547,7 +550,8 @@
return array_ops.zeros_like(input)
-@tf_export("angle")
+@tf_export("math.angle", "angle")
+@deprecation.deprecated_endpoints("angle")
def angle(input, name=None):
r"""Returns the element-wise argument of a complex (or real) tensor.
@@ -586,7 +590,7 @@
# pylint: enable=redefined-outer-name,redefined-builtin
-@tf_export("round")
+@tf_export("math.round", "round")
def round(x, name=None): # pylint: disable=redefined-builtin
"""Rounds the values of a tensor to the nearest integer, element-wise.
@@ -613,7 +617,7 @@
return gen_math_ops.round(x, name=name)
-@tf_export("cast")
+@tf_export("dtypes.cast", "cast")
def cast(x, dtype, name=None):
"""Casts a tensor to a new type.
@@ -676,7 +680,7 @@
return x
-@tf_export("saturate_cast")
+@tf_export("dtypes.saturate_cast", "saturate_cast")
def saturate_cast(value, dtype, name=None):
"""Performs a safe saturating cast of `value` to `dtype`.
@@ -995,7 +999,7 @@
return gen_math_ops.floor_div(x, y, name=name)
-@tf_export("truediv")
+@tf_export("math.truediv", "truediv")
def truediv(x, y, name=None):
"""Divides x / y elementwise (using Python 3 division operator semantics).
@@ -1006,7 +1010,7 @@
arguments are cast to floating types first. This op is generated by normal
`x / y` division in Python 3 and in Python 2.7 with
`from __future__ import division`. If you want integer division that rounds
- down, use `x // y` or `tf.floordiv`.
+ down, use `x // y` or `tf.math.floordiv`.
`x` and `y` must have the same numeric type. If the inputs are floating
point, the output will have the same type. If the inputs are integral, the
@@ -1078,7 +1082,8 @@
# TODO(aselle): Deprecate this once all internal functionality uses
# tf.truncatediv
-@tf_export("floordiv")
+@tf_export("math.floordiv", "floordiv")
+@deprecation.deprecated_endpoints("floordiv")
def floordiv(x, y, name=None):
"""Divides `x / y` elementwise, rounding toward the most negative integer.
@@ -1151,7 +1156,8 @@
_OverrideBinaryOperatorHelper(pow, "pow")
-@tf_export("logical_xor")
+@tf_export("math.logical_xor", "logical_xor")
+@deprecation.deprecated_endpoints("logical_xor")
def logical_xor(x, y, name="LogicalXor"):
"""x ^ y = (x | y) & ~(x & y)."""
# TODO(alemi) Make this a cwise op if people end up relying on it.
@@ -1277,7 +1283,7 @@
return output
-@tf_export("reduce_sum")
+@tf_export("math.reduce_sum", "reduce_sum")
@deprecation.deprecated_args(
None, "keep_dims is deprecated, use keepdims instead", "keep_dims")
def reduce_sum(input_tensor,
@@ -1339,7 +1345,7 @@
name=name))
-@tf_export("count_nonzero")
+@tf_export("math.count_nonzero", "count_nonzero")
@deprecation.deprecated_args(
None, "keep_dims is deprecated, use keepdims instead", "keep_dims")
def count_nonzero(input_tensor,
@@ -1417,7 +1423,7 @@
dtype=dtype)
-@tf_export("reduce_mean")
+@tf_export("math.reduce_mean", "reduce_mean")
@deprecation.deprecated_args(
None, "keep_dims is deprecated, use keepdims instead", "keep_dims")
def reduce_mean(input_tensor,
@@ -1489,7 +1495,7 @@
name=name))
-@tf_export("reduce_prod")
+@tf_export("math.reduce_prod", "reduce_prod")
@deprecation.deprecated_args(
None, "keep_dims is deprecated, use keepdims instead", "keep_dims")
def reduce_prod(input_tensor,
@@ -1539,7 +1545,7 @@
name=name))
-@tf_export("reduce_min")
+@tf_export("math.reduce_min", "reduce_min")
@deprecation.deprecated_args(
None, "keep_dims is deprecated, use keepdims instead", "keep_dims")
def reduce_min(input_tensor,
@@ -1588,7 +1594,7 @@
name=name))
-@tf_export("reduce_max")
+@tf_export("math.reduce_max", "reduce_max")
@deprecation.deprecated_args(
None, "keep_dims is deprecated, use keepdims instead", "keep_dims")
def reduce_max(input_tensor,
@@ -1637,7 +1643,7 @@
name=name))
-@tf_export("reduce_all")
+@tf_export("math.reduce_all", "reduce_all")
@deprecation.deprecated_args(
None, "keep_dims is deprecated, use keepdims instead", "keep_dims")
def reduce_all(input_tensor,
@@ -1695,7 +1701,7 @@
name=name))
-@tf_export("reduce_any")
+@tf_export("math.reduce_any", "reduce_any")
@deprecation.deprecated_args(
None, "keep_dims is deprecated, use keepdims instead", "keep_dims")
def reduce_any(input_tensor,
@@ -1753,7 +1759,7 @@
name=name))
-@tf_export("reduce_logsumexp")
+@tf_export("math.reduce_logsumexp", "reduce_logsumexp")
@deprecation.deprecated_args(
None, "keep_dims is deprecated, use keepdims instead", "keep_dims")
def reduce_logsumexp(input_tensor,
@@ -1827,7 +1833,8 @@
return _may_reduce_to_scalar(keepdims, axis, reduction_indices, result)
-@tf_export("trace", "linalg.trace")
+@tf_export("linalg.trace", "trace")
+@deprecation.deprecated_endpoints("trace")
def trace(x, name=None):
"""Compute the trace of a tensor `x`.
@@ -1841,12 +1848,12 @@
```python
x = tf.constant([[1, 2], [3, 4]])
- tf.trace(x) # 5
+ tf.linalg.trace(x) # 5
x = tf.constant([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
- tf.trace(x) # 15
+ tf.linalg.trace(x) # 15
x = tf.constant([[[1, 2, 3],
[4, 5, 6],
@@ -1854,7 +1861,7 @@
[[-1, -2, -3],
[-4, -5, -6],
[-7, -8, -9]]])
- tf.trace(x) # [15, -15]
+ tf.linalg.trace(x) # [15, -15]
```
Args:
@@ -1869,7 +1876,7 @@
return reduce_sum(array_ops.matrix_diag_part(x), [-1], name=name)
-@tf_export("matmul")
+@tf_export("linalg.matmul", "matmul")
def matmul(a,
b,
transpose_a=False,
@@ -2131,7 +2138,7 @@
return casted_outputs
-@tf_export("add_n")
+@tf_export("math.add_n", "add_n")
def add_n(inputs, name=None):
"""Adds all input tensors element-wise.
@@ -2166,14 +2173,15 @@
return gen_math_ops.add_n(inputs, name=name)
-@tf_export("accumulate_n")
+@tf_export("math.accumulate_n", "accumulate_n")
+@deprecation.deprecated_endpoints("accumulate_n")
def accumulate_n(inputs, shape=None, tensor_dtype=None, name=None):
"""Returns the element-wise sum of a list of tensors.
Optionally, pass `shape` and `tensor_dtype` for shape and type checking,
otherwise, these are inferred.
- `tf.accumulate_n` performs the same operation as `tf.add_n`, but does not
+ `tf.math.accumulate_n` performs the same operation as `tf.add_n`, but does not
wait for all of its inputs to be ready before beginning to sum. This can
save memory if inputs are ready at different times, since minimum temporary
storage is proportional to the output size rather than the inputs size.
@@ -2185,10 +2193,10 @@
```python
a = tf.constant([[1, 2], [3, 4]])
b = tf.constant([[5, 0], [0, 6]])
- tf.accumulate_n([a, b, a]) # [[7, 4], [6, 14]]
+ tf.math.accumulate_n([a, b, a]) # [[7, 4], [6, 14]]
# Explicitly pass shape and type
- tf.accumulate_n([a, b, a], shape=[2, 2], tensor_dtype=tf.int32)
+ tf.math.accumulate_n([a, b, a], shape=[2, 2], tensor_dtype=tf.int32)
# [[7, 4],
# [6, 14]]
```
@@ -2252,7 +2260,7 @@
return [grad] * len(op.inputs)
-@tf_export("nn.sigmoid", "sigmoid")
+@tf_export("math.sigmoid", "nn.sigmoid", "sigmoid")
def sigmoid(x, name=None):
"""Computes sigmoid of `x` element-wise.
@@ -2275,7 +2283,8 @@
return gen_math_ops.sigmoid(x, name=name)
-@tf_export("log_sigmoid")
+@tf_export("math.log_sigmoid", "log_sigmoid")
+@deprecation.deprecated_endpoints("log_sigmoid")
def log_sigmoid(x, name=None):
"""Computes log sigmoid of `x` element-wise.
@@ -2294,7 +2303,7 @@
return gen_math_ops.neg(gen_nn_ops.softplus(-x), name=name)
-@tf_export("nn.tanh", "tanh")
+@tf_export("math.tanh", "nn.tanh", "tanh")
def tanh(x, name=None):
"""Computes hyperbolic tangent of `x` element-wise.
@@ -2315,7 +2324,8 @@
return gen_math_ops.tanh(x, name=name)
-@tf_export("bincount")
+@tf_export("math.bincount", "bincount")
+@deprecation.deprecated_endpoints("bincount")
def bincount(arr,
weights=None,
minlength=None,
@@ -2362,7 +2372,7 @@
return gen_math_ops.bincount(arr, output_size, weights)
-@tf_export("cumsum")
+@tf_export("math.cumsum", "cumsum")
def cumsum(x, axis=0, exclusive=False, reverse=False, name=None):
"""Compute the cumulative sum of the tensor `x` along `axis`.
@@ -2414,7 +2424,8 @@
x, axis, exclusive=exclusive, reverse=reverse, name=name)
-@tf_export("cumprod")
+@tf_export("math.cumprod", "cumprod")
+@deprecation.deprecated_endpoints("cumprod")
def cumprod(x, axis=0, exclusive=False, reverse=False, name=None):
"""Compute the cumulative product of the tensor `x` along `axis`.
@@ -2422,7 +2433,7 @@
first element of the input is identical to the first element of the output:
```python
- tf.cumprod([a, b, c]) # [a, a * b, a * b * c]
+ tf.math.cumprod([a, b, c]) # [a, a * b, a * b * c]
```
By setting the `exclusive` kwarg to `True`, an exclusive cumprod is
@@ -2430,21 +2441,21 @@
instead:
```python
- tf.cumprod([a, b, c], exclusive=True) # [1, a, a * b]
+ tf.math.cumprod([a, b, c], exclusive=True) # [1, a, a * b]
```
By setting the `reverse` kwarg to `True`, the cumprod is performed in the
opposite direction:
```python
- tf.cumprod([a, b, c], reverse=True) # [a * b * c, b * c, c]
+ tf.math.cumprod([a, b, c], reverse=True) # [a * b * c, b * c, c]
```
This is more efficient than using separate `tf.reverse` ops.
The `reverse` and `exclusive` kwargs can also be combined:
```python
- tf.cumprod([a, b, c], exclusive=True, reverse=True) # [b * c, c, 1]
+ tf.math.cumprod([a, b, c], exclusive=True, reverse=True) # [b * c, c, 1]
```
Args:
@@ -2466,7 +2477,8 @@
x, axis, exclusive=exclusive, reverse=reverse, name=name)
-@tf_export("conj")
+@tf_export("math.conj", "conj")
+@deprecation.deprecated_endpoints("conj")
def conj(x, name=None):
r"""Returns the complex conjugate of a complex number.
@@ -2480,7 +2492,7 @@
For example:
# tensor 'input' is [-2.25 + 4.75j, 3.25 + 5.75j]
- tf.conj(input) ==> [-2.25 - 4.75j, 3.25 - 5.75j]
+ tf.math.conj(input) ==> [-2.25 - 4.75j, 3.25 - 5.75j]
If `x` is real, it is returned unchanged.
@@ -2566,7 +2578,8 @@
return gen_math_ops.maximum(N, 1)
-@tf_export("unsorted_segment_mean")
+@tf_export("math.unsorted_segment_mean", "unsorted_segment_mean")
+@deprecation.deprecated_endpoints("unsorted_segment_mean")
def unsorted_segment_mean(data, segment_ids, num_segments, name=None):
r"""Computes the mean along segments of a tensor.
@@ -2608,7 +2621,8 @@
return summed / N
-@tf_export("unsorted_segment_sqrt_n")
+@tf_export("math.unsorted_segment_sqrt_n", "unsorted_segment_sqrt_n")
+@deprecation.deprecated_endpoints("unsorted_segment_sqrt_n")
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).
@@ -2653,7 +2667,8 @@
return summed / gen_math_ops.sqrt(N)
-@tf_export("sparse_segment_sum")
+@tf_export("sparse.segment_sum", "sparse_segment_sum")
+@deprecation.deprecated_endpoints("sparse_segment_sum")
def sparse_segment_sum(data, indices, segment_ids, name=None,
num_segments=None):
r"""Computes the sum along sparse segments of a tensor.
@@ -2674,16 +2689,16 @@
c = tf.constant([[1,2,3,4], [-1,-2,-3,-4], [5,6,7,8]])
# Select two rows, one segment.
- tf.sparse_segment_sum(c, tf.constant([0, 1]), tf.constant([0, 0]))
+ tf.sparse.segment_sum(c, tf.constant([0, 1]), tf.constant([0, 0]))
# => [[0 0 0 0]]
# Select two rows, two segment.
- tf.sparse_segment_sum(c, tf.constant([0, 1]), tf.constant([0, 1]))
+ tf.sparse.segment_sum(c, tf.constant([0, 1]), tf.constant([0, 1]))
# => [[ 1 2 3 4]
# [-1 -2 -3 -4]]
# With missing segment ids.
- tf.sparse_segment_sum(c, tf.constant([0, 1]), tf.constant([0, 2]),
+ tf.sparse.segment_sum(c, tf.constant([0, 1]), tf.constant([0, 2]),
num_segments=4)
# => [[ 1 2 3 4]
# [ 0 0 0 0]
@@ -2691,7 +2706,7 @@
# [ 0 0 0 0]]
# Select all rows, two segments.
- tf.sparse_segment_sum(c, tf.constant([0, 1, 2]), tf.constant([0, 0, 1]))
+ tf.sparse.segment_sum(c, tf.constant([0, 1, 2]), tf.constant([0, 0, 1]))
# => [[0 0 0 0]
# [5 6 7 8]]
@@ -2726,7 +2741,8 @@
data=data, indices=indices, segment_ids=segment_ids, name=name)
-@tf_export("sparse_segment_mean")
+@tf_export("sparse.segment_mean", "sparse_segment_mean")
+@deprecation.deprecated_endpoints("sparse_segment_mean")
def sparse_segment_mean(data,
indices,
segment_ids,
@@ -2771,7 +2787,8 @@
data=data, indices=indices, segment_ids=segment_ids, name=name)
-@tf_export("sparse_segment_sqrt_n")
+@tf_export("sparse.segment_sqrt_n", "sparse_segment_sqrt_n")
+@deprecation.deprecated_endpoints("sparse_segment_sqrt_n")
def sparse_segment_sqrt_n(data,
indices,
segment_ids,
diff --git a/tensorflow/python/ops/math_ops_test.py b/tensorflow/python/ops/math_ops_test.py
index 1b01d1d..f051850 100644
--- a/tensorflow/python/ops/math_ops_test.py
+++ b/tensorflow/python/ops/math_ops_test.py
@@ -21,6 +21,7 @@
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
@@ -488,5 +489,75 @@
self.assertAllEqual(tf_result, np_result)
+class XlogyTest(test_util.TensorFlowTestCase):
+
+ @test_util.run_in_graph_and_eager_modes
+ def testXlogyNoZero(self):
+ for dtype in [dtypes.float16, dtypes.float32, dtypes.float64]:
+ x = constant_op.constant([[0.1, 0.2, 3.5], [-2., -5., 30.]], dtype=dtype)
+ y = constant_op.constant([[0.1, 0.2, 3.5], [3.1, 4., 2.]], dtype=dtype)
+ with self.cached_session(use_gpu=True):
+ xlogy = self.evaluate(math_ops.xlogy(x, y))
+ xtimeslogy = self.evaluate(x * math_ops.log(y))
+ self.assertAllClose(xlogy, xtimeslogy)
+
+ @test_util.run_in_graph_and_eager_modes
+ def testXlogyWithZero(self):
+ for dtype in [dtypes.float16, dtypes.float32, dtypes.float64]:
+ x = constant_op.constant(np.zeros((2, 3)), dtype=dtype)
+ y = constant_op.constant([[0.1, 0.2, 3.5], [0., 1., 2.]], dtype=dtype)
+ with self.cached_session(use_gpu=True):
+ xlogy_tf_np = self.evaluate(math_ops.xlogy(x, y))
+ zeros_np = self.evaluate(array_ops.zeros_like(y))
+ self.assertAllClose(xlogy_tf_np, zeros_np)
+
+ @test_util.run_in_graph_and_eager_modes
+ def testXlogyWithZeroBroadcast(self):
+ for dtype in [dtypes.float16, dtypes.float32, dtypes.float64]:
+ x = constant_op.constant([[0.], [1.]], dtype=dtype)
+ y = constant_op.constant([[0.1, 0.2, 3.5], [0., 1., 2.]], dtype=dtype)
+ with self.cached_session(use_gpu=True):
+ xlogy_tf_np = self.evaluate(math_ops.xlogy(x, y))
+ zeros_np = self.evaluate(array_ops.zeros_like(y[0]))
+ xtimes_logy = self.evaluate(math_ops.log(y[1]))
+ self.assertAllClose(zeros_np, xlogy_tf_np[0])
+ self.assertAllClose(xtimes_logy, xlogy_tf_np[1])
+
+
+class XdivyTest(test_util.TensorFlowTestCase):
+
+ @test_util.run_in_graph_and_eager_modes
+ def testXdivyNoZero(self):
+ for dtype in [dtypes.float16, dtypes.float32, dtypes.float64]:
+ x = constant_op.constant([[0.1, 0.2, 3.5], [-2., -5., 30.]], dtype=dtype)
+ y = constant_op.constant([[0.1, 0.2, 3.5], [3.1, 4., 2.]], dtype=dtype)
+ with self.cached_session(use_gpu=True):
+ xdivy = self.evaluate(math_ops.xdivy(x, y))
+ x_over_y = self.evaluate(x / y)
+ self.assertAllClose(xdivy, x_over_y)
+
+ @test_util.run_in_graph_and_eager_modes
+ def testXdivyWithZero(self):
+ for dtype in [dtypes.float16, dtypes.float32, dtypes.float64]:
+ x = constant_op.constant(np.zeros((2, 3)), dtype=dtype)
+ y = constant_op.constant([[0.1, 0.2, 3.5], [0., 1., 2.]], dtype=dtype)
+ with self.cached_session(use_gpu=True):
+ xdivy_tf_np = self.evaluate(math_ops.xdivy(x, y))
+ zeros_np = self.evaluate(array_ops.zeros_like(y))
+ self.assertAllClose(xdivy_tf_np, zeros_np)
+
+ @test_util.run_in_graph_and_eager_modes
+ def testXdivyWithZeroBroadcast(self):
+ for dtype in [dtypes.float16, dtypes.float32, dtypes.float64]:
+ x = constant_op.constant([[0.], [1.]], dtype=dtype)
+ y = constant_op.constant([[0.1, 0.2, 3.5], [0., 1., 2.]], dtype=dtype)
+ with self.cached_session(use_gpu=True):
+ xdivy_tf_np = self.evaluate(math_ops.xdivy(x, y))
+ zeros_np = self.evaluate(array_ops.zeros_like(y[0]))
+ x_over_y = self.evaluate(1 / y[1])
+ self.assertAllClose(zeros_np, xdivy_tf_np[0])
+ self.assertAllClose(x_over_y, xdivy_tf_np[1])
+
+
if __name__ == "__main__":
googletest.main()
diff --git a/tensorflow/python/ops/matmul_benchmark.py b/tensorflow/python/ops/matmul_benchmark.py
index 6e5fe74..138149e 100644
--- a/tensorflow/python/ops/matmul_benchmark.py
+++ b/tensorflow/python/ops/matmul_benchmark.py
@@ -49,13 +49,13 @@
"""
with ops.device('%s' % device):
if not transpose_a:
- x = variables.Variable(random_ops.random_uniform([n, m], dtype=dtype))
+ x = variables.VariableV1(random_ops.random_uniform([n, m], dtype=dtype))
else:
- x = variables.Variable(random_ops.random_uniform([m, n], dtype=dtype))
+ x = variables.VariableV1(random_ops.random_uniform([m, n], dtype=dtype))
if not transpose_b:
- y = variables.Variable(random_ops.random_uniform([m, k], dtype=dtype))
+ y = variables.VariableV1(random_ops.random_uniform([m, k], dtype=dtype))
else:
- y = variables.Variable(random_ops.random_uniform([k, m], dtype=dtype))
+ y = variables.VariableV1(random_ops.random_uniform([k, m], dtype=dtype))
z = math_ops.matmul(x, y, transpose_a=transpose_a, transpose_b=transpose_b)
return control_flow_ops.group(z)
diff --git a/tensorflow/python/ops/nn_impl.py b/tensorflow/python/ops/nn_impl.py
index 2a1919e..453848f 100644
--- a/tensorflow/python/ops/nn_impl.py
+++ b/tensorflow/python/ops/nn_impl.py
@@ -328,7 +328,7 @@
return features * math_ops.sigmoid(features)
-@tf_export("nn.l2_normalize")
+@tf_export("math.l2_normalize", "linalg.l2_normalize", "nn.l2_normalize")
@deprecated_args(None, "dim is deprecated, use axis instead", "dim")
def l2_normalize(x, axis=None, epsilon=1e-12, name=None, dim=None):
"""Normalizes along dimension `axis` using an L2 norm.
@@ -360,7 +360,7 @@
return math_ops.multiply(x, x_inv_norm, name=name)
-@tf_export("nn.zero_fraction")
+@tf_export("math.zero_fraction", "nn.zero_fraction")
def zero_fraction(value, name=None):
"""Returns the fraction of zeros in `value`.
@@ -689,7 +689,7 @@
# Compute true mean while keeping the dims for proper broadcasting.
mean = math_ops.reduce_mean(y, axes, keepdims=True, name="mean")
# sample variance, not unbiased variance
- # Note: stop_gradient does not change the gradient that gets
+ # Note: stop_gradient does not change the gradient that gets
# backpropagated to the mean from the variance calculation,
# because that gradient is zero
variance = math_ops.reduce_mean(
diff --git a/tensorflow/python/ops/nn_ops.py b/tensorflow/python/ops/nn_ops.py
index 9ef177e..fd71e7c 100644
--- a/tensorflow/python/ops/nn_ops.py
+++ b/tensorflow/python/ops/nn_ops.py
@@ -1692,7 +1692,7 @@
return output
-@tf_export("nn.softmax")
+@tf_export("nn.softmax", "math.softmax")
@deprecation.deprecated_args(None, "dim is deprecated, use axis instead", "dim")
def softmax(logits, axis=None, name=None, dim=None):
"""Computes softmax activations.
@@ -1722,7 +1722,7 @@
return _softmax(logits, gen_nn_ops.softmax, axis, name)
-@tf_export("nn.log_softmax")
+@tf_export("nn.log_softmax", "math.log_softmax")
@deprecation.deprecated_args(None, "dim is deprecated, use axis instead", "dim")
def log_softmax(logits, axis=None, name=None, dim=None):
"""Computes log softmax activations.
@@ -2329,7 +2329,7 @@
return ret
-@tf_export("nn.top_k")
+@tf_export("math.top_k", "nn.top_k")
def top_k(input, k=1, sorted=True, name=None): # pylint: disable=redefined-builtin
"""Finds values and indices of the `k` largest entries for the last dimension.
@@ -2644,7 +2644,7 @@
name=name))
-@tf_export("nn.in_top_k")
+@tf_export("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.
diff --git a/tensorflow/python/ops/numerics.py b/tensorflow/python/ops/numerics.py
index 8fcbd7d..002e87b 100644
--- a/tensorflow/python/ops/numerics.py
+++ b/tensorflow/python/ops/numerics.py
@@ -24,10 +24,12 @@
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
-@tf_export("verify_tensor_all_finite")
+@tf_export("debugging.assert_all_finite", "verify_tensor_all_finite")
+@deprecation.deprecated_endpoints("verify_tensor_all_finite")
def verify_tensor_all_finite(t, msg, name=None):
"""Assert that the tensor does not contain any NaN's or Inf's.
diff --git a/tensorflow/python/ops/parallel_for/pfor.py b/tensorflow/python/ops/parallel_for/pfor.py
index e0f6d51..83cbe64 100644
--- a/tensorflow/python/ops/parallel_for/pfor.py
+++ b/tensorflow/python/ops/parallel_for/pfor.py
@@ -1987,14 +1987,12 @@
@RegisterPForWithArgs("Pow", math_ops.pow)
@RegisterPForWithArgs("RealDiv", math_ops.divide)
@RegisterPForWithArgs("Real", math_ops.real)
-@RegisterPForWithArgs("ReciprocalGrad", math_ops.reciprocal_grad)
@RegisterPForWithArgs("Reciprocal", math_ops.reciprocal)
@RegisterPForWithArgs("Relu6", nn_ops.relu6)
@RegisterPForWithArgs("Relu", nn_ops.relu)
@RegisterPForWithArgs("RightShift", bitwise_ops.right_shift)
@RegisterPForWithArgs("Rint", math_ops.rint)
@RegisterPForWithArgs("Round", math_ops.round)
-@RegisterPForWithArgs("RsqrtGrad", math_ops.rsqrt_grad)
@RegisterPForWithArgs("Rsqrt", math_ops.rsqrt)
@RegisterPForWithArgs("Selu", nn_ops.selu)
@RegisterPForWithArgs("Sigmoid", math_ops.sigmoid)
@@ -2003,7 +2001,6 @@
@RegisterPForWithArgs("Sin", math_ops.sin)
@RegisterPForWithArgs("Softplus", nn_ops.softplus)
@RegisterPForWithArgs("Softsign", nn_ops.softsign)
-@RegisterPForWithArgs("SqrtGrad", math_ops.sqrt_grad)
@RegisterPForWithArgs("Sqrt", math_ops.sqrt)
@RegisterPForWithArgs("SquaredDifference", math_ops.squared_difference)
@RegisterPForWithArgs("Square", math_ops.square)
@@ -2095,6 +2092,9 @@
@RegisterPForWithArgs("SoftplusGrad")
@RegisterPForWithArgs("SoftsignGrad")
@RegisterPForWithArgs("TanhGrad")
+@RegisterPForWithArgs("SqrtGrad")
+@RegisterPForWithArgs("RsqrtGrad")
+@RegisterPForWithArgs("ReciprocalGrad")
def _convert_grads(pfor_input, op_type, *args, **kw_args):
del args
del kw_args
diff --git a/tensorflow/python/ops/parsing_ops.py b/tensorflow/python/ops/parsing_ops.py
index b3e03a0..ff50fe0 100644
--- a/tensorflow/python/ops/parsing_ops.py
+++ b/tensorflow/python/ops/parsing_ops.py
@@ -36,6 +36,7 @@
from tensorflow.python.ops.gen_parsing_ops import *
# pylint: enable=wildcard-import,undefined-variable
from tensorflow.python.platform import tf_logging
+from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
@@ -45,7 +46,7 @@
ops.NotDifferentiable("StringToNumber")
-@tf_export("VarLenFeature")
+@tf_export("io.VarLenFeature", "VarLenFeature")
class VarLenFeature(collections.namedtuple("VarLenFeature", ["dtype"])):
"""Configuration for parsing a variable-length input feature.
@@ -55,7 +56,7 @@
pass
-@tf_export("SparseFeature")
+@tf_export("io.SparseFeature", "SparseFeature")
class SparseFeature(
collections.namedtuple(
"SparseFeature",
@@ -130,7 +131,7 @@
cls, index_key, value_key, dtype, size, already_sorted)
-@tf_export("FixedLenFeature")
+@tf_export("io.FixedLenFeature", "FixedLenFeature")
class FixedLenFeature(collections.namedtuple(
"FixedLenFeature", ["shape", "dtype", "default_value"])):
"""Configuration for parsing a fixed-length input feature.
@@ -150,7 +151,7 @@
cls, shape, dtype, default_value)
-@tf_export("FixedLenSequenceFeature")
+@tf_export("io.FixedLenSequenceFeature", "FixedLenSequenceFeature")
class FixedLenSequenceFeature(collections.namedtuple(
"FixedLenSequenceFeature",
["shape", "dtype", "allow_missing", "default_value"])):
@@ -360,7 +361,7 @@
return features
-@tf_export("parse_example")
+@tf_export("io.parse_example", "parse_example")
def parse_example(serialized, features, name=None, example_names=None):
# pylint: disable=line-too-long
"""Parses `Example` protos into a `dict` of tensors.
@@ -761,7 +762,7 @@
dense_shapes_as_proto, dense_shapes)
-@tf_export("parse_single_example")
+@tf_export("io.parse_single_example", "parse_single_example")
def parse_single_example(serialized, features, name=None, example_names=None):
"""Parses a single `Example` proto.
@@ -1244,7 +1245,7 @@
# TODO(sundberg): rewrite this method to call the batch version, which is more
# efficient especially for large inputs.
-@tf_export("parse_single_sequence_example")
+@tf_export("io.parse_single_sequence_example", "parse_single_sequence_example")
def parse_single_sequence_example(
serialized, context_features=None, sequence_features=None,
example_name=None, name=None):
@@ -1564,7 +1565,8 @@
# Swap `name` and `na_value` for backward compatibility.
-@tf_export("decode_csv")
+@tf_export("io.decode_csv", "decode_csv")
+@deprecation.deprecated_endpoints("decode_csv")
def decode_csv(records,
record_defaults,
field_delim=",",
diff --git a/tensorflow/python/ops/random_ops.py b/tensorflow/python/ops/random_ops.py
index 4baf506..c2eb9df 100644
--- a/tensorflow/python/ops/random_ops.py
+++ b/tensorflow/python/ops/random_ops.py
@@ -29,6 +29,7 @@
# go/tf-wildcard-import
# pylint: disable=wildcard-import
from tensorflow.python.ops.gen_random_ops import *
+from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
# pylint: enable=wildcard-import
@@ -43,7 +44,7 @@
return ops.convert_to_tensor(shape, dtype=dtype, name="shape")
-@tf_export("random_normal")
+@tf_export("random.normal", "random_normal")
def random_normal(shape,
mean=0.0,
stddev=1.0,
@@ -136,7 +137,7 @@
return rnd
-@tf_export("truncated_normal")
+@tf_export("random.truncated_normal", "truncated_normal")
def truncated_normal(shape,
mean=0.0,
stddev=1.0,
@@ -181,7 +182,7 @@
ops.NotDifferentiable("TruncatedNormal")
-@tf_export("random_uniform")
+@tf_export("random.uniform", "random_uniform")
def random_uniform(shape,
minval=0,
maxval=None,
@@ -246,7 +247,7 @@
ops.NotDifferentiable("RandomUniform")
-@tf_export("random_shuffle")
+@tf_export("random.shuffle", "random_shuffle")
def random_shuffle(value, seed=None, name=None):
"""Randomly shuffles a tensor along its first dimension.
@@ -277,7 +278,7 @@
value, seed=seed1, seed2=seed2, name=name)
-@tf_export("random_crop")
+@tf_export("image.random_crop", "random_crop")
def random_crop(value, size, seed=None, name=None):
"""Randomly crops a tensor to a given size.
@@ -320,7 +321,7 @@
return array_ops.slice(value, offset, size, name=name)
-@tf_export("multinomial")
+@tf_export("random.multinomial", "multinomial")
def multinomial(logits, num_samples, seed=None, name=None, output_dtype=None):
"""Draws samples from a multinomial distribution.
@@ -356,7 +357,8 @@
ops.NotDifferentiable("Multinomial")
-@tf_export("random_gamma")
+@tf_export("random.gamma", "random_gamma")
+@deprecation.deprecated_endpoints("random_gamma")
def random_gamma(shape,
alpha,
beta=None,
@@ -439,7 +441,8 @@
shape, alpha_broadcast, seed=seed1, seed2=seed2) / beta)
-@tf_export("random_poisson")
+@tf_export("random.poisson", "random_poisson")
+@deprecation.deprecated_endpoints("random_poisson")
def random_poisson(lam, shape, dtype=dtypes.float32, seed=None, name=None):
"""Draws `shape` samples from each of the given Poisson distribution(s).
diff --git a/tensorflow/python/ops/rnn_cell_impl.py b/tensorflow/python/ops/rnn_cell_impl.py
index 43cca1a..dd4f3d7 100644
--- a/tensorflow/python/ops/rnn_cell_impl.py
+++ b/tensorflow/python/ops/rnn_cell_impl.py
@@ -370,7 +370,7 @@
*args, **kwargs)
-@tf_export("nn.rnn_cell.BasicRNNCell")
+@tf_export(v1=["nn.rnn_cell.BasicRNNCell"])
class BasicRNNCell(LayerRNNCell):
"""The most basic RNN cell.
@@ -393,6 +393,8 @@
`trainable` etc when constructing the cell from configs of get_config().
"""
+ @deprecated(None, "This class is equivalent as tf.keras.layers.SimpleRNNCell,"
+ " and will be replaced by that in Tensorflow 2.0.")
def __init__(self,
num_units,
activation=None,
@@ -611,7 +613,7 @@
# TODO(scottzhu): Stop exporting this class in TF 2.0.
@tf_export("nn.rnn_cell.BasicLSTMCell")
class BasicLSTMCell(LayerRNNCell):
- """DEPRECATED: Please use @{tf.nn.rnn_cell.LSTMCell} instead.
+ """DEPRECATED: Please use `tf.nn.rnn_cell.LSTMCell` instead.
Basic LSTM recurrent network cell.
diff --git a/tensorflow/python/ops/sparse_ops.py b/tensorflow/python/ops/sparse_ops.py
index 400a42a..7e3dbdb 100644
--- a/tensorflow/python/ops/sparse_ops.py
+++ b/tensorflow/python/ops/sparse_ops.py
@@ -185,7 +185,8 @@
# pylint: disable=protected-access
-@tf_export("sparse_concat")
+@tf_export("sparse.concat", "sparse_concat")
+@deprecation.deprecated_endpoints("sparse_concat")
@deprecation.deprecated_args(
None, "concat_dim is deprecated, use axis instead", "concat_dim")
def sparse_concat(axis,
@@ -317,7 +318,8 @@
return sparse_tensor.SparseTensor(output_ind, output_val, output_shape)
-@tf_export("sparse_add")
+@tf_export("sparse.add", "sparse_add")
+@deprecation.deprecated_endpoints("sparse_add")
def sparse_add(a, b, thresh=0):
"""Adds two tensors, at least one of each is a `SparseTensor`.
@@ -557,7 +559,8 @@
return sparse_tensor.SparseTensor(sp_t.indices, result, sp_t.dense_shape)
-@tf_export("sparse_reorder")
+@tf_export("sparse.reorder", "sparse_reorder")
+@deprecation.deprecated_endpoints("sparse_reorder")
def sparse_reorder(sp_input, name=None):
"""Reorders a `SparseTensor` into the canonical, row-major ordering.
@@ -607,7 +610,8 @@
return sparse_tensor.SparseTensor(reordered_ind, reordered_val, dense_shape)
-@tf_export("sparse_reshape")
+@tf_export("sparse.reshape", "sparse_reshape")
+@deprecation.deprecated_endpoints("sparse_reshape")
def sparse_reshape(sp_input, shape, name=None):
"""Reshapes a `SparseTensor` to represent values in a new dense shape.
@@ -700,7 +704,8 @@
return "KeywordRequired()"
-@tf_export("sparse_split")
+@tf_export("sparse.split", "sparse_split")
+@deprecation.deprecated_endpoints("sparse_split")
@deprecation.deprecated_args(
None, "split_dim is deprecated, use axis instead", "split_dim")
def sparse_split(keyword_required=KeywordRequired(),
@@ -773,7 +778,8 @@
return sparse_tensors
-@tf_export("sparse_slice")
+@tf_export("sparse.slice", "sparse_slice")
+@deprecation.deprecated_endpoints("sparse_slice")
def sparse_slice(sp_input, start, size, name=None):
"""Slice a `SparseTensor` based on the `start` and `size.
@@ -785,11 +791,11 @@
Graphically the output tensors are:
- sparse_slice([0, 0], [2, 4]) = shape = [2, 4]
+ sparse.slice([0, 0], [2, 4]) = shape = [2, 4]
[ a ]
[b c ]
- sparse_slice([0, 4], [2, 3]) = shape = [2, 3]
+ sparse.slice([0, 4], [2, 3]) = shape = [2, 3]
[ d e ]
[ ]
@@ -823,6 +829,9 @@
@tf_export("sparse_to_dense")
+@deprecation.deprecated(
+ None,
+ "Create a `tf.sparse.SparseTensor` and use `tf.sparse.to_dense` instead.")
def sparse_to_dense(sparse_indices,
output_shape,
sparse_values,
@@ -878,7 +887,8 @@
name=name)
-@tf_export("sparse_reduce_max")
+@tf_export("sparse.reduce_max", "sparse_reduce_max")
+@deprecation.deprecated_endpoints("sparse_reduce_max")
@deprecation.deprecated_args(
None, "keep_dims is deprecated, use keepdims instead", "keep_dims")
def sparse_reduce_max(sp_input, axis=None, keepdims=None,
@@ -912,16 +922,16 @@
# 'x' represents [[1, ?, 2]
# [?, 3, ?]]
# where ? is implicitly-zero.
- tf.sparse_reduce_max(x) ==> 3
- tf.sparse_reduce_max(x, 0) ==> [1, 3, 2]
- tf.sparse_reduce_max(x, 1) ==> [2, 3] # Can also use -1 as the axis.
- tf.sparse_reduce_max(x, 1, keepdims=True) ==> [[2], [3]]
- tf.sparse_reduce_max(x, [0, 1]) ==> 3
+ tf.sparse.reduce_max(x) ==> 3
+ tf.sparse.reduce_max(x, 0) ==> [1, 3, 2]
+ tf.sparse.reduce_max(x, 1) ==> [2, 3] # Can also use -1 as the axis.
+ tf.sparse.reduce_max(x, 1, keepdims=True) ==> [[2], [3]]
+ tf.sparse.reduce_max(x, [0, 1]) ==> 3
# 'y' represents [[-7, ?]
# [ 4, 3]
# [ ?, ?]
- tf.sparse_reduce_max(x, 1) ==> [-7, 4, 0]
+ tf.sparse.reduce_max(x, 1) ==> [-7, 4, 0]
```
Args:
@@ -945,7 +955,8 @@
math_ops._ReductionDims(sp_input, axis, reduction_axes), keepdims)
-@tf_export("sparse_reduce_max_sparse")
+@tf_export("sparse.reduce_max_sparse", "sparse_reduce_max_sparse")
+@deprecation.deprecated_endpoints("sparse_reduce_max_sparse")
@deprecation.deprecated_args(
None, "keep_dims is deprecated, use keepdims instead", "keep_dims")
def sparse_reduce_max_sparse(sp_input,
@@ -995,7 +1006,8 @@
return sparse_tensor.SparseTensor(output_ind, output_val, output_shape)
-@tf_export("sparse_reduce_sum")
+@tf_export("sparse.reduce_sum", "sparse_reduce_sum")
+@deprecation.deprecated_endpoints("sparse_reduce_sum")
@deprecation.deprecated_args(
None, "keep_dims is deprecated, use keepdims instead", "keep_dims")
def sparse_reduce_sum(sp_input, axis=None, keepdims=None,
@@ -1021,11 +1033,11 @@
# 'x' represents [[1, ?, 1]
# [?, 1, ?]]
# where ? is implicitly-zero.
- tf.sparse_reduce_sum(x) ==> 3
- tf.sparse_reduce_sum(x, 0) ==> [1, 1, 1]
- tf.sparse_reduce_sum(x, 1) ==> [2, 1] # Can also use -1 as the axis.
- tf.sparse_reduce_sum(x, 1, keepdims=True) ==> [[2], [1]]
- tf.sparse_reduce_sum(x, [0, 1]) ==> 3
+ tf.sparse.reduce_sum(x) ==> 3
+ tf.sparse.reduce_sum(x, 0) ==> [1, 1, 1]
+ tf.sparse.reduce_sum(x, 1) ==> [2, 1] # Can also use -1 as the axis.
+ tf.sparse.reduce_sum(x, 1, keepdims=True) ==> [[2], [1]]
+ tf.sparse.reduce_sum(x, [0, 1]) ==> 3
```
Args:
@@ -1049,7 +1061,8 @@
math_ops._ReductionDims(sp_input, axis, reduction_axes), keepdims)
-@tf_export("sparse_reduce_sum_sparse")
+@tf_export("sparse.reduce_sum_sparse", "sparse_reduce_sum_sparse")
+@deprecation.deprecated_endpoints("sparse_reduce_sum_sparse")
@deprecation.deprecated_args(
None, "keep_dims is deprecated, use keepdims instead", "keep_dims")
def sparse_reduce_sum_sparse(sp_input,
@@ -1099,7 +1112,8 @@
return sparse_tensor.SparseTensor(output_ind, output_val, output_shape)
-@tf_export("sparse_tensor_to_dense")
+@tf_export("sparse.to_dense", "sparse_tensor_to_dense")
+@deprecation.deprecated_endpoints("sparse_tensor_to_dense")
def sparse_tensor_to_dense(sp_input,
default_value=0,
validate_indices=True,
@@ -1151,7 +1165,8 @@
name=name)
-@tf_export("sparse_to_indicator")
+@tf_export("sparse.to_indicator", "sparse_to_indicator")
+@deprecation.deprecated_endpoints("sparse_to_indicator")
def sparse_to_indicator(sp_input, vocab_size, name=None):
"""Converts a `SparseTensor` of ids into a dense bool indicator tensor.
@@ -1214,7 +1229,8 @@
sp_new, default_value=False, validate_indices=False, name=name)
-@tf_export("sparse_merge")
+@tf_export("sparse.merge", "sparse_merge")
+@deprecation.deprecated_endpoints("sparse_merge")
def sparse_merge(sp_ids, sp_values, vocab_size, name=None,
already_sorted=False):
"""Combines a batch of feature ids and values into a single `SparseTensor`.
@@ -1358,7 +1374,8 @@
sorted_result.indices, sorted_result.values, new_shape)
-@tf_export("sparse_retain")
+@tf_export("sparse.retain", "sparse_retain")
+@deprecation.deprecated_endpoints("sparse_retain")
def sparse_retain(sp_input, to_retain):
"""Retains specified non-empty values within a `SparseTensor`.
@@ -1402,7 +1419,8 @@
array_ops.identity(sp_input.dense_shape))
-@tf_export("sparse_reset_shape")
+@tf_export("sparse.reset_shape", "sparse_reset_shape")
+@deprecation.deprecated_endpoints("sparse_reset_shape")
def sparse_reset_shape(sp_input, new_shape=None):
"""Resets the shape of a `SparseTensor` with indices and values unchanged.
@@ -1503,7 +1521,8 @@
return sparse_tensor.SparseTensor(in_indices, in_values, output_shape_tensor)
-@tf_export("sparse_fill_empty_rows")
+@tf_export("sparse.fill_empty_rows", "sparse_fill_empty_rows")
+@deprecation.deprecated_endpoints("sparse_fill_empty_rows")
def sparse_fill_empty_rows(sp_input, default_value, name=None):
"""Fills empty rows in the input 2-D `SparseTensor` with a default value.
@@ -1567,7 +1586,8 @@
dense_shape=sp_input.dense_shape), empty_row_indicator)
-@tf_export("serialize_sparse")
+@tf_export("io.serialize_sparse", "serialize_sparse")
+@deprecation.deprecated_endpoints("serialize_sparse")
def serialize_sparse(sp_input, name=None, out_type=dtypes.string):
"""Serialize a `SparseTensor` into a 3-vector (1-D `Tensor`) object.
@@ -1593,7 +1613,8 @@
out_type=out_type)
-@tf_export("serialize_many_sparse")
+@tf_export("io.serialize_many_sparse", "serialize_many_sparse")
+@deprecation.deprecated_endpoints("serialize_many_sparse")
def serialize_many_sparse(sp_input, name=None, out_type=dtypes.string):
"""Serialize `N`-minibatch `SparseTensor` into an `[N, 3]` `Tensor`.
@@ -1694,7 +1715,8 @@
return sparse_tensor.SparseTensor(output_indices, output_values, output_shape)
-@tf_export("deserialize_many_sparse")
+@tf_export("io.deserialize_many_sparse", "deserialize_many_sparse")
+@deprecation.deprecated_endpoints("deserialize_many_sparse")
def deserialize_many_sparse(serialized_sparse, dtype, rank=None, name=None):
"""Deserialize and concatenate `SparseTensors` from a serialized minibatch.
@@ -1712,7 +1734,7 @@
The input `SparseTensor` objects' indices are assumed ordered in
standard lexicographic order. If this is not the case, after this
- step run `sparse_reorder` to restore index ordering.
+ step run `sparse.reorder` to restore index ordering.
For example, if the serialized input is a `[2, 3]` matrix representing two
original `SparseTensor` objects:
@@ -1764,7 +1786,8 @@
return sparse_tensor.SparseTensor(output_indices, output_values, output_shape)
-@tf_export("sparse_tensor_dense_matmul")
+@tf_export("sparse.matmul", "sparse_tensor_dense_matmul")
+@deprecation.deprecated_endpoints("sparse_tensor_dense_matmul")
def sparse_tensor_dense_matmul(sp_a,
b,
adjoint_a=False,
@@ -1777,7 +1800,7 @@
following input format is recommended for optimal behavior:
* If `adjoint_a == false`: `A` should be sorted in lexicographically
- increasing order. Use `sparse_reorder` if you're not sure.
+ increasing order. Use `sparse.reorder` if you're not sure.
* If `adjoint_a == true`: `A` should be sorted in order of increasing
dimension 1 (i.e., "column major" order instead of "row major" order).
@@ -1981,7 +2004,8 @@
adjoint_b=adjoint_b)
-@tf_export("sparse_softmax")
+@tf_export("sparse.softmax", "sparse_softmax")
+@deprecation.deprecated_endpoints("sparse_softmax")
def sparse_softmax(sp_input, name=None):
"""Applies softmax to a batched N-D `SparseTensor`.
@@ -2036,7 +2060,8 @@
sp_input.dense_shape)
-@tf_export("sparse_maximum")
+@tf_export("sparse.maximum", "sparse_maximum")
+@deprecation.deprecated_endpoints("sparse_maximum")
def sparse_maximum(sp_a, sp_b, name=None):
"""Returns the element-wise max of two SparseTensors.
@@ -2073,7 +2098,8 @@
return sparse_tensor.SparseTensor(out_indices, out_values, sp_a.dense_shape)
-@tf_export("sparse_minimum")
+@tf_export("sparse.minimum", "sparse_minimum")
+@deprecation.deprecated_endpoints("sparse_minimum")
def sparse_minimum(sp_a, sp_b, name=None):
"""Returns the element-wise min of two SparseTensors.
@@ -2110,7 +2136,8 @@
return sparse_tensor.SparseTensor(out_indices, out_values, sp_a.dense_shape)
-@tf_export("sparse_transpose")
+@tf_export("sparse.transpose", "sparse_transpose")
+@deprecation.deprecated_endpoints("sparse_transpose")
def sparse_transpose(sp_input, perm=None, name=None):
"""Transposes a `SparseTensor`
@@ -2259,7 +2286,7 @@
The input `SparseTensor` objects' indices are assumed ordered in
standard lexicographic order. If this is not the case, after this
- step run `sparse_reorder` to restore index ordering.
+ step run `sparse.reorder` to restore index ordering.
For example, if the serialized input is a `[2, 3]` matrix representing two
original `SparseTensor` objects:
diff --git a/tensorflow/python/ops/special_math_ops.py b/tensorflow/python/ops/special_math_ops.py
index 9a10abf..cfab943 100644
--- a/tensorflow/python/ops/special_math_ops.py
+++ b/tensorflow/python/ops/special_math_ops.py
@@ -29,11 +29,13 @@
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
# TODO(b/27419586) Change docstring for required dtype of x once int allowed
-@tf_export('lbeta')
+@tf_export('math.lbeta', 'lbeta')
+@deprecation.deprecated_endpoints('lbeta')
def lbeta(x, name=None):
r"""Computes \\(ln(|Beta(x)|)\\), reducing along the last dimension.
diff --git a/tensorflow/python/ops/string_ops.py b/tensorflow/python/ops/string_ops.py
index 046a48d..e83c08f 100644
--- a/tensorflow/python/ops/string_ops.py
+++ b/tensorflow/python/ops/string_ops.py
@@ -310,8 +310,9 @@
return math_ops.range(array_ops.rank(x) - 1, -1, -1)
-@tf_export("reduce_join")
-def reduce_join(inputs, axis=None,
+@tf_export("strings.reduce_join", "reduce_join")
+@deprecation.deprecated_endpoints("reduce_join")
+def reduce_join(inputs, axis=None, # pylint: disable=missing-docstring
keep_dims=False,
separator="",
name=None,
@@ -329,6 +330,8 @@
reduce_join.__doc__ = deprecation.rewrite_argument_docstring(
gen_string_ops.reduce_join.__doc__, "reduction_indices", "axis")
+reduce_join.__doc__ = reduce_join.__doc__.replace("tf.reduce_join(",
+ "tf.strings.reduce_join(")
# This wrapper provides backwards compatibility for code that predates the
diff --git a/tensorflow/python/ops/variable_scope.py b/tensorflow/python/ops/variable_scope.py
index a43676c..5032ca7 100644
--- a/tensorflow/python/ops/variable_scope.py
+++ b/tensorflow/python/ops/variable_scope.py
@@ -198,7 +198,7 @@
VariableAggregation = variables.VariableAggregation # pylint: disable=invalid-name
AUTO_REUSE = _ReuseMode.AUTO_REUSE
-tf_export("AUTO_REUSE").export_constant(__name__, "AUTO_REUSE")
+tf_export(v1=["AUTO_REUSE"]).export_constant(__name__, "AUTO_REUSE")
AUTO_REUSE.__doc__ = """
When passed in as the value for the `reuse` flag, AUTO_REUSE indicates that
get_variable() should create the requested variable if it doesn't exist or, if
@@ -515,8 +515,10 @@
"synchronization": synchronization,
"aggregation": aggregation,
}
- # `fn_args` can handle functions, `functools.partial`, `lambda`.
- if "constraint" in function_utils.fn_args(custom_getter):
+ # `fn_args` and `has_kwargs` can handle functions, `functools.partial`,
+ # `lambda`.
+ if ("constraint" in function_utils.fn_args(custom_getter) or
+ function_utils.has_kwargs(custom_getter)):
custom_getter_kwargs["constraint"] = constraint
return custom_getter(**custom_getter_kwargs)
else:
@@ -906,7 +908,7 @@
if use_resource is None:
# Set the default value if unspecified.
use_resource = _DEFAULT_USE_RESOURCE
- v = variable(
+ v = variables.VariableV1(
initial_value=init_val,
name=name,
trainable=trainable,
@@ -937,7 +939,8 @@
if regularizer:
with ops.colocate_with(v):
with ops.name_scope(name + "/Regularizer/"):
- loss = regularizer(v)
+ with ops.init_scope():
+ loss = regularizer(v)
if loss is not None:
if context.executing_eagerly():
v_name = "v_%s" % type(v)
@@ -992,7 +995,7 @@
# TODO(alive): support caching devices and partitioned variables in Eager mode.
-@tf_export("VariableScope")
+@tf_export(v1=["VariableScope"])
class VariableScope(object):
"""Variable scope object to carry defaults to provide to `get_variable`.
@@ -1340,7 +1343,7 @@
return scope_store
-@tf_export("get_variable_scope")
+@tf_export(v1=["get_variable_scope"])
def get_variable_scope():
"""Returns the current variable scope."""
return get_variable_scope_store().current_scope
@@ -1449,7 +1452,7 @@
# The argument list for get_variable must match arguments to get_local_variable.
# So, if you are updating the arguments, also update arguments to
# get_local_variable below.
-@tf_export("get_variable")
+@tf_export(v1=["get_variable"])
def get_variable(name,
shape=None,
dtype=None,
@@ -1594,7 +1597,7 @@
# The argument list for get_local_variable must match arguments to get_variable.
# So, if you are updating the arguments, also update arguments to get_variable.
-@tf_export("get_local_variable")
+@tf_export(v1=["get_local_variable"])
def get_local_variable( # pylint: disable=missing-docstring
name,
shape=None,
@@ -1939,7 +1942,7 @@
# Named like a function for backwards compatibility with the
# @tf_contextlib.contextmanager version, which was switched to a class to avoid
# some object creation overhead.
-@tf_export("variable_scope") # pylint: disable=invalid-name
+@tf_export(v1=["variable_scope"]) # pylint: disable=invalid-name
class variable_scope(object):
"""A context manager for defining ops that creates variables (layers).
@@ -2320,7 +2323,7 @@
# pylint: disable=g-doc-return-or-yield
-@tf_export("variable_op_scope")
+@tf_export(v1=["variable_op_scope"])
@tf_contextlib.contextmanager
def variable_op_scope(values,
name_or_scope,
@@ -2441,7 +2444,33 @@
expected_shape=expected_shape, import_scope=import_scope)
+def default_variable_creator_v2(next_creator=None, **kwargs):
+ """Default variable creator."""
+ assert next_creator is None
+ initial_value = kwargs.get("initial_value", None)
+ trainable = kwargs.get("trainable", None)
+ validate_shape = kwargs.get("validate_shape", True)
+ caching_device = kwargs.get("caching_device", None)
+ name = kwargs.get("name", None)
+ variable_def = kwargs.get("variable_def", None)
+ dtype = kwargs.get("dtype", None)
+ import_scope = kwargs.get("import_scope", None)
+ constraint = kwargs.get("constraint", None)
+
+ # Set trainable value based on synchronization value.
+ synchronization = kwargs.get("synchronization", VariableSynchronization.AUTO)
+ trainable = _get_trainable_value(
+ synchronization=synchronization, trainable=trainable)
+
+ return resource_variable_ops.ResourceVariable(
+ initial_value=initial_value, trainable=trainable,
+ validate_shape=validate_shape, caching_device=caching_device,
+ name=name, dtype=dtype, constraint=constraint, variable_def=variable_def,
+ import_scope=import_scope)
+
+
variables.default_variable_creator = default_variable_creator
+variables.default_variable_creator_v2 = default_variable_creator_v2
def _make_getter(captured_getter, captured_previous):
@@ -2450,11 +2479,12 @@
# TODO(apassos) remove forwarding symbol
-variable = variables.Variable
+variable = variables.VariableV1
+@tf_export(v1=["variable_creator_scope"])
@tf_contextlib.contextmanager
-def variable_creator_scope(variable_creator):
+def variable_creator_scope_v1(variable_creator):
"""Scope which defines a variable creation function to be used by variable().
variable_creator is expected to be a function with the following signature:
@@ -2525,3 +2555,73 @@
"""
with ops.get_default_graph()._variable_creator_scope(variable_creator): # pylint: disable=protected-access
yield
+
+
+# Note: only the docstrings differ between this and v1.
+@tf_export(v2=["variable_creator_scope"])
+@tf_contextlib.contextmanager
+def variable_creator_scope(variable_creator):
+ """Scope which defines a variable creation function to be used by variable().
+
+ variable_creator is expected to be a function with the following signature:
+
+ ```
+ def variable_creator(next_creator, **kwargs)
+ ```
+
+ The creator is supposed to eventually call the next_creator to create a
+ variable if it does want to create a variable and not call Variable or
+ ResourceVariable directly. This helps make creators composable. A creator may
+ choose to create multiple variables, return already existing variables, or
+ simply register that a variable was created and defer to the next creators in
+ line. Creators can also modify the keyword arguments seen by the next
+ creators.
+
+ Custom getters in the variable scope will eventually resolve down to these
+ custom creators when they do create variables.
+
+ The valid keyword arguments in kwds are:
+ initial_value: A `Tensor`, or Python object convertible to a `Tensor`,
+ which is the initial value for the Variable. The initial value must have
+ a shape specified unless `validate_shape` is set to False. Can also be a
+ callable with no argument that returns the initial value when called. In
+ that case, `dtype` must be specified. (Note that initializer functions
+ from init_ops.py must first be bound to a shape before being used here.)
+ trainable: If `True`, the default, GradientTapes automatically watch
+ uses of this Variable.
+ validate_shape: If `False`, allows the variable to be initialized with a
+ value of unknown shape. If `True`, the default, the shape of
+ `initial_value` must be known.
+ caching_device: Optional device string describing where the Variable
+ should be cached for reading. Defaults to the Variable's device.
+ If not `None`, caches on another device. Typical use is to cache
+ on the device where the Ops using the Variable reside, to deduplicate
+ copying through `Switch` and other conditional statements.
+ name: Optional name for the variable. Defaults to `'Variable'` and gets
+ uniquified automatically.
+ dtype: If set, initial_value will be converted to the given type.
+ If `None`, either the datatype will be kept (if `initial_value` is
+ a Tensor), or `convert_to_tensor` will decide.
+ constraint: A constraint function to be applied to the variable after
+ updates by some algorithms.
+ synchronization: Indicates when a distributed a variable will be
+ aggregated. Accepted values are constants defined in the class
+ `tf.VariableSynchronization`. By default the synchronization is set to
+ `AUTO` and the current `DistributionStrategy` chooses
+ when to synchronize. If `synchronization` is set to `ON_READ`,
+ `trainable` must not be set to `True`.
+ aggregation: Indicates how a distributed variable will be aggregated.
+ Accepted values are constants defined in the class
+ `tf.VariableAggregation`.
+
+ This set may grow over time, so it's important the signature of creators is as
+ mentioned above.
+
+ Args:
+ variable_creator: the passed creator
+
+ Yields:
+ A scope in which the creator is active
+ """
+ with ops.get_default_graph()._variable_creator_scope(variable_creator): # pylint: disable=protected-access
+ yield
diff --git a/tensorflow/python/ops/variables.py b/tensorflow/python/ops/variables.py
index 7a46157..45c8618 100644
--- a/tensorflow/python/ops/variables.py
+++ b/tensorflow/python/ops/variables.py
@@ -46,6 +46,11 @@
raise NotImplementedError("variable_scope needs to be imported")
+def default_variable_creator_v2(_, **kwds):
+ del kwds
+ raise NotImplementedError("variable_scope needs to be imported")
+
+
def _make_getter(captured_getter, captured_previous):
"""To avoid capturing loop variables."""
def getter(**kwargs):
@@ -101,21 +106,21 @@
class VariableMetaclass(type):
"""Metaclass to allow construction of tf.Variable to be overridden."""
- def _variable_call(cls,
- initial_value=None,
- trainable=None,
- collections=None,
- validate_shape=True,
- caching_device=None,
- name=None,
- variable_def=None,
- dtype=None,
- expected_shape=None,
- import_scope=None,
- constraint=None,
- use_resource=None,
- synchronization=VariableSynchronization.AUTO,
- aggregation=VariableAggregation.NONE):
+ def _variable_v1_call(cls,
+ initial_value=None,
+ trainable=None,
+ collections=None,
+ validate_shape=True,
+ caching_device=None,
+ name=None,
+ variable_def=None,
+ dtype=None,
+ expected_shape=None,
+ import_scope=None,
+ constraint=None,
+ use_resource=None,
+ synchronization=VariableSynchronization.AUTO,
+ aggregation=VariableAggregation.NONE):
"""Call on Variable class. Useful to force the signature."""
previous_getter = lambda **kwargs: default_variable_creator(None, **kwargs)
for getter in ops.get_default_graph()._variable_creator_stack: # pylint: disable=protected-access
@@ -140,14 +145,49 @@
synchronization=synchronization,
aggregation=aggregation)
+ def _variable_v2_call(cls,
+ initial_value=None,
+ trainable=None,
+ validate_shape=True,
+ caching_device=None,
+ name=None,
+ variable_def=None,
+ dtype=None,
+ import_scope=None,
+ constraint=None,
+ synchronization=VariableSynchronization.AUTO,
+ aggregation=VariableAggregation.NONE):
+ """Call on Variable class. Useful to force the signature."""
+ previous_getter = lambda **kws: default_variable_creator_v2(None, **kws)
+ for getter in ops.get_default_graph()._variable_creator_stack: # pylint: disable=protected-access
+ previous_getter = _make_getter(getter, previous_getter)
+
+ # Reset `aggregation` that is explicitly set as `None` to the enum NONE.
+ if aggregation is None:
+ aggregation = VariableAggregation.NONE
+ return previous_getter(
+ initial_value=initial_value,
+ trainable=trainable,
+ validate_shape=validate_shape,
+ caching_device=caching_device,
+ name=name,
+ variable_def=variable_def,
+ dtype=dtype,
+ import_scope=import_scope,
+ constraint=constraint,
+ synchronization=synchronization,
+ aggregation=aggregation)
+
def __call__(cls, *args, **kwargs):
- if cls is Variable:
- return cls._variable_call(*args, **kwargs)
+ if cls is VariableV1:
+ return cls._variable_v1_call(*args, **kwargs)
+ elif cls is Variable:
+ return cls._variable_v2_call(*args, **kwargs)
else:
return super(VariableMetaclass, cls).__call__(*args, **kwargs)
-@tf_export("Variable")
+@tf_export(v2=["Variable"])
class Variable(six.with_metaclass(VariableMetaclass,
checkpointable.CheckpointableBase)):
"""See the [Variables Guide](https://tensorflow.org/guide/variables).
@@ -267,16 +307,13 @@
def __init__(self,
initial_value=None,
trainable=True,
- collections=None,
validate_shape=True,
caching_device=None,
name=None,
variable_def=None,
dtype=None,
- expected_shape=None,
import_scope=None,
constraint=None,
- use_resource=None,
synchronization=VariableSynchronization.AUTO,
aggregation=VariableAggregation.NONE):
"""Creates a new variable with value `initial_value`.
@@ -297,11 +334,8 @@
callable with no argument that returns the initial value when called. In
that case, `dtype` must be specified. (Note that initializer functions
from init_ops.py must first be bound to a shape before being used here.)
- trainable: If `True`, the default, also adds the variable to the graph
- collection `GraphKeys.TRAINABLE_VARIABLES`. This collection is used as
- the default list of variables to use by the `Optimizer` classes.
- collections: List of graph collections keys. The new variable is added to
- these collections. Defaults to `[GraphKeys.GLOBAL_VARIABLES]`.
+ trainable: If `True`, the default, GradientTapes automatically watch uses
+ of this variable.
validate_shape: If `False`, allows the variable to be initialized with a
value of unknown shape. If `True`, the default, the shape of
`initial_value` must be known.
@@ -319,8 +353,6 @@
dtype: If set, initial_value will be converted to the given type.
If `None`, either the datatype will be kept (if `initial_value` is
a Tensor), or `convert_to_tensor` will decide.
- expected_shape: A TensorShape. If set, initial_value is expected
- to have this shape.
import_scope: Optional `string`. Name scope to add to the
`Variable.` Only used when initializing from protocol buffer.
constraint: An optional projection function to be applied to the variable
@@ -330,9 +362,6 @@
variable and return the Tensor for the projected value
(which must have the same shape). Constraints are not safe to
use when doing asynchronous distributed training.
- use_resource: if True, a ResourceVariable is created; otherwise an
- old-style ref-based variable is created. When eager execution is enabled
- a resource variable is always created.
synchronization: Indicates when a distributed a variable will be
aggregated. Accepted values are constants defined in the class
`tf.VariableSynchronization`. By default the synchronization is set to
@@ -1009,11 +1038,207 @@
raise NotImplementedError
+@tf_export(v1=["Variable"])
+class VariableV1(Variable):
+ """See the [Variables Guide](https://tensorflow.org/guide/variables).
+
+ A variable maintains state in the graph across calls to `run()`. You add a
+ variable to the graph by constructing an instance of the class `Variable`.
+
+ The `Variable()` constructor requires an initial value for the variable,
+ which can be a `Tensor` of any type and shape. The initial value defines the
+ type and shape of the variable. After construction, the type and shape of
+ the variable are fixed. The value can be changed using one of the assign
+ methods.
+
+ If you want to change the shape of a variable later you have to use an
+ `assign` Op with `validate_shape=False`.
+
+ Just like any `Tensor`, variables created with `Variable()` can be used as
+ inputs for other Ops in the graph. Additionally, all the operators
+ overloaded for the `Tensor` class are carried over to variables, so you can
+ also add nodes to the graph by just doing arithmetic on variables.
+
+ ```python
+ import tensorflow as tf
+
+ # Create a variable.
+ w = tf.Variable(<initial-value>, name=<optional-name>)
+
+ # Use the variable in the graph like any Tensor.
+ y = tf.matmul(w, ...another variable or tensor...)
+
+ # The overloaded operators are available too.
+ z = tf.sigmoid(w + y)
+
+ # Assign a new value to the variable with `assign()` or a related method.
+ w.assign(w + 1.0)
+ w.assign_add(1.0)
+ ```
+
+ When you launch the graph, variables have to be explicitly initialized before
+ you can run Ops that use their value. You can initialize a variable by
+ running its *initializer op*, restoring the variable from a save file, or
+ simply running an `assign` Op that assigns a value to the variable. In fact,
+ the variable *initializer op* is just an `assign` Op that assigns the
+ variable's initial value to the variable itself.
+
+ ```python
+ # Launch the graph in a session.
+ with tf.Session() as sess:
+ # Run the variable initializer.
+ sess.run(w.initializer)
+ # ...you now can run ops that use the value of 'w'...
+ ```
+
+ The most common initialization pattern is to use the convenience function
+ `global_variables_initializer()` to add an Op to the graph that initializes
+ all the variables. You then run that Op after launching the graph.
+
+ ```python
+ # Add an Op to initialize global variables.
+ init_op = tf.global_variables_initializer()
+
+ # Launch the graph in a session.
+ with tf.Session() as sess:
+ # Run the Op that initializes global variables.
+ sess.run(init_op)
+ # ...you can now run any Op that uses variable values...
+ ```
+
+ If you need to create a variable with an initial value dependent on another
+ variable, use the other variable's `initialized_value()`. This ensures that
+ variables are initialized in the right order.
+
+ All variables are automatically collected in the graph where they are
+ created. By default, the constructor adds the new variable to the graph
+ collection `GraphKeys.GLOBAL_VARIABLES`. The convenience function
+ `global_variables()` returns the contents of that collection.
+
+ When building a machine learning model it is often convenient to distinguish
+ between variables holding the trainable model parameters and other variables
+ such as a `global step` variable used to count training steps. To make this
+ easier, the variable constructor supports a `trainable=<bool>` parameter. If
+ `True`, the new variable is also added to the graph collection
+ `GraphKeys.TRAINABLE_VARIABLES`. The convenience function
+ `trainable_variables()` returns the contents of this collection. The
+ various `Optimizer` classes use this collection as the default list of
+ variables to optimize.
+
+ WARNING: tf.Variable objects by default have a non-intuitive memory model. A
+ Variable is represented internally as a mutable Tensor which can
+ non-deterministically alias other Tensors in a graph. The set of operations
+ which consume a Variable and can lead to aliasing is undetermined and can
+ change across TensorFlow versions. Avoid writing code which relies on the
+ value of a Variable either changing or not changing as other operations
+ happen. For example, using Variable objects or simple functions thereof as
+ predicates in a `tf.cond` is dangerous and error-prone:
+
+ ```
+ v = tf.Variable(True)
+ tf.cond(v, lambda: v.assign(False), my_false_fn) # Note: this is broken.
+ ```
+
+ Here replacing adding `use_resource=True` when constructing the variable will
+ fix any nondeterminism issues:
+ ```
+ v = tf.Variable(True, use_resource=True)
+ tf.cond(v, lambda: v.assign(False), my_false_fn)
+ ```
+
+ To use the replacement for variables which does
+ not have these issues:
+
+ * Add `use_resource=True` when constructing `tf.Variable`;
+ * Call `tf.get_variable_scope().set_use_resource(True)` inside a
+ `tf.variable_scope` before the `tf.get_variable()` call.
+ """
+
+ def __init__(self, # pylint: disable=super-init-not-called
+ initial_value=None,
+ trainable=True,
+ collections=None,
+ validate_shape=True,
+ caching_device=None,
+ name=None,
+ variable_def=None,
+ dtype=None,
+ expected_shape=None,
+ import_scope=None,
+ constraint=None,
+ use_resource=None,
+ synchronization=VariableSynchronization.AUTO,
+ aggregation=VariableAggregation.NONE):
+ """Creates a new variable with value `initial_value`.
+
+ The new variable is added to the graph collections listed in `collections`,
+ which defaults to `[GraphKeys.GLOBAL_VARIABLES]`.
+
+ If `trainable` is `True` the variable is also added to the graph collection
+ `GraphKeys.TRAINABLE_VARIABLES`.
+
+ This constructor creates both a `variable` Op and an `assign` Op to set the
+ variable to its initial value.
+
+ Args:
+ initial_value: A `Tensor`, or Python object convertible to a `Tensor`,
+ which is the initial value for the Variable. The initial value must have
+ a shape specified unless `validate_shape` is set to False. Can also be a
+ callable with no argument that returns the initial value when called. In
+ that case, `dtype` must be specified. (Note that initializer functions
+ from init_ops.py must first be bound to a shape before being used here.)
+ trainable: If `True`, the default, also adds the variable to the graph
+ collection `GraphKeys.TRAINABLE_VARIABLES`. This collection is used as
+ the default list of variables to use by the `Optimizer` classes.
+ collections: List of graph collections keys. The new variable is added to
+ these collections. Defaults to `[GraphKeys.GLOBAL_VARIABLES]`.
+ validate_shape: If `False`, allows the variable to be initialized with a
+ value of unknown shape. If `True`, the default, the shape of
+ `initial_value` must be known.
+ caching_device: Optional device string describing where the Variable
+ should be cached for reading. Defaults to the Variable's device.
+ If not `None`, caches on another device. Typical use is to cache
+ on the device where the Ops using the Variable reside, to deduplicate
+ copying through `Switch` and other conditional statements.
+ name: Optional name for the variable. Defaults to `'Variable'` and gets
+ uniquified automatically.
+ variable_def: `VariableDef` protocol buffer. If not `None`, recreates
+ the Variable object with its contents, referencing the variable's nodes
+ in the graph, which must already exist. The graph is not changed.
+ `variable_def` and the other arguments are mutually exclusive.
+ dtype: If set, initial_value will be converted to the given type.
+ If `None`, either the datatype will be kept (if `initial_value` is
+ a Tensor), or `convert_to_tensor` will decide.
+ expected_shape: A TensorShape. If set, initial_value is expected
+ to have this shape.
+ import_scope: Optional `string`. Name scope to add to the
+ `Variable.` Only used when initializing from protocol buffer.
+ constraint: An optional projection function to be applied to the variable
+ after being updated by an `Optimizer` (e.g. used to implement norm
+ constraints or value constraints for layer weights). The function must
+ take as input the unprojected Tensor representing the value of the
+ variable and return the Tensor for the projected value
+ (which must have the same shape). Constraints are not safe to
+ use when doing asynchronous distributed training.
+ use_resource: whether to use resource variables.
+ synchronization: unused
+ aggregation: unused
+
+ Raises:
+ ValueError: If both `variable_def` and initial_value are specified.
+ ValueError: If the initial value is not specified, or does not have a
+ shape and `validate_shape` is `True`.
+ RuntimeError: If eager execution is enabled.
+ """
+
+ SaveSliceInfo = Variable.SaveSliceInfo
+
+
# TODO(apassos): do not repeat all comments here
-class RefVariable(Variable):
+class RefVariable(VariableV1):
"""Ref-based implementation of variables."""
- def __init__(self,
+ def __init__(self, # pylint: disable=super-init-not-called
initial_value=None,
trainable=True,
collections=None,
@@ -1873,7 +2098,7 @@
def _OverloadAllOperators(): # pylint: disable=invalid-name
"""Register overloads for all operators."""
for operator in ops.Tensor.OVERLOADABLE_OPERATORS:
- Variable._OverloadOperator(operator)
+ Variable._OverloadOperator(operator) # pylint: disable=protected-access
# For slicing, bind getitem differently than a tensor (use SliceHelperVar
# instead)
# pylint: disable=protected-access
@@ -2395,13 +2620,53 @@
def _get_partitions(self):
return self._partitions
- def assign(self, value, use_locking=False):
- _ = value, use_locking
- raise NotImplementedError(
- "assign() has not been implemented for PartitionedVariable.")
+ def _apply_assign_fn(self, assign_fn, value):
+ partition_axes = self._partition_axes()
+ if len(partition_axes) > 1:
+ raise NotImplementedError(
+ "Cannot do assign action along more than one dimension: %s. "
+ "Multi-axis partition assign action is not supported " %
+ str(partition_axes))
+ partition_ix = partition_axes[0]
+ size_splits_list = [
+ var.shape[partition_ix].value for var in self._variable_list
+ ]
+ value_list = array_ops.split(value, size_splits_list, axis=partition_ix)
+ op_list = [
+ assign_fn(var, value_list[idx], idx)
+ for idx, var in enumerate(self._variable_list)
+ ]
+ return op_list
+
+ def assign(self, value, use_locking=False, name=None, read_value=True):
+ assign_fn = lambda var, r_value, idx: var.assign(
+ r_value, use_locking=use_locking,
+ name="%s_%d" % (name, idx), read_value=read_value)
+ assign_list = self._apply_assign_fn(assign_fn, value)
+ if read_value:
+ return assign_list
+ return [assign.op for assign in assign_list]
+
+ def assign_add(self, value, use_locking=False, name=None, read_value=True):
+ assign_fn = lambda var, r_value, idx: var.assign_add(
+ r_value, use_locking=use_locking,
+ name="%s_%d" % (name, idx), read_value=read_value)
+ assign_list = self._apply_assign_fn(assign_fn, value)
+ if read_value:
+ return assign_list
+ return [assign.op for assign in assign_list]
+
+ def assign_sub(self, value, use_locking=False, name=None, read_value=True):
+ assign_fn = lambda var, r_value, idx: var.assign_sub(
+ r_value, use_locking=use_locking,
+ name="%s_%d" % (name, idx), read_value=read_value)
+ assign_list = self._apply_assign_fn(assign_fn, value)
+ if read_value:
+ return assign_list
+ return [assign.op for assign in assign_list]
-@tf_export("global_variables")
+@tf_export(v1=["global_variables"])
def global_variables(scope=None):
"""Returns global variables.
@@ -2427,7 +2692,7 @@
return ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES, scope)
-@tf_export("all_variables")
+@tf_export(v1=["all_variables"])
@deprecated("2017-03-02", "Please use tf.global_variables instead.")
def all_variables():
"""See `tf.global_variables`."""
@@ -2452,7 +2717,7 @@
ops.get_collection(ops.GraphKeys.SAVEABLE_OBJECTS, scope))
-@tf_export("local_variables")
+@tf_export(v1=["local_variables"])
def local_variables(scope=None):
"""Returns local variables.
@@ -2480,7 +2745,7 @@
return ops.get_collection(ops.GraphKeys.LOCAL_VARIABLES, scope)
-@tf_export("model_variables")
+@tf_export(v1=["model_variables"])
def model_variables(scope=None):
"""Returns all variables in the MODEL_VARIABLES collection.
@@ -2497,7 +2762,7 @@
return ops.get_collection(ops.GraphKeys.MODEL_VARIABLES, scope)
-@tf_export("trainable_variables")
+@tf_export(v1=["trainable_variables"])
def trainable_variables(scope=None):
"""Returns all variables created with `trainable=True`.
@@ -2519,7 +2784,7 @@
return ops.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES, scope)
-@tf_export("moving_average_variables")
+@tf_export(v1=["moving_average_variables"])
def moving_average_variables(scope=None):
"""Returns all variables that maintain their moving averages.
@@ -2541,7 +2806,7 @@
return ops.get_collection(ops.GraphKeys.MOVING_AVERAGE_VARIABLES, scope)
-@tf_export("initializers.variables", "variables_initializer")
+@tf_export(v1=["initializers.variables", "variables_initializer"])
def variables_initializer(var_list, name="init"):
"""Returns an Op that initializes a list of variables.
@@ -2567,7 +2832,7 @@
return control_flow_ops.no_op(name=name)
-@tf_export("initialize_variables")
+@tf_export(v1=["initialize_variables"])
@tf_should_use.should_use_result
@deprecated("2017-03-02", "Use `tf.variables_initializer` instead.")
def initialize_variables(var_list, name="init"):
@@ -2575,7 +2840,7 @@
return variables_initializer(var_list, name=name)
-@tf_export("initializers.global_variables", "global_variables_initializer")
+@tf_export(v1=["initializers.global_variables", "global_variables_initializer"])
def global_variables_initializer():
"""Returns an Op that initializes global variables.
@@ -2589,7 +2854,7 @@
return variables_initializer(global_variables())
-@tf_export("initialize_all_variables")
+@tf_export(v1=["initialize_all_variables"])
@tf_should_use.should_use_result
@deprecated("2017-03-02", "Use `tf.global_variables_initializer` instead.")
def initialize_all_variables():
@@ -2597,7 +2862,7 @@
return global_variables_initializer()
-@tf_export("initializers.local_variables", "local_variables_initializer")
+@tf_export(v1=["initializers.local_variables", "local_variables_initializer"])
def local_variables_initializer():
"""Returns an Op that initializes all local variables.
@@ -2611,7 +2876,7 @@
return variables_initializer(local_variables())
-@tf_export("initialize_local_variables")
+@tf_export(v1=["initialize_local_variables"])
@tf_should_use.should_use_result
@deprecated("2017-03-02", "Use `tf.local_variables_initializer` instead.")
def initialize_local_variables():
@@ -2619,7 +2884,7 @@
return local_variables_initializer()
-@tf_export("is_variable_initialized")
+@tf_export(v1=["is_variable_initialized"])
@tf_should_use.should_use_result
def is_variable_initialized(variable):
"""Tests if a variable has been initialized.
@@ -2634,7 +2899,7 @@
return state_ops.is_variable_initialized(variable)
-@tf_export("assert_variables_initialized")
+@tf_export(v1=["assert_variables_initialized"])
@tf_should_use.should_use_result
def assert_variables_initialized(var_list=None):
"""Returns an Op to check if variables are initialized.
@@ -2677,7 +2942,7 @@
return array_ops.stack(ranks)
-@tf_export("report_uninitialized_variables")
+@tf_export(v1=["report_uninitialized_variables"])
@tf_should_use.should_use_result
def report_uninitialized_variables(var_list=None,
name="report_uninitialized_variables"):
diff --git a/tensorflow/python/ops/while_v2.py b/tensorflow/python/ops/while_v2.py
index 875be31..6791e1c 100644
--- a/tensorflow/python/ops/while_v2.py
+++ b/tensorflow/python/ops/while_v2.py
@@ -24,6 +24,7 @@
from __future__ import print_function
import collections
+import sys
from tensorflow.core.framework import attr_value_pb2
from tensorflow.python.eager import function
@@ -33,6 +34,7 @@
from tensorflow.python.framework import tensor_shape
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import cond_v2_impl as cond_v2
+from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import control_flow_util
from tensorflow.python.ops import gen_functional_ops
from tensorflow.python.ops import gradients_impl
@@ -41,6 +43,8 @@
# pylint: disable=protected-access
+control_flow_ops._while_v2 = sys.modules[__name__]
+
# TODO(b/79881896): Handle external control dependencies. tf.while_loop allows
# control dependencies on external nodes with at least 1 output.
# Another idea is to create const nodes outside the loop and add control edges
diff --git a/tensorflow/python/pywrap_tfe.i b/tensorflow/python/pywrap_tfe.i
index c411a58..61e0abb 100755
--- a/tensorflow/python/pywrap_tfe.i
+++ b/tensorflow/python/pywrap_tfe.i
@@ -67,6 +67,7 @@
%rename("%s") TFE_ContextStartStep;
%rename("%s") TFE_ContextEndStep;
%rename("%s") TFE_Py_RegisterVSpace;
+%rename("%s") TFE_Py_EncodeArg;
%{
#include "tensorflow/python/eager/pywrap_tfe.h"
diff --git a/tensorflow/python/saved_model/builder_impl.py b/tensorflow/python/saved_model/builder_impl.py
index 8e7f123..8bf057f 100644
--- a/tensorflow/python/saved_model/builder_impl.py
+++ b/tensorflow/python/saved_model/builder_impl.py
@@ -36,10 +36,13 @@
from tensorflow.python.training import saver as tf_saver
from tensorflow.python.util import compat
from tensorflow.python.util.deprecation import deprecated_args
+from tensorflow.python.util.deprecation import deprecated_endpoints
from tensorflow.python.util.tf_export import tf_export
-@tf_export("saved_model.builder.SavedModelBuilder")
+@tf_export("saved_model.Builder",
+ "saved_model.builder.SavedModelBuilder")
+@deprecated_endpoints("saved_model.builder.SavedModelBuilder")
class SavedModelBuilder(object):
"""Builds the `SavedModel` protocol buffer and saves variables and assets.
@@ -61,7 +64,7 @@
Typical usage for the `SavedModelBuilder`:
```python
...
- builder = tf.saved_model.builder.SavedModelBuilder(export_dir)
+ builder = tf.saved_model.Builder(export_dir)
with tf.Session(graph=tf.Graph()) as sess:
...
diff --git a/tensorflow/python/saved_model/loader_impl.py b/tensorflow/python/saved_model/loader_impl.py
index e853610..895644a 100644
--- a/tensorflow/python/saved_model/loader_impl.py
+++ b/tensorflow/python/saved_model/loader_impl.py
@@ -34,6 +34,7 @@
from tensorflow.python.saved_model import utils_impl as saved_model_utils
from tensorflow.python.training import saver as tf_saver
from tensorflow.python.util import compat
+from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
@@ -144,7 +145,10 @@
return main_op_tensor
-@tf_export("saved_model.loader.maybe_saved_model_directory")
+@tf_export("saved_model.maybe_saved_model_directory",
+ "saved_model.loader.maybe_saved_model_directory")
+@deprecation.deprecated_endpoints(
+ "saved_model.loader.maybe_saved_model_directory")
def maybe_saved_model_directory(export_dir):
"""Checks whether the provided export directory could contain a SavedModel.
@@ -165,7 +169,7 @@
return file_io.file_exists(txt_path) or file_io.file_exists(pb_path)
-@tf_export("saved_model.loader.load")
+@tf_export("saved_model.load", "saved_model.loader.load")
def load(sess, tags, export_dir, import_scope=None, **saver_kwargs):
"""Loads the model from a SavedModel as specified by tags.
diff --git a/tensorflow/python/saved_model/loader_test.py b/tensorflow/python/saved_model/loader_test.py
index b7e217a..924b2e7 100644
--- a/tensorflow/python/saved_model/loader_test.py
+++ b/tensorflow/python/saved_model/loader_test.py
@@ -47,8 +47,8 @@
def setUp(self):
"""Write test SavedModels to a temp directory."""
with session.Session(graph=ops.Graph()) as sess:
- x = variables.Variable(5, name="x")
- y = variables.Variable(11, name="y")
+ x = variables.VariableV1(5, name="x")
+ y = variables.VariableV1(11, name="y")
z = x + y
sess.run(variables.global_variables_initializer())
@@ -134,8 +134,8 @@
def test_restore_variables(self):
loader = loader_impl.SavedModelLoader(SAVED_MODEL_WITH_MAIN_OP)
with self.session(graph=ops.Graph()) as sess:
- x = variables.Variable(0, name="x")
- y = variables.Variable(0, name="y")
+ x = variables.VariableV1(0, name="x")
+ y = variables.VariableV1(0, name="y")
z = x * y
sess.run(variables.global_variables_initializer())
@@ -186,8 +186,10 @@
"""
path = _get_export_dir("no_variable_saved_model")
with session.Session(graph=ops.Graph()) as sess:
- x = variables.Variable(5, name="x", collections=["not_global_variable"])
- y = variables.Variable(11, name="y", collections=["not_global_variable"])
+ x = variables.VariableV1(
+ 5, name="x", collections=["not_global_variable"])
+ y = variables.VariableV1(
+ 11, name="y", collections=["not_global_variable"])
self.assertFalse(variables._all_saveable_objects())
z = x + y
sess.run(variables.variables_initializer([x, y]))
diff --git a/tensorflow/python/saved_model/main_op_impl.py b/tensorflow/python/saved_model/main_op_impl.py
index 631ee63..ad4511b 100644
--- a/tensorflow/python/saved_model/main_op_impl.py
+++ b/tensorflow/python/saved_model/main_op_impl.py
@@ -22,6 +22,7 @@
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import lookup_ops
from tensorflow.python.ops import variables
+from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
@@ -42,7 +43,9 @@
# TODO(sukritiramesh): Integrate with Saver for complete restore functionality.
-@tf_export('saved_model.main_op.main_op_with_restore')
+@tf_export('saved_model.main_op_with_restore',
+ 'saved_model.main_op.main_op_with_restore')
+@deprecation.deprecated_endpoints('saved_model.main_op.main_op_with_restore')
def main_op_with_restore(restore_op_name):
"""Returns a main op to init variables, tables and restore the graph.
diff --git a/tensorflow/python/saved_model/saved_model_test.py b/tensorflow/python/saved_model/saved_model_test.py
index 49d52d3..80b75b7 100644
--- a/tensorflow/python/saved_model/saved_model_test.py
+++ b/tensorflow/python/saved_model/saved_model_test.py
@@ -60,7 +60,7 @@
return os.path.join(test.get_temp_dir(), label)
def _init_and_validate_variable(self, sess, variable_name, variable_value):
- v = variables.Variable(variable_value, name=variable_name)
+ v = variables.VariableV1(variable_value, name=variable_name)
sess.run(variables.global_variables_initializer())
self.assertEqual(variable_value, v.eval())
@@ -458,7 +458,7 @@
# Graph with a single variable added to a collection. SavedModel invoked to:
# - add with weights.
with self.session(graph=ops.Graph()) as sess:
- v = variables.Variable(42, name="v")
+ v = variables.VariableV1(42, name="v")
ops.add_to_collection("foo_vars", v)
sess.run(variables.global_variables_initializer())
self.assertEqual(42, v.eval())
@@ -468,7 +468,7 @@
# SavedModel invoked to:
# - simply add the model (weights are not updated).
with self.session(graph=ops.Graph()) as sess:
- v = variables.Variable(43, name="v")
+ v = variables.VariableV1(43, name="v")
ops.add_to_collection("bar_vars", v)
sess.run(variables.global_variables_initializer())
self.assertEqual(43, v.eval())
@@ -780,13 +780,13 @@
with self.session(graph=ops.Graph()) as sess:
# Add `v1` and `v2` variables to the graph.
- v1 = variables.Variable(1, name="v1")
+ v1 = variables.VariableV1(1, name="v1")
ops.add_to_collection("v", v1)
- v2 = variables.Variable(2, name="v2")
+ v2 = variables.VariableV1(2, name="v2")
ops.add_to_collection("v", v2)
# Initialize another variable `v3` to 42.
- v3 = variables.Variable(42, name="v3")
+ v3 = variables.VariableV1(42, name="v3")
ops.add_to_collection("v", v3)
# Set up an assignment op to be run as part of the main_op.
@@ -815,13 +815,13 @@
with self.session(graph=ops.Graph()) as sess:
# Add `v1` and `v2` variables to the graph.
- v1 = variables.Variable(1, name="v1")
+ v1 = variables.VariableV1(1, name="v1")
ops.add_to_collection("v", v1)
- v2 = variables.Variable(2, name="v2")
+ v2 = variables.VariableV1(2, name="v2")
ops.add_to_collection("v", v2)
# Initialize another variable `v3` to 42.
- v3 = variables.Variable(42, name="v3", trainable=False, collections=[])
+ v3 = variables.VariableV1(42, name="v3", trainable=False, collections=[])
ops.add_to_collection("v", v3)
# Set up an assignment op to be run as part of the legacy_init_op.
@@ -860,11 +860,11 @@
g = ops.Graph()
with self.session(graph=g) as sess:
# Initialize variable `v1` to 1.
- v1 = variables.Variable(1, name="v1")
+ v1 = variables.VariableV1(1, name="v1")
ops.add_to_collection("v", v1)
# Initialize another variable `v2` to 42.
- v2 = variables.Variable(42, name="v2", trainable=False, collections=[])
+ v2 = variables.VariableV1(42, name="v2", trainable=False, collections=[])
ops.add_to_collection("v", v2)
# Set up an assignment op to be run as part of the init op.
@@ -889,9 +889,9 @@
with self.session(graph=ops.Graph()) as sess:
# Add `v1` and `v2` variables to the graph.
- v1 = variables.Variable(1, name="v1")
+ v1 = variables.VariableV1(1, name="v1")
ops.add_to_collection("v", v1)
- v2 = variables.Variable(2, name="v2")
+ v2 = variables.VariableV1(2, name="v2")
ops.add_to_collection("v", v2)
sess.run(variables.global_variables_initializer())
@@ -918,9 +918,9 @@
with self.session(graph=ops.Graph()) as sess:
# Add `v1` and `v2` variables to the graph.
- v1 = variables.Variable(1, name="v1")
+ v1 = variables.VariableV1(1, name="v1")
ops.add_to_collection("v", v1)
- v2 = variables.Variable(2, name="v2")
+ v2 = variables.VariableV1(2, name="v2")
ops.add_to_collection("v", v2)
sess.run(variables.global_variables_initializer())
@@ -947,9 +947,9 @@
with self.session(graph=ops.Graph()) as sess:
# Add `v1` and `v2` variables to the graph.
- v1 = variables.Variable(1, name="v1")
+ v1 = variables.VariableV1(1, name="v1")
ops.add_to_collection("v", v1)
- v2 = variables.Variable(2, name="v2")
+ v2 = variables.VariableV1(2, name="v2")
ops.add_to_collection("v", v2)
sess.run(variables.global_variables_initializer())
@@ -1071,13 +1071,13 @@
graph=ops.Graph(),
config=config_pb2.ConfigProto(device_count={"CPU": 2})) as sess:
with sess.graph.device("/cpu:0"):
- v1 = variables.Variable(1, name="v1")
+ v1 = variables.VariableV1(1, name="v1")
with sess.graph.device("/cpu:1"):
- v2 = variables.Variable(2, name="v2")
+ v2 = variables.VariableV1(2, name="v2")
# v3 is an unsaved variable derived from v1 and v2. It is used to
# exercise the ability to run an init op when restoring a graph.
- v3 = variables.Variable(1, name="v3", trainable=False, collections=[])
+ v3 = variables.VariableV1(1, name="v3", trainable=False, collections=[])
assign_v3 = state_ops.assign(v3, math_ops.add(v1, v2))
init_op = control_flow_ops.group(assign_v3, name="init_op")
@@ -1140,7 +1140,7 @@
builder = saved_model_builder.SavedModelBuilder(export_dir)
with self.session(graph=ops.Graph()) as sess:
- variables.Variable(1, name="v1")
+ variables.VariableV1(1, name="v1")
sess.run(variables.global_variables_initializer())
custom_saver = training.Saver(name="my_saver")
builder.add_meta_graph_and_variables(sess, ["tag"], saver=custom_saver)
@@ -1162,7 +1162,7 @@
builder = saved_model_builder.SavedModelBuilder(export_dir)
with self.session(graph=ops.Graph()) as sess:
- variables.Variable(1, name="v1")
+ variables.VariableV1(1, name="v1")
sess.run(variables.global_variables_initializer())
training.Saver(name="my_saver")
builder.add_meta_graph_and_variables(sess, ["tag"])
@@ -1184,7 +1184,7 @@
builder = saved_model_builder.SavedModelBuilder(export_dir)
with self.session(graph=ops.Graph()) as sess:
- variables.Variable(1, name="v1")
+ variables.VariableV1(1, name="v1")
sess.run(variables.global_variables_initializer())
builder.add_meta_graph_and_variables(sess, ["tag_0"])
@@ -1293,8 +1293,8 @@
# Add a graph with two float32 variables and a Complex Op composing them
# with strip_default_attrs enabled.
with session.Session(graph=ops.Graph()) as sess:
- real_num = variables.Variable(1.0, dtype=dtypes.float32, name="real")
- imag_num = variables.Variable(2.0, dtype=dtypes.float32, name="imag")
+ real_num = variables.VariableV1(1.0, dtype=dtypes.float32, name="real")
+ imag_num = variables.VariableV1(2.0, dtype=dtypes.float32, name="imag")
math_ops.complex(real_num, imag_num, name="complex")
sess.run(variables.global_variables_initializer())
builder.add_meta_graph_and_variables(
@@ -1303,8 +1303,8 @@
# Add a graph with the same float32 variables and a Complex Op composing
# them with strip_default_attrs disabled.
with session.Session(graph=ops.Graph()) as sess:
- real_num = variables.Variable(1.0, dtype=dtypes.float32, name="real")
- imag_num = variables.Variable(2.0, dtype=dtypes.float32, name="imag")
+ real_num = variables.VariableV1(1.0, dtype=dtypes.float32, name="real")
+ imag_num = variables.VariableV1(2.0, dtype=dtypes.float32, name="imag")
math_ops.complex(real_num, imag_num, name="complex")
sess.run(variables.global_variables_initializer())
builder.add_meta_graph(["bar"], strip_default_attrs=False)
@@ -1366,7 +1366,7 @@
# Add a graph with a single variable and a test op with a defaultless
# float32 attr, "test_attr".
with session.Session(graph=ops.Graph()) as sess:
- variables.Variable(1.0, dtype=dtypes.float64, name="var")
+ variables.VariableV1(1.0, dtype=dtypes.float64, name="var")
test_ops.test_attr(T=dtypes.float32, name="test_attr")
sess.run(variables.global_variables_initializer())
builder.add_meta_graph_and_variables(sess, ["foo"])
diff --git a/tensorflow/python/saved_model/signature_def_utils_impl.py b/tensorflow/python/saved_model/signature_def_utils_impl.py
index 37f927f..a103441 100644
--- a/tensorflow/python/saved_model/signature_def_utils_impl.py
+++ b/tensorflow/python/saved_model/signature_def_utils_impl.py
@@ -24,10 +24,14 @@
from tensorflow.python.framework import ops
from tensorflow.python.saved_model import signature_constants
from tensorflow.python.saved_model import utils
+from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
-@tf_export('saved_model.signature_def_utils.build_signature_def')
+@tf_export('saved_model.build_signature_def',
+ 'saved_model.signature_def_utils.build_signature_def')
+@deprecation.deprecated_endpoints(
+ 'saved_model.signature_def_utils.build_signature_def')
def build_signature_def(inputs=None, outputs=None, method_name=None):
"""Utility function to build a SignatureDef protocol buffer.
@@ -53,7 +57,10 @@
return signature_def
-@tf_export('saved_model.signature_def_utils.regression_signature_def')
+@tf_export('saved_model.regression_signature_def',
+ 'saved_model.signature_def_utils.regression_signature_def')
+@deprecation.deprecated_endpoints(
+ 'saved_model.signature_def_utils.regression_signature_def')
def regression_signature_def(examples, predictions):
"""Creates regression signature from given examples and predictions.
@@ -95,7 +102,10 @@
return signature_def
-@tf_export('saved_model.signature_def_utils.classification_signature_def')
+@tf_export('saved_model.classification_signature_def',
+ 'saved_model.signature_def_utils.classification_signature_def')
+@deprecation.deprecated_endpoints(
+ 'saved_model.signature_def_utils.classification_signature_def')
def classification_signature_def(examples, classes, scores):
"""Creates classification signature from given examples and predictions.
@@ -148,7 +158,10 @@
return signature_def
-@tf_export('saved_model.signature_def_utils.predict_signature_def')
+@tf_export('saved_model.predict_signature_def',
+ 'saved_model.signature_def_utils.predict_signature_def')
+@deprecation.deprecated_endpoints(
+ 'saved_model.signature_def_utils.predict_signature_def')
def predict_signature_def(inputs, outputs):
"""Creates prediction signature from given inputs and outputs.
@@ -239,7 +252,10 @@
return signature_def
-@tf_export('saved_model.signature_def_utils.is_valid_signature')
+@tf_export('saved_model.is_valid_signature',
+ 'saved_model.signature_def_utils.is_valid_signature')
+@deprecation.deprecated_endpoints(
+ 'saved_model.signature_def_utils.is_valid_signature')
def is_valid_signature(signature_def):
"""Determine whether a SignatureDef can be served by TensorFlow Serving."""
if signature_def is None:
@@ -313,4 +329,3 @@
return False
return True
-
diff --git a/tensorflow/python/saved_model/utils_impl.py b/tensorflow/python/saved_model/utils_impl.py
index 06d0932..0bba7b6 100644
--- a/tensorflow/python/saved_model/utils_impl.py
+++ b/tensorflow/python/saved_model/utils_impl.py
@@ -27,13 +27,16 @@
from tensorflow.python.lib.io import file_io
from tensorflow.python.saved_model import constants
from tensorflow.python.util import compat
+from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
# TensorInfo helpers.
-@tf_export("saved_model.utils.build_tensor_info")
+@tf_export("saved_model.build_tensor_info",
+ "saved_model.utils.build_tensor_info")
+@deprecation.deprecated_endpoints("saved_model.utils.build_tensor_info")
def build_tensor_info(tensor):
"""Utility function to build TensorInfo proto.
@@ -57,7 +60,10 @@
return tensor_info
-@tf_export("saved_model.utils.get_tensor_from_tensor_info")
+@tf_export("saved_model.get_tensor_from_tensor_info",
+ "saved_model.utils.get_tensor_from_tensor_info")
+@deprecation.deprecated_endpoints(
+ "saved_model.utils.get_tensor_from_tensor_info")
def get_tensor_from_tensor_info(tensor_info, graph=None, import_scope=None):
"""Returns the Tensor or SparseTensor described by a TensorInfo proto.
diff --git a/tensorflow/python/tools/BUILD b/tensorflow/python/tools/BUILD
index 75824d8..384c7a8 100644
--- a/tensorflow/python/tools/BUILD
+++ b/tensorflow/python/tools/BUILD
@@ -8,6 +8,7 @@
exports_files(["LICENSE"])
load("//tensorflow:tensorflow.bzl", "py_test")
+load("//tensorflow:tensorflow.bzl", "py_binary")
# Transitive dependencies of this target will be included in the pip package.
py_library(
@@ -21,6 +22,13 @@
":saved_model_cli",
":saved_model_utils",
":strip_unused",
+ # The following py_library are needed because
+ # py_binary may not depend on them when --define=no_tensorflow_py_deps=true
+ # is specified. See https://github.com/tensorflow/tensorflow/issues/22390
+ ":freeze_graph_lib",
+ ":optimize_for_inference_lib",
+ ":selective_registration_header_lib",
+ ":strip_unused_lib",
],
)
diff --git a/tensorflow/python/tools/api/generator/api_init_files.bzl b/tensorflow/python/tools/api/generator/api_init_files.bzl
index 92446e2..533a138 100644
--- a/tensorflow/python/tools/api/generator/api_init_files.bzl
+++ b/tensorflow/python/tools/api/generator/api_init_files.bzl
@@ -8,6 +8,7 @@
"bitwise/__init__.py",
"compat/__init__.py",
"data/__init__.py",
+ "data/experimental/__init__.py",
"debugging/__init__.py",
"distributions/__init__.py",
"dtypes/__init__.py",
@@ -69,6 +70,7 @@
"profiler/__init__.py",
"python_io/__init__.py",
"quantization/__init__.py",
+ "random/__init__.py",
"resource_loader/__init__.py",
"strings/__init__.py",
"saved_model/__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 bc2f351..0747424 100644
--- a/tensorflow/python/tools/api/generator/api_init_files_v1.bzl
+++ b/tensorflow/python/tools/api/generator/api_init_files_v1.bzl
@@ -8,6 +8,7 @@
"bitwise/__init__.py",
"compat/__init__.py",
"data/__init__.py",
+ "data/experimental/__init__.py",
"debugging/__init__.py",
"distributions/__init__.py",
"dtypes/__init__.py",
@@ -69,6 +70,7 @@
"profiler/__init__.py",
"python_io/__init__.py",
"quantization/__init__.py",
+ "random/__init__.py",
"resource_loader/__init__.py",
"strings/__init__.py",
"saved_model/__init__.py",
diff --git a/tensorflow/python/tools/freeze_graph_test.py b/tensorflow/python/tools/freeze_graph_test.py
index e38945f..5dc14a6 100644
--- a/tensorflow/python/tools/freeze_graph_test.py
+++ b/tensorflow/python/tools/freeze_graph_test.py
@@ -60,7 +60,7 @@
# We'll create an input graph that has a single variable containing 1.0,
# and that then multiplies it by 2.
with ops.Graph().as_default():
- variable_node = variables.Variable(1.0, name="variable_node")
+ variable_node = variables.VariableV1(1.0, name="variable_node")
output_node = math_ops.multiply(variable_node, 2.0, name="output_node")
sess = session.Session()
init = variables.global_variables_initializer()
@@ -138,7 +138,7 @@
features = parsing_ops.parse_example(examples, feature_configs)
feature = features[feature_name]
- variable_node = variables.Variable(1.0, name="variable_node")
+ variable_node = variables.VariableV1(1.0, name="variable_node")
scores = math_ops.multiply(variable_node, feature, name="output_node")
class_feature = array_ops.fill(array_ops.shape(feature),
"class_%s" % feature_name)
@@ -174,7 +174,7 @@
output_graph_filename = os.path.join(tmp_dir, "output_graph.pb")
with ops.Graph().as_default():
- variable_node = variables.Variable(1.0, name="variable_node")
+ variable_node = variables.VariableV1(1.0, name="variable_node")
output_node = math_ops.multiply(variable_node, 2.0, name="output_node")
sess = session.Session()
init = variables.global_variables_initializer()
diff --git a/tensorflow/python/tools/saved_model_cli.py b/tensorflow/python/tools/saved_model_cli.py
index 3dbccd1..2fcb0fa 100644
--- a/tensorflow/python/tools/saved_model_cli.py
+++ b/tensorflow/python/tools/saved_model_cli.py
@@ -267,7 +267,8 @@
def run_saved_model_with_feed_dict(saved_model_dir, tag_set, signature_def_key,
input_tensor_key_feed_dict, outdir,
- overwrite_flag, worker=None, tf_debug=False):
+ overwrite_flag, worker=None, init_tpu=False,
+ tf_debug=False):
"""Runs SavedModel and fetch all outputs.
Runs the input dictionary through the MetaGraphDef within a SavedModel
@@ -287,6 +288,8 @@
the same name exists.
worker: If provided, the session will be run on the worker. Valid worker
specification is a bns or gRPC path.
+ init_tpu: If true, the TPU system will be initialized after the session
+ is created.
tf_debug: A boolean flag to use TensorFlow Debugger (TFDBG) to observe the
intermediate Tensor values and runtime GraphDefs while running the
SavedModel.
@@ -328,6 +331,12 @@
]
with session.Session(worker, graph=ops_lib.Graph()) as sess:
+ if init_tpu:
+ print('Initializing TPU System ...')
+ # This is needed for freshly started worker, or if the job
+ # restarts after a preemption.
+ sess.run(tf.contrib.tpu.initialize_system())
+
loader.load(sess, tag_set.split(','), saved_model_dir)
if tf_debug:
@@ -632,7 +641,7 @@
run_saved_model_with_feed_dict(args.dir, args.tag_set, args.signature_def,
tensor_key_feed_dict, args.outdir,
args.overwrite, worker=args.worker,
- tf_debug=args.tf_debug)
+ init_tpu=args.init_tpu, tf_debug=args.tf_debug)
def scan(args):
@@ -775,6 +784,12 @@
default=None,
help='if specified, a Session will be run on the worker. '
'Valid worker specification is a bns or gRPC path.')
+ parser_run.add_argument(
+ '--init_tpu',
+ action='store_true',
+ default=None,
+ help='if specified, tpu.initialize_system will be called on the Session. '
+ 'This option should be only used if the worker is a TPU job.')
parser_run.set_defaults(func=run)
# scan command
diff --git a/tensorflow/python/training/basic_session_run_hooks.py b/tensorflow/python/training/basic_session_run_hooks.py
index 3bd4bd7..1efabcd 100644
--- a/tensorflow/python/training/basic_session_run_hooks.py
+++ b/tensorflow/python/training/basic_session_run_hooks.py
@@ -344,7 +344,7 @@
raise ValueError("steps_per_run should be greater than 0")
self._num_steps = num_steps
self._last_step = last_step
- self._steps_per_run = steps_per_run
+ self._steps_per_run_initial_value = steps_per_run
def begin(self):
self._global_step_tensor = training_util.get_global_step()
@@ -353,7 +353,8 @@
self._steps_per_run_variable = get_or_create_steps_per_run_variable()
def _update_steps_per_run_variable(self, global_step, session):
- steps = min(self._last_step - global_step, self._steps_per_run)
+ steps = min(self._last_step - global_step,
+ self._steps_per_run_initial_value)
self._steps_per_run_variable.load(steps, session=session)
def after_create_session(self, session, coord):
diff --git a/tensorflow/python/training/checkpointable/util.py b/tensorflow/python/training/checkpointable/util.py
index 56c4043..eff15b2 100644
--- a/tensorflow/python/training/checkpointable/util.py
+++ b/tensorflow/python/training/checkpointable/util.py
@@ -247,7 +247,7 @@
def initial_value():
return initializer(
shape_object.as_list(), dtype=dtype, partition_info=partition_info)
- return variables.Variable(
+ return variables.VariableV1(
initial_value=initial_value,
name=name,
dtype=variable_dtype,
diff --git a/tensorflow/python/training/distribute.py b/tensorflow/python/training/distribute.py
index 419a9ec..a92a1bd 100644
--- a/tensorflow/python/training/distribute.py
+++ b/tensorflow/python/training/distribute.py
@@ -26,7 +26,6 @@
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.ops import state_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops.losses import losses_impl
from tensorflow.python.platform import tf_logging
@@ -807,15 +806,22 @@
var: Variable, possibly mirrored to multiple devices, to operate on.
fn: Function to call. Should take the variable as the first argument.
*args: Additional positional arguments to pass to `fn()`.
- **kwargs: Keyword arguments to pass to `fn()`.
+ **kwargs: Keyword arguments to pass to `fn()`. If "grouped=False" is
+ specified, the return value will be unwrapped.
Returns:
- Merged return value of `fn` across all towers.
+ By default, the merged return value of `fn` across all towers. The merged
+ result has dependencies to make sure that if it is evaluated at all, the
+ side effects (updates) will happen on every tower. If instead
+ "grouped=False" is specified, this function will return a nest of lists
+ where each list has an element per tower, and the caller is responsible
+ for ensuring all elements are executed.
"""
_require_cross_tower_context(self)
- return self._update(var, fn, *args, **kwargs)
+ options = {"grouped": kwargs.pop("grouped", True)}
+ return self._update(var, options, fn, *args, **kwargs)
- def _update(self, var, fn, *args, **kwargs):
+ def _update(self, var, options, fn, *args, **kwargs):
raise NotImplementedError("must be implemented in descendants")
def update_non_slot(self, colocate_with, fn, *args, **kwargs):
@@ -825,15 +831,18 @@
colocate_with: The return value of `non_slot_devices()`.
fn: Function to execute.
*args: Positional arguments to pass to `fn()`.
- **kwargs: Keyword arguments to pass to `fn()`.
+ **kwargs: Keyword arguments to pass to `fn()`. If "grouped=False" is
+ specified, the return value will be unwrapped and the caller is
+ responsible for ensuring all elements are executed.
Returns:
Return value of `fn`, possibly merged across devices.
"""
_require_cross_tower_context(self)
- return self._update_non_slot(colocate_with, fn, *args, **kwargs)
+ options = {"grouped": kwargs.pop("grouped", True)}
+ return self._update_non_slot(colocate_with, options, fn, *args, **kwargs)
- def _update_non_slot(self, colocate_with, fn, *args, **kwargs):
+ def _update_non_slot(self, colocate_with, options, fn, *args, **kwargs):
raise NotImplementedError("must be implemented in descendants")
def unwrap(self, value):
@@ -1134,17 +1143,22 @@
del aggregation, destinations
return value
- def _update(self, var, fn, *args, **kwargs):
- # TODO(josh11b): Figure out what we should be passing to UpdateContext()
- # once that value is used for something.
- with ops.colocate_with(var), UpdateContext(var):
- return fn(var, *args, **kwargs)
+ def _update(self, var, options, fn, *args, **kwargs):
+ # The implementations of _update() and _update_non_slot() are identical
+ # except _update() passes `var` as the first argument to `fn()`.
+ return self._update_non_slot(var, options, fn, var, *args, **kwargs)
- def _update_non_slot(self, colocate_with, fn, *args, **kwargs):
+ def _update_non_slot(self, colocate_with, options, fn, *args, **kwargs):
+ should_group = options.pop("grouped")
+ assert not options # Validate that we are processing all of the options.
# TODO(josh11b): Figure out what we should be passing to UpdateContext()
# once that value is used for something.
with ops.colocate_with(colocate_with), UpdateContext(colocate_with):
- return fn(*args, **kwargs)
+ result = fn(*args, **kwargs)
+ if should_group:
+ return result
+ else:
+ return nest.map_structure(self._unwrap, result)
def read_var(self, tower_local_var):
return array_ops.identity(tower_local_var)
@@ -1193,13 +1207,10 @@
def increment_var(v, amount=1):
"""`v += amount`, distributed-aware version."""
def update(vu):
- if isinstance(vu, resource_variable_ops.ResourceVariable):
- return vu.assign_add(amount, read_value=False)
- else:
- return state_ops.assign_add(vu, amount)
+ return vu.assign_add(amount, read_value=False)
def merge_fn(dist, vm):
- return dist.group(dist.update(vm, update))
+ return dist.update(vm, update)
tower_context = distribution_strategy_context.get_tower_context()
return tower_context.merge_call(merge_fn, v)
diff --git a/tensorflow/python/training/distribution_strategy_context.py b/tensorflow/python/training/distribution_strategy_context.py
index 998b5c3..ce580a4 100644
--- a/tensorflow/python/training/distribution_strategy_context.py
+++ b/tensorflow/python/training/distribution_strategy_context.py
@@ -89,6 +89,7 @@
"""Returns the current TowerContext or None if in a cross-tower context.
Note that execution:
+
1. starts in the default (single-tower) tower context (this function
will return the default TowerContext object);
2. switches to cross-tower context (in which case this will return
@@ -121,6 +122,7 @@
"""Returns the current DistributionStrategy if in a cross-tower context.
Note that execution:
+
1. starts in the default (single-tower) tower context;
2. switches to cross-tower context when entering a
`with DistributionStrategy.scope():` block;
diff --git a/tensorflow/python/training/evaluation.py b/tensorflow/python/training/evaluation.py
index b36444a..2c4eb02 100644
--- a/tensorflow/python/training/evaluation.py
+++ b/tensorflow/python/training/evaluation.py
@@ -18,13 +18,14 @@
from __future__ import division
from __future__ import print_function
-import time
import math
+import time
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import init_ops
+from tensorflow.python.ops import math_ops
from tensorflow.python.ops import state_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.platform import tf_logging as logging
@@ -77,6 +78,59 @@
return array_ops.identity(_get_or_create_eval_step().read_value())
+class _MultiStepStopAfterNEvalsHook(session_run_hook.SessionRunHook):
+ """Run hook used by the evaluation routines to run the `eval_ops` N times."""
+
+ def __init__(self, num_evals, steps_per_run=1):
+ """Constructs the run hook.
+
+ Args:
+ num_evals: The number of evaluations to run for. if set to None, will
+ iterate the dataset until all inputs are exhausted.
+ steps_per_run: Number of steps executed per run call.
+ """
+ self._num_evals = num_evals
+ self._evals_completed = None
+ self._steps_per_run_initial_value = steps_per_run
+
+ def _set_evals_completed_tensor(self, updated_eval_step):
+ self._evals_completed = updated_eval_step
+
+ def begin(self):
+ self._steps_per_run_variable = \
+ basic_session_run_hooks.get_or_create_steps_per_run_variable()
+
+ def after_create_session(self, session, coord):
+ # Update number of steps to run in the first run call
+ if self._num_evals is None:
+ steps = self._steps_per_run_initial_value
+ else:
+ steps = min(self._steps_per_run_initial_value, self._num_evals)
+ self._steps_per_run_variable.load(steps, session=session)
+
+ def before_run(self, run_context):
+ return session_run_hook.SessionRunArgs({
+ 'evals_completed': self._evals_completed
+ })
+
+ def after_run(self, run_context, run_values):
+ evals_completed = run_values.results['evals_completed']
+ # Update number of steps to run in the next iteration
+ if self._num_evals is None:
+ steps = self._steps_per_run_initial_value
+ else:
+ steps = min(self._num_evals - evals_completed,
+ self._steps_per_run_initial_value)
+ self._steps_per_run_variable.load(steps, session=run_context.session)
+
+ if self._num_evals is None:
+ logging.info('Evaluation [%d]', evals_completed)
+ else:
+ logging.info('Evaluation [%d/%d]', evals_completed, self._num_evals)
+ if self._num_evals is not None and evals_completed >= self._num_evals:
+ run_context.request_stop()
+
+
class _StopAfterNEvalsHook(session_run_hook.SessionRunHook):
"""Run hook used by the evaluation routines to run the `eval_ops` N times."""
@@ -176,7 +230,15 @@
hooks = list(hooks or [])
if eval_ops is not None:
- update_eval_step = state_ops.assign_add(eval_step, 1, use_locking=True)
+ if any([isinstance(h, _MultiStepStopAfterNEvalsHook) for h in hooks]):
+ steps_per_run_variable = \
+ basic_session_run_hooks.get_or_create_steps_per_run_variable()
+ update_eval_step = state_ops.assign_add(
+ eval_step,
+ math_ops.cast(steps_per_run_variable, dtype=eval_step.dtype),
+ use_locking=True)
+ else:
+ update_eval_step = state_ops.assign_add(eval_step, 1, use_locking=True)
if isinstance(eval_ops, dict):
eval_ops['update_eval_step'] = update_eval_step
@@ -188,7 +250,7 @@
eval_step_value = _get_latest_eval_step_value(eval_ops)
for h in hooks:
- if isinstance(h, _StopAfterNEvalsHook):
+ 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',
diff --git a/tensorflow/python/training/input.py b/tensorflow/python/training/input.py
index 9d9db70..eb131ac 100644
--- a/tensorflow/python/training/input.py
+++ b/tensorflow/python/training/input.py
@@ -56,7 +56,8 @@
# pylint: enable=protected-access
-@tf_export("train.match_filenames_once")
+@tf_export("io.match_filenames_once", "train.match_filenames_once")
+@deprecation.deprecated_endpoints("train.match_filenames_once")
def match_filenames_once(pattern, name=None):
"""Save the list of files matching pattern, so it is only computed once.
diff --git a/tensorflow/python/training/learning_rate_decay_test.py b/tensorflow/python/training/learning_rate_decay_test.py
index 5a92157..03a32f6 100644
--- a/tensorflow/python/training/learning_rate_decay_test.py
+++ b/tensorflow/python/training/learning_rate_decay_test.py
@@ -63,7 +63,7 @@
def testVariables(self):
with self.cached_session():
- step = variables.Variable(1)
+ step = variables.VariableV1(1)
assign_1 = step.assign(1)
assign_2 = step.assign(2)
assign_100 = step.assign(100)
@@ -121,7 +121,7 @@
# Test that ref types are valid.
if not context.executing_eagerly():
- x = variables.Variable(0.0)
+ x = variables.VariableV1(0.0)
x_ref = x.op.outputs[0] # float32_ref tensor should be accepted
boundaries, values = [1.0, 2.0], [1, 2, 3]
learning_rate_decay.piecewise_constant(x_ref, boundaries, values)
diff --git a/tensorflow/python/training/monitored_session_test.py b/tensorflow/python/training/monitored_session_test.py
index 2d7799d..c870d99 100644
--- a/tensorflow/python/training/monitored_session_test.py
+++ b/tensorflow/python/training/monitored_session_test.py
@@ -69,8 +69,8 @@
def test_defaults_empty_graph(self):
with ops.Graph().as_default():
scaffold = monitored_session.Scaffold()
- variables.Variable(1, name='my_var')
- variables.Variable(
+ variables.VariableV1(1, name='my_var')
+ variables.VariableV1(
2, name='my_local_var', collections=[ops.GraphKeys.LOCAL_VARIABLES])
scaffold.finalize()
self.assertTrue(isinstance(scaffold.init_op, ops.Operation))
@@ -105,7 +105,7 @@
def test_caches_values(self):
with ops.Graph().as_default():
- variables.Variable([1])
+ variables.VariableV1([1])
scaffold1 = monitored_session.Scaffold()
scaffold1.finalize()
scaffold2 = monitored_session.Scaffold()
@@ -119,7 +119,7 @@
def test_raise_error_if_more_than_one_cached_item(self):
with ops.Graph().as_default():
- variables.Variable([1])
+ variables.VariableV1([1])
ops.add_to_collection(ops.GraphKeys.SAVERS, saver_lib.Saver())
ops.add_to_collection(ops.GraphKeys.SAVERS, saver_lib.Saver())
with self.assertRaisesRegexp(RuntimeError, 'More than one item'):
@@ -127,7 +127,7 @@
def test_uses_passed_values(self):
with ops.Graph().as_default():
- variables.Variable([1])
+ variables.VariableV1([1])
saver = saver_lib.Saver()
scaffold = monitored_session.Scaffold(
init_op=2,
@@ -148,7 +148,7 @@
def test_graph_is_finalized(self):
with ops.Graph().as_default():
- variables.Variable([1])
+ variables.VariableV1([1])
monitored_session.Scaffold().finalize()
with self.assertRaisesRegexp(RuntimeError,
'Graph is finalized and cannot be modified'):
@@ -157,7 +157,7 @@
def test_new_scaffold_from_default_scaffold(self):
scaffold1 = monitored_session.Scaffold()
with ops.Graph().as_default():
- variables.Variable([1])
+ variables.VariableV1([1])
saver = saver_lib.Saver()
scaffold2 = monitored_session.Scaffold(
init_op=2,
@@ -180,7 +180,7 @@
def test_new_scaffold_from_existing_scaffold(self):
with ops.Graph().as_default():
- variables.Variable([1])
+ variables.VariableV1([1])
saver = saver_lib.Saver()
scaffold1 = monitored_session.Scaffold(
init_op=2,
@@ -1374,7 +1374,7 @@
def test_defaults(self):
with ops.Graph().as_default():
- a_var = variables.Variable(0)
+ a_var = variables.VariableV1(0)
with monitored_session.MonitoredSession() as session:
self.assertEqual(0, session.run(a_var))
@@ -1700,7 +1700,7 @@
def test_graph_finalized_during_run_unfinalized_after_exit(self):
with ops.Graph().as_default() as g:
- a_var = variables.Variable(0)
+ a_var = variables.VariableV1(0)
with monitored_session.MonitoredSession() as session:
self.assertEqual(0, session.run(a_var))
self.assertTrue(g.finalized)
@@ -1708,7 +1708,7 @@
def test_keep_finalized_graph_as_finalized(self):
with ops.Graph().as_default() as g:
- a_var = variables.Variable(0)
+ a_var = variables.VariableV1(0)
monitored_session.Scaffold().finalize()
with monitored_session.MonitoredSession() as session:
self.assertEqual(0, session.run(a_var))
@@ -2032,7 +2032,7 @@
with ops.Graph().as_default():
c = array_ops.placeholder(dtypes.float32)
v = array_ops.identity(c)
- graph_state = variables.Variable(0.0)
+ graph_state = variables.VariableV1(0.0)
graph_side_effect = state_ops.assign_add(graph_state, 0.31)
def step_fn(step_context):
@@ -2088,7 +2088,7 @@
c = array_ops.placeholder(dtypes.float32)
v = array_ops.identity(c)
vv = constant_op.constant(3.2)
- graph_state = variables.Variable(0.0)
+ graph_state = variables.VariableV1(0.0)
graph_side_effect = state_ops.assign_add(graph_state, 0.31)
class Hook(session_run_hook.SessionRunHook):
@@ -2125,7 +2125,7 @@
def test_handles_initialization(self):
with ops.Graph().as_default():
- a_var = variables.Variable(0)
+ a_var = variables.VariableV1(0)
with monitored_session.SingularMonitoredSession() as session:
# If it's not initialized, following statement raises an error.
self.assertEqual(0, session.run(a_var))
diff --git a/tensorflow/python/training/moving_averages.py b/tensorflow/python/training/moving_averages.py
index 177a7dd..041266d 100644
--- a/tensorflow/python/training/moving_averages.py
+++ b/tensorflow/python/training/moving_averages.py
@@ -372,13 +372,13 @@
Args:
var_list: A list of Variable or Tensor objects. The variables
- and Tensors must be of types float16, float32, or float64.
+ and Tensors must be of types bfloat16, float16, float32, or float64.
Returns:
An Operation that updates the moving averages.
Raises:
- TypeError: If the arguments are not all float16, float32, or float64.
+ TypeError: If the arguments are not an allowed type.
ValueError: If the moving average of one of the variables is already
being computed.
"""
@@ -387,8 +387,9 @@
var_list = variables.trainable_variables()
zero_debias_true = set() # set of vars to set `zero_debias=True`
for var in var_list:
- if var.dtype.base_dtype not in [dtypes.float16, dtypes.float32,
- dtypes.float64]:
+ if var.dtype.base_dtype not in [
+ dtypes.bfloat16, dtypes.float16, dtypes.float32, dtypes.float64
+ ]:
raise TypeError("The variables must be half, float, or double: %s" %
var.name)
diff --git a/tensorflow/python/training/moving_averages_test.py b/tensorflow/python/training/moving_averages_test.py
index 93991d0..bb2fca6 100644
--- a/tensorflow/python/training/moving_averages_test.py
+++ b/tensorflow/python/training/moving_averages_test.py
@@ -18,6 +18,7 @@
from __future__ import division
from __future__ import print_function
+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
@@ -110,6 +111,32 @@
denominator_2 = denominator_1 * decay + weight_2 * (1.0 - decay)
self.assertAllClose(numerator_2 / denominator_2, wma_array)
+ def testWeightedMovingAverageBfloat16(self):
+ bfloat16 = pywrap_tensorflow.TF_bfloat16_type()
+ with self.cached_session() as sess:
+ decay = 0.5
+ weight = array_ops.placeholder(dtypes.bfloat16, [])
+ val = array_ops.placeholder(dtypes.bfloat16, [])
+
+ wma = moving_averages.weighted_moving_average(val, decay, weight)
+ variables.global_variables_initializer().run()
+
+ # Get the first weighted moving average.
+ val_1 = 3.0
+ weight_1 = 4.0
+ wma_array = sess.run(wma, feed_dict={val: val_1, weight: weight_1})
+ numerator_1 = val_1 * weight_1 * (1.0 - decay)
+ denominator_1 = weight_1 * (1.0 - decay)
+ self.assertAllClose(numerator_1 / denominator_1, wma_array)
+
+ # Get the second weighted moving average.
+ val_2 = 11.0
+ weight_2 = 22.0
+ wma_array = sess.run(wma, feed_dict={val: val_2, weight: weight_2})
+ numerator_2 = numerator_1 * decay + val_2 * weight_2 * (1.0 - decay)
+ denominator_2 = denominator_1 * decay + weight_2 * (1.0 - decay)
+ self.assertAllClose(bfloat16(numerator_2 / denominator_2), wma_array)
+
def _Repeat(value, dim):
if dim == 1:
diff --git a/tensorflow/python/training/optimizer.py b/tensorflow/python/training/optimizer.py
index 699162b..4703491 100644
--- a/tensorflow/python/training/optimizer.py
+++ b/tensorflow/python/training/optimizer.py
@@ -471,7 +471,10 @@
if var_list is None:
var_list = tape.watched_variables()
- grads = tape.gradient(loss_value, var_list, grad_loss)
+ # TODO(jhseu): Figure out why GradientTape's gradients don't require loss
+ # to be executed.
+ with ops.control_dependencies([loss_value]):
+ grads = tape.gradient(loss_value, var_list, grad_loss)
return list(zip(grads, var_list))
# Non-callable/Tensor loss case
@@ -585,7 +588,7 @@
var_list = [v for g, v, _ in converted_grads_and_vars if g is not None]
if not var_list:
raise ValueError("No gradients provided for any variable: %s." %
- ([str(v) for _, _, v in converted_grads_and_vars],))
+ ([str(v) for _, v, _ in converted_grads_and_vars],))
with ops.init_scope():
self._create_slots(var_list)
update_ops = []
@@ -689,7 +692,7 @@
update_ops = [
op
for grad, var in grads_and_vars
- for op in distribution.unwrap(distribution.update(var, update, grad))
+ for op in distribution.update(var, update, grad, grouped=False)
]
def finish(self, update_ops):
@@ -697,13 +700,13 @@
non_slot_devices = distribution.non_slot_devices(var_list)
finish_updates = distribution.update_non_slot(
- non_slot_devices, finish, self, update_ops)
+ non_slot_devices, finish, self, update_ops, grouped=False)
if global_step is None:
apply_updates = distribution.group(finish_updates, name=name)
else:
- with ops.control_dependencies(distribution.unwrap(finish_updates)):
- apply_updates = distribution.group(distribution.update(
- global_step, state_ops.assign_add, 1, name=name))
+ with ops.control_dependencies(finish_updates):
+ apply_updates = distribution.update(
+ global_step, state_ops.assign_add, 1, name=name)
if not context.executing_eagerly():
if isinstance(apply_updates, ops.Tensor):
diff --git a/tensorflow/python/training/quantize_training_test.py b/tensorflow/python/training/quantize_training_test.py
index 9754ade..6edbf76 100644
--- a/tensorflow/python/training/quantize_training_test.py
+++ b/tensorflow/python/training/quantize_training_test.py
@@ -58,7 +58,8 @@
g = ops.Graph()
with session.Session(graph=g) as sess:
a = constant_op.constant(6.0, shape=[1, 1], name='a')
- b = variables.Variable(constant_op.constant(7.0, shape=[1, 1]), name='b')
+ b = variables.VariableV1(
+ constant_op.constant(7.0, shape=[1, 1]), name='b')
c = math_ops.matmul(a, b, name='matmul')
init_op = variables.global_variables_initializer()
diff --git a/tensorflow/python/training/queue_runner_test.py b/tensorflow/python/training/queue_runner_test.py
index 9b9e28a..15fe42b 100644
--- a/tensorflow/python/training/queue_runner_test.py
+++ b/tensorflow/python/training/queue_runner_test.py
@@ -44,7 +44,7 @@
with self.cached_session() as sess:
# CountUpTo will raise OUT_OF_RANGE when it reaches the count.
zero64 = constant_op.constant(0, dtype=dtypes.int64)
- var = variables.Variable(zero64)
+ 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()
@@ -64,9 +64,9 @@
with self.cached_session() as sess:
# CountUpTo will raise OUT_OF_RANGE when it reaches the count.
zero64 = constant_op.constant(0, dtype=dtypes.int64)
- var0 = variables.Variable(zero64)
+ var0 = variables.VariableV1(zero64)
count_up_to_3 = var0.count_up_to(3)
- var1 = variables.Variable(zero64)
+ var1 = variables.VariableV1(zero64)
count_up_to_30 = var1.count_up_to(30)
queue = data_flow_ops.FIFOQueue(10, dtypes.float32)
qr = queue_runner_impl.QueueRunner(queue, [count_up_to_3, count_up_to_30])
@@ -131,7 +131,7 @@
with self.cached_session() as sess:
# CountUpTo will raise OUT_OF_RANGE when it reaches the count.
zero64 = constant_op.constant(0, dtype=dtypes.int64)
- var = variables.Variable(zero64)
+ 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()
@@ -184,7 +184,7 @@
with self.cached_session() as sess:
with session.Session() as other_sess:
zero64 = constant_op.constant(0, dtype=dtypes.int64)
- var = variables.Variable(zero64)
+ 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()
@@ -199,7 +199,7 @@
with self.cached_session() as sess:
# CountUpTo will raise OUT_OF_RANGE when it reaches the count.
zero64 = constant_op.constant(0, dtype=dtypes.int64)
- var = variables.Variable(zero64)
+ 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()
@@ -215,7 +215,7 @@
with self.cached_session() as sess:
# CountUpTo will raise OUT_OF_RANGE when it reaches the count.
zero64 = constant_op.constant(0, dtype=dtypes.int64)
- var = variables.Variable(zero64)
+ 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()
@@ -250,7 +250,7 @@
def testStartQueueRunners(self):
# CountUpTo will raise OUT_OF_RANGE when it reaches the count.
zero64 = constant_op.constant(0, dtype=dtypes.int64)
- var = variables.Variable(zero64)
+ var = variables.VariableV1(zero64)
count_up_to = var.count_up_to(3)
queue = data_flow_ops.FIFOQueue(10, dtypes.float32)
init_op = variables.global_variables_initializer()
@@ -267,7 +267,7 @@
def testStartQueueRunnersRaisesIfNotASession(self):
zero64 = constant_op.constant(0, dtype=dtypes.int64)
- var = variables.Variable(zero64)
+ var = variables.VariableV1(zero64)
count_up_to = var.count_up_to(3)
queue = data_flow_ops.FIFOQueue(10, dtypes.float32)
init_op = variables.global_variables_initializer()
@@ -280,7 +280,7 @@
def testStartQueueRunnersIgnoresMonitoredSession(self):
zero64 = constant_op.constant(0, dtype=dtypes.int64)
- var = variables.Variable(zero64)
+ var = variables.VariableV1(zero64)
count_up_to = var.count_up_to(3)
queue = data_flow_ops.FIFOQueue(10, dtypes.float32)
init_op = variables.global_variables_initializer()
@@ -297,7 +297,7 @@
graph = ops.Graph()
with graph.as_default():
zero64 = constant_op.constant(0, dtype=dtypes.int64)
- var = variables.Variable(zero64)
+ var = variables.VariableV1(zero64)
count_up_to = var.count_up_to(3)
queue = data_flow_ops.FIFOQueue(10, dtypes.float32)
init_op = variables.global_variables_initializer()
diff --git a/tensorflow/python/training/saver_test.py b/tensorflow/python/training/saver_test.py
index 69b1055..49e6e65 100644
--- a/tensorflow/python/training/saver_test.py
+++ b/tensorflow/python/training/saver_test.py
@@ -311,8 +311,8 @@
# Build a graph with 2 parameter nodes, and Save and
# Restore nodes for them.
- v0 = variables.Variable(10.0, name="v0")
- v1 = variables.Variable(20.0, name="v1")
+ v0 = variables.VariableV1(10.0, name="v0")
+ v1 = variables.VariableV1(20.0, name="v1")
v2 = saver_test_utils.CheckpointedOp(name="v2")
v2_init = v2.insert("k1", 30.0)
save = saver_module.Saver(
@@ -350,8 +350,8 @@
# Start a second session. In that session the parameter nodes
# have not been initialized either.
with self.cached_session() as sess:
- v0 = variables.Variable(-1.0, name="v0")
- v1 = variables.Variable(-1.0, name="v1")
+ v0 = variables.VariableV1(-1.0, name="v0")
+ v1 = variables.VariableV1(-1.0, name="v1")
v2 = saver_test_utils.CheckpointedOp(name="v2")
save = saver_module.Saver({"v0": v0, "v1": v1, "v2": v2.saveable})
@@ -370,7 +370,7 @@
self.assertEqual(30.0, v2.values().eval())
def testFilenameTensor(self):
- v0 = variables.Variable(0, name="v0")
+ v0 = variables.VariableV1(0, name="v0")
filename = b"somerandomfilename"
save = saver_module.Saver({"v0": v0}, filename=filename)
with self.cached_session() as sess:
@@ -379,7 +379,7 @@
self.assertEqual(sess.run(tensor), filename)
def testInvalidPath(self):
- v0 = variables.Variable(0, name="v0")
+ v0 = variables.VariableV1(0, name="v0")
for ver in (saver_pb2.SaverDef.V1, saver_pb2.SaverDef.V2):
with self.cached_session() as sess:
save = saver_module.Saver({"v0": v0}, write_version=ver)
@@ -392,7 +392,7 @@
with self.cached_session() as sess:
# Build a graph with 1 node, and save and restore for them.
- v = variables.Variable(np.int64(15), name="v")
+ v = variables.VariableV1(np.int64(15), name="v")
save = saver_module.Saver({"v": v}, restore_sequentially=True)
variables.global_variables_initializer().run()
@@ -402,7 +402,7 @@
self.assertEqual(save_path, val)
with self.cached_session() as sess:
- v = variables.Variable(np.int64(-1), name="v")
+ v = variables.VariableV1(np.int64(-1), name="v")
save = saver_module.Saver({"v": v})
with self.assertRaisesWithPredicateMatch(
@@ -416,9 +416,9 @@
def testSomeErrors(self):
with ops_lib.Graph().as_default():
- v0 = variables.Variable([10.0], name="v0")
- v1 = variables.Variable([20.0], name="v1")
- v2 = variables.Variable([20.0], name="v2")
+ v0 = variables.VariableV1([10.0], name="v0")
+ v1 = variables.VariableV1([20.0], name="v1")
+ v2 = variables.VariableV1([20.0], name="v2")
v2._set_save_slice_info(
variables.Variable.SaveSliceInfo("v1", [1], [0], [1]))
@@ -446,7 +446,7 @@
def testSameName(self):
with ops_lib.Graph().as_default():
- v0 = variables.Variable([10.0], name="v0")
+ v0 = variables.VariableV1([10.0], name="v0")
v2 = saver_test_utils.CheckpointedOp(name="v2")
# Saving one variable under two names raises an error.
@@ -468,8 +468,8 @@
with self.session(graph=ops_lib.Graph()) as sess:
# Build a graph with 2 parameter nodes, and Save and
# Restore nodes for them.
- v0 = variables.Variable(10.0, name="v0")
- v1 = variables.Variable(20.0, name="v1")
+ v0 = variables.VariableV1(10.0, name="v0")
+ v1 = variables.VariableV1(20.0, name="v1")
v2 = saver_test_utils.CheckpointedOp(name="v2")
v2_init = v2.insert("k1", 30.0)
save = saver_module.Saver([v0, v1, v2.saveable])
@@ -490,8 +490,8 @@
# Start a second session. In that session the variables
# have not been initialized either.
with self.session(graph=ops_lib.Graph()) as sess:
- v0 = variables.Variable(-1.0, name="v0")
- v1 = variables.Variable(-1.0, name="v1")
+ v0 = variables.VariableV1(-1.0, name="v0")
+ v1 = variables.VariableV1(-1.0, name="v1")
v2 = saver_test_utils.CheckpointedOp(name="v2")
save = saver_module.Saver([v0, v1, v2.saveable])
@@ -515,8 +515,8 @@
# Build another graph with 2 nodes, initialized
# differently, and a Restore node for them.
with self.session(graph=ops_lib.Graph()) as sess:
- v0_2 = variables.Variable(1000.0, name="v0")
- v1_2 = variables.Variable(2000.0, name="v1")
+ v0_2 = variables.VariableV1(1000.0, name="v0")
+ v1_2 = variables.VariableV1(2000.0, name="v1")
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()
@@ -574,14 +574,14 @@
save_path = os.path.join(self.get_temp_dir(), "gpu")
with session.Session("", graph=ops_lib.Graph()) as sess:
with sess.graph.device(test.gpu_device_name()):
- v0_1 = variables.Variable(123.45)
+ v0_1 = variables.VariableV1(123.45)
save = saver_module.Saver({"v0": v0_1})
variables.global_variables_initializer().run()
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.Variable(543.21)
+ v0_2 = variables.VariableV1(543.21)
save = saver_module.Saver({"v0": v0_2})
variables.global_variables_initializer().run()
@@ -591,22 +591,22 @@
save_path = os.path.join(self.get_temp_dir(), "gpu")
with session.Session("", graph=ops_lib.Graph()) as sess:
with sess.graph.device(test.gpu_device_name()):
- v0_1 = variables.Variable(123.45)
+ v0_1 = variables.VariableV1(123.45)
save = saver_module.Saver({"v0": v0_1}, sharded=True, allow_empty=True)
variables.global_variables_initializer().run()
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.Variable(543.21)
+ v0_2 = variables.VariableV1(543.21)
save = saver_module.Saver({"v0": v0_2}, sharded=True, allow_empty=True)
variables.global_variables_initializer().run()
def testVariables(self):
save_path = os.path.join(self.get_temp_dir(), "variables")
with session.Session("", graph=ops_lib.Graph()) as sess:
- one = variables.Variable(1.0)
- twos = variables.Variable([2.0, 2.0, 2.0])
+ one = variables.VariableV1(1.0)
+ twos = variables.VariableV1([2.0, 2.0, 2.0])
v2 = saver_test_utils.CheckpointedOp(name="v2")
init = variables.global_variables_initializer()
save = saver_module.Saver()
@@ -615,8 +615,8 @@
save.save(sess, save_path)
with session.Session("", graph=ops_lib.Graph()) as sess:
- one = variables.Variable(0.0)
- twos = variables.Variable([0.0, 0.0, 0.0])
+ one = variables.VariableV1(0.0)
+ twos = variables.VariableV1([0.0, 0.0, 0.0])
v2 = saver_test_utils.CheckpointedOp(name="v2")
# Saver with no arg, defaults to 'all variables'.
save = saver_module.Saver()
@@ -628,14 +628,14 @@
def testVarListShouldBeEmptyInDeferredBuild(self):
with ops_lib.Graph().as_default():
- v = variables.Variable(1.0)
+ v = variables.VariableV1(1.0)
with self.assertRaisesRegexp(ValueError, "defer_build"):
saver_module.Saver([v], defer_build=True)
def testBuildShouldBeCalledBeforeSaveInCaseOfDeferBuild(self):
save_path = os.path.join(self.get_temp_dir(), "error_deferred_build")
with ops_lib.Graph().as_default(), session.Session() as sess:
- variables.Variable(1.0)
+ variables.VariableV1(1.0)
saver = saver_module.Saver(defer_build=True)
with self.assertRaisesRegexp(RuntimeError, "build"):
saver.save(sess, save_path)
@@ -643,18 +643,18 @@
def testDeferredBuild(self):
save_path = os.path.join(self.get_temp_dir(), "deferred_build")
with session.Session("", graph=ops_lib.Graph()) as sess:
- one = variables.Variable(1.0)
+ one = variables.VariableV1(1.0)
save = saver_module.Saver(defer_build=True)
# if build is not deferred, saver cannot save the `twos`.
- twos = variables.Variable([2.0, 2.0, 2.0])
+ twos = variables.VariableV1([2.0, 2.0, 2.0])
init = variables.global_variables_initializer()
save.build()
init.run()
save.save(sess, save_path)
with session.Session("", graph=ops_lib.Graph()) as sess:
- one = variables.Variable(0.0)
- twos = variables.Variable([0.0, 0.0, 0.0])
+ one = variables.VariableV1(0.0)
+ twos = variables.VariableV1([0.0, 0.0, 0.0])
# Saver with no arg, defaults to 'all variables'.
save = saver_module.Saver()
save.restore(sess, save_path)
@@ -664,7 +664,7 @@
def testReshape(self):
save_path = os.path.join(self.get_temp_dir(), "variables_reshape")
with session.Session("", graph=ops_lib.Graph()) as sess:
- var = variables.Variable([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+ var = variables.VariableV1([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
init = variables.global_variables_initializer()
save = saver_module.Saver()
init.run()
@@ -672,7 +672,7 @@
# Error when restoring with default reshape=False
with session.Session("", graph=ops_lib.Graph()) as sess:
- var = variables.Variable([[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]])
+ var = variables.VariableV1([[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]])
save = saver_module.Saver()
with self.assertRaisesRegexp(
errors_impl.InvalidArgumentError,
@@ -681,7 +681,7 @@
# Restored to new shape with reshape=True
with session.Session("", graph=ops_lib.Graph()) as sess:
- var = variables.Variable([[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]])
+ 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())
@@ -731,8 +731,8 @@
for save_path in paths:
# Build a graph with 2 parameter nodes, and Save and
# Restore nodes for them.
- v0 = variables.Variable(10.0, name="v0")
- v1 = variables.Variable(20.0, name="v1")
+ v0 = variables.VariableV1(10.0, name="v0")
+ v1 = variables.VariableV1(20.0, name="v1")
save = saver_module.Saver({"v0": v0, "v1": v1}, restore_sequentially=True)
init_all_op = variables.global_variables_initializer()
@@ -770,8 +770,8 @@
# Build a graph with 2 parameter nodes, and Save and
# Restore nodes for them.
- v0 = variables.Variable(10.0, name="v0")
- v1 = variables.Variable(20.0, name="v1")
+ v0 = variables.VariableV1(10.0, name="v0")
+ v1 = variables.VariableV1(20.0, name="v1")
save = saver_module.Saver({"v0": v0, "v1": v1}, restore_sequentially=True)
init_all_op = variables.global_variables_initializer()
@@ -859,10 +859,10 @@
target="",
config=config_pb2.ConfigProto(device_count={"CPU": 2})) as sess:
with sess.graph.device("/cpu:0"):
- v0 = variables.Variable(10, name="v0")
+ v0 = variables.VariableV1(10, name="v0")
t0 = saver_test_utils.CheckpointedOp(name="t0")
with sess.graph.device("/cpu:1"):
- v1 = variables.Variable(20, name="v1")
+ v1 = variables.VariableV1(20, name="v1")
t1 = saver_test_utils.CheckpointedOp(name="t1")
save = saver_module.Saver(
{
@@ -890,7 +890,7 @@
target="",
config=config_pb2.ConfigProto(device_count={"CPU": 2})) as sess:
with sess.graph.device("/cpu:0"):
- v0 = variables.Variable(111, name="v0")
+ v0 = variables.VariableV1(111, name="v0")
t0 = saver_test_utils.CheckpointedOp(name="t0")
save = saver_module.Saver(
{
@@ -914,7 +914,7 @@
target="",
config=config_pb2.ConfigProto(device_count={"CPU": 2})) as sess:
with sess.graph.device("/cpu:0"):
- v1 = variables.Variable(222)
+ v1 = variables.VariableV1(222)
t1 = saver_test_utils.CheckpointedOp(name="t1")
save = saver_module.Saver(
{
@@ -938,10 +938,10 @@
target="",
config=config_pb2.ConfigProto(device_count={"CPU": 2})) as sess:
with sess.graph.device("/cpu:0"):
- v0 = variables.Variable(111, name="v0")
+ v0 = variables.VariableV1(111, name="v0")
t0 = saver_test_utils.CheckpointedOp(name="t0")
with sess.graph.device("/cpu:1"):
- v1 = variables.Variable(222, name="v1")
+ v1 = variables.VariableV1(222, name="v1")
t1 = saver_test_utils.CheckpointedOp(name="t1")
save = saver_module.Saver(
{
@@ -984,7 +984,7 @@
def testSaverDef(self):
with self.cached_session():
- v0 = variables.Variable(123, name="v0")
+ v0 = variables.VariableV1(123, name="v0")
save = saver_module.Saver({"v0": v0}, sharded=True)
sd = save.as_saver_def()
self.assertTrue(sd.sharded)
@@ -1023,7 +1023,7 @@
if use_resource:
vs = [resource_variable_ops.ResourceVariable(rnd, name=var_name)]
else:
- vs = [variables.Variable(rnd, name=var_name)]
+ vs = [variables.VariableV1(rnd, name=var_name)]
variables.global_variables_initializer().run()
if call_saver_with_dict:
@@ -1054,7 +1054,7 @@
]
else:
new_vs = [
- variables.Variable(
+ variables.VariableV1(
array_ops.zeros(
shape=var_full_shape), # != original contents.
name=var_name)
@@ -1210,7 +1210,7 @@
save_dir = self._get_test_dir("max_to_keep_non_sharded")
with self.cached_session() as sess:
- v = variables.Variable(10.0, name="v")
+ v = variables.VariableV1(10.0, name="v")
save = saver_module.Saver({"v": v}, max_to_keep=2)
variables.global_variables_initializer().run()
self.assertEqual([], save.last_checkpoints)
@@ -1389,9 +1389,9 @@
target="",
config=config_pb2.ConfigProto(device_count={"CPU": 2})) as sess:
with sess.graph.device("/cpu:0"):
- v0 = variables.Variable(111, name="v0")
+ v0 = variables.VariableV1(111, name="v0")
with sess.graph.device("/cpu:1"):
- v1 = variables.Variable(222, name="v1")
+ v1 = variables.VariableV1(222, name="v1")
save = saver_module.Saver(
{
"v0": v0,
@@ -1448,7 +1448,7 @@
save_dir2 = self._get_test_dir("max_to_keep_0")
with self.cached_session() as sess:
- v = variables.Variable(10.0, name="v")
+ v = variables.VariableV1(10.0, name="v")
variables.global_variables_initializer().run()
# Test max_to_keep being None.
@@ -1475,7 +1475,7 @@
save_dir = self._get_test_dir("no_meta_graph")
with self.cached_session() as sess:
- v = variables.Variable(10.0, name="v")
+ v = variables.VariableV1(10.0, name="v")
save = saver_module.Saver({"v": v})
variables.global_variables_initializer().run()
@@ -1632,13 +1632,13 @@
filename = os.path.join(test_dir, "metafile")
with self.cached_session():
# Creates a graph.
- v0 = variables.Variable(1.0, name="v0")
+ v0 = variables.VariableV1(1.0, name="v0")
control_flow_ops.cond(
math_ops.less(v0, 10), lambda: math_ops.add(v0, 1),
lambda: math_ops.subtract(v0, 1))
control_flow_ops.while_loop(lambda i: math_ops.less(i, 10),
lambda i: math_ops.add(i, 1), [v0])
- var = variables.Variable(constant_op.constant(0, dtype=dtypes.int64))
+ var = variables.VariableV1(constant_op.constant(0, dtype=dtypes.int64))
count_up_to = var.count_up_to(3)
input_queue = data_flow_ops.FIFOQueue(
30, dtypes.float32, shared_name="collection_queue")
@@ -1687,7 +1687,7 @@
def testAddCollectionDefFails(self):
with self.cached_session():
# Creates a graph.
- v0 = variables.Variable(10.0, name="v0")
+ v0 = variables.VariableV1(10.0, name="v0")
# Creates a saver.
save = saver_module.Saver({"v0": v0})
# Generates MetaGraphDef.
@@ -1711,8 +1711,8 @@
saver1_ckpt = os.path.join(test_dir, "saver1.ckpt")
with self.session(graph=ops_lib.Graph()) as sess:
# Creates a graph.
- v0 = variables.Variable([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], name="v0")
- v1 = variables.Variable(11.0, name="v1")
+ v0 = variables.VariableV1([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], name="v0")
+ v1 = variables.VariableV1(11.0, name="v1")
# Creates 2 savers.
saver0 = saver_module.Saver({"v0": v0}, name="saver0")
saver1 = saver_module.Saver({"v1": v1}, name="saver1")
@@ -1788,8 +1788,8 @@
saver1_ckpt = os.path.join(test_dir, "saver1.ckpt")
with self.session(graph=ops_lib.Graph()) as sess:
# Creates a graph.
- v0 = variables.Variable([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], name="v0")
- v1 = variables.Variable(11.0, name="v1")
+ v0 = variables.VariableV1([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], name="v0")
+ v1 = variables.VariableV1(11.0, name="v1")
# Creates 2 savers.
saver0 = saver_module.Saver({"v0": v0}, name="saver0")
@@ -1840,7 +1840,7 @@
filename = os.path.join(test_dir, "metafile")
with self.session(graph=ops_lib.Graph()):
# Creates a graph.
- variables.Variable(10.0, name="v0")
+ variables.VariableV1(10.0, name="v0")
# Exports the graph as binary format.
saver_module.export_meta_graph(filename, as_text=False)
with self.session(graph=ops_lib.Graph()):
@@ -1871,8 +1871,8 @@
test_dir = self._get_test_dir("slice_saver")
filename = os.path.join(test_dir, "metafile")
with self.cached_session():
- v1 = variables.Variable([20.0], name="v1")
- v2 = variables.Variable([20.0], name="v2")
+ v1 = variables.VariableV1([20.0], name="v1")
+ v2 = variables.VariableV1([20.0], name="v2")
v2._set_save_slice_info(
variables.Variable.SaveSliceInfo("v1", [1], [0], [1]))
@@ -1899,7 +1899,7 @@
# Hidden 1
images = constant_op.constant(1.2, dtypes.float32, shape=[100, 28])
with ops_lib.name_scope("hidden1"):
- weights = variables.Variable(
+ weights = variables.VariableV1(
random_ops.truncated_normal(
[28, 128], stddev=1.0 / math.sqrt(float(28))),
name="weights")
@@ -1907,7 +1907,7 @@
# the save and restore of control flow context (which doesn't make any
# sense here from a machine learning perspective). The typical biases is
# a simple Variable without the conditions.
- biases = variables.Variable(
+ biases = variables.VariableV1(
control_flow_ops.cond(
math_ops.less(random.random(), 0.5),
lambda: array_ops.ones([128]), lambda: array_ops.zeros([128])),
@@ -1915,7 +1915,7 @@
hidden1 = nn_ops.relu(math_ops.matmul(images, weights) + biases)
# Hidden 2
with ops_lib.name_scope("hidden2"):
- weights = variables.Variable(
+ weights = variables.VariableV1(
random_ops.truncated_normal(
[128, 32], stddev=1.0 / math.sqrt(float(128))),
name="weights")
@@ -1933,15 +1933,16 @@
_, biases = control_flow_ops.while_loop(
loop_cond, loop_body,
- [constant_op.constant(0), variables.Variable(array_ops.zeros([32]))])
+ [constant_op.constant(0),
+ variables.VariableV1(array_ops.zeros([32]))])
hidden2 = nn_ops.relu(math_ops.matmul(hidden1, weights) + biases)
# Linear
with ops_lib.name_scope("softmax_linear"):
- weights = variables.Variable(
+ weights = variables.VariableV1(
random_ops.truncated_normal(
[32, 10], stddev=1.0 / math.sqrt(float(32))),
name="weights")
- biases = variables.Variable(array_ops.zeros([10]), name="biases")
+ biases = variables.VariableV1(array_ops.zeros([10]), name="biases")
logits = math_ops.matmul(hidden2, weights) + biases
ops_lib.add_to_collection("logits", logits)
init_all_op = variables.global_variables_initializer()
@@ -2028,7 +2029,7 @@
# Create while loop using `outer_body_fn`.
with ops_lib.Graph().as_default():
- var = variables.Variable(0.0)
+ var = variables.VariableV1(0.0)
var_name = var.name
output = graph_fn(var)
output_name = output.name
@@ -2122,8 +2123,8 @@
def testStrippedOpListDef(self):
with self.cached_session():
# Creates a graph.
- v0 = variables.Variable(0.0)
- var = variables.Variable(10.0)
+ v0 = variables.VariableV1(0.0)
+ var = variables.VariableV1(10.0)
math_ops.add(v0, var)
@function.Defun(dtypes.float32)
@@ -2161,8 +2162,8 @@
# With strip_default_attrs enabled, attributes "T" (float32) and "Tout"
# (complex64) in the "Complex" op must be removed.
with self.cached_session():
- real_num = variables.Variable(1.0, dtype=dtypes.float32, name="real")
- imag_num = variables.Variable(2.0, dtype=dtypes.float32, name="imag")
+ real_num = variables.VariableV1(1.0, dtype=dtypes.float32, name="real")
+ imag_num = variables.VariableV1(2.0, dtype=dtypes.float32, name="imag")
math_ops.complex(real_num, imag_num, name="complex")
save = saver_module.Saver({"real_num": real_num, "imag_num": imag_num})
@@ -2178,8 +2179,8 @@
# (complex64) in the "Complex" op must *not* be removed, even if they map
# to their defaults.
with self.session(graph=ops_lib.Graph()):
- real_num = variables.Variable(1.0, dtype=dtypes.float32, name="real")
- imag_num = variables.Variable(2.0, dtype=dtypes.float32, name="imag")
+ real_num = variables.VariableV1(1.0, dtype=dtypes.float32, name="real")
+ imag_num = variables.VariableV1(2.0, dtype=dtypes.float32, name="imag")
math_ops.complex(real_num, imag_num, name="complex")
save = saver_module.Saver({"real_num": real_num, "imag_num": imag_num})
@@ -2198,9 +2199,9 @@
image = array_ops.placeholder(dtypes.float32, [None, 784], name="image")
label = array_ops.placeholder(dtypes.float32, [None, 10], name="label")
with session.Session() as sess:
- weights = variables.Variable(
+ weights = variables.VariableV1(
random_ops.random_uniform([784, 10]), name="weights")
- bias = variables.Variable(array_ops.zeros([10]), name="bias")
+ bias = variables.VariableV1(array_ops.zeros([10]), name="bias")
logit = nn_ops.relu(math_ops.matmul(image, weights) + bias, name="logits")
nn_ops.softmax(logit, name="prediction")
cost = nn_ops.softmax_cross_entropy_with_logits(labels=label,
@@ -2243,7 +2244,7 @@
self.assertIsNone(new_saver_1)
# Create a variable in graph_2 under scope "my_scope".
- variables.Variable(array_ops.zeros([10]), name="my_scope/my_var")
+ variables.VariableV1(array_ops.zeros([10]), name="my_scope/my_var")
sess.run(variables.global_variables_initializer())
# Restore the checkpoint into a different scope "subgraph_2".
new_saver_2 = saver_module.import_meta_graph(
@@ -2268,9 +2269,9 @@
image = array_ops.placeholder(dtypes.float32, [None, 784], name="image")
label = array_ops.placeholder(dtypes.float32, [None, 10], name="label")
with session.Session() as sess:
- weights = variables.Variable(
+ weights = variables.VariableV1(
random_ops.random_uniform([784, 10]), name="weights")
- bias = variables.Variable(array_ops.zeros([10]), name="bias")
+ bias = variables.VariableV1(array_ops.zeros([10]), name="bias")
logit = nn_ops.relu(math_ops.matmul(image, weights) + bias, name="logits")
nn_ops.softmax(logit, name="prediction")
cost = nn_ops.softmax_cross_entropy_with_logits(labels=label,
@@ -2299,9 +2300,9 @@
with ops_lib.device("/job:ps/replica:0/task:0/device:GPU:0"):
image = array_ops.placeholder(dtypes.float32, [None, 784], name="image")
label = array_ops.placeholder(dtypes.float32, [None, 10], name="label")
- weights = variables.Variable(
+ weights = variables.VariableV1(
random_ops.random_uniform([784, 10]), name="weights")
- bias = variables.Variable(array_ops.zeros([10]), name="bias")
+ bias = variables.VariableV1(array_ops.zeros([10]), name="bias")
logit = nn_ops.relu(math_ops.matmul(image, weights) + bias)
nn_ops.softmax(logit, name="prediction")
cost = nn_ops.softmax_cross_entropy_with_logits(labels=label,
@@ -2332,9 +2333,9 @@
with ops_lib.device("/job:ps/replica:0/task:0/device:GPU:0"):
image = array_ops.placeholder(dtypes.float32, [None, 784], name="image")
label = array_ops.placeholder(dtypes.float32, [None, 10], name="label")
- weights = variables.Variable(
+ weights = variables.VariableV1(
random_ops.random_uniform([784, 10]), name="weights")
- bias = variables.Variable(array_ops.zeros([10]), name="bias")
+ bias = variables.VariableV1(array_ops.zeros([10]), name="bias")
logit = nn_ops.relu(math_ops.matmul(image, weights) + bias)
nn_ops.softmax(logit, name="prediction")
cost = nn_ops.softmax_cross_entropy_with_logits(labels=label,
@@ -2385,9 +2386,9 @@
def testDebugString(self):
# Builds a graph.
- v0 = variables.Variable(
+ v0 = variables.VariableV1(
[[1, 2, 3], [4, 5, 6]], dtype=dtypes.float32, name="v0")
- v1 = variables.Variable(
+ v1 = variables.VariableV1(
[[[1], [2]], [[3], [4]], [[5], [6]]], dtype=dtypes.float32, name="v1")
init_all_op = variables.global_variables_initializer()
save = saver_module.Saver(
@@ -2444,7 +2445,8 @@
def testWriteGraph(self):
test_dir = self._get_test_dir("write_graph_dir")
- variables.Variable([[1, 2, 3], [4, 5, 6]], dtype=dtypes.float32, name="v0")
+ variables.VariableV1(
+ [[1, 2, 3], [4, 5, 6]], dtype=dtypes.float32, name="v0")
path = graph_io.write_graph(ops_lib.get_default_graph(),
os.path.join(test_dir, "l1"), "graph.pbtxt")
truth = os.path.join(test_dir, "l1", "graph.pbtxt")
@@ -2453,7 +2455,8 @@
def testRecursiveCreate(self):
test_dir = self._get_test_dir("deep_dir")
- variables.Variable([[1, 2, 3], [4, 5, 6]], dtype=dtypes.float32, name="v0")
+ variables.VariableV1(
+ [[1, 2, 3], [4, 5, 6]], dtype=dtypes.float32, name="v0")
path = graph_io.write_graph(ops_lib.get_default_graph().as_graph_def(),
os.path.join(test_dir, "l1", "l2", "l3"),
"graph.pbtxt")
@@ -2477,7 +2480,7 @@
images = constant_op.constant(
1.2, dtypes.float32, shape=[100, 28], name="images")
with ops_lib.name_scope("hidden1"):
- weights1 = variables.Variable(
+ weights1 = variables.VariableV1(
random_ops.truncated_normal(
[28, 128], stddev=1.0 / math.sqrt(float(28))),
name="weights")
@@ -2485,7 +2488,7 @@
# coverage the save and restore of control flow context (which doesn't
# make any sense here from a machine learning perspective). The typical
# biases is a simple Variable without the conditions.
- biases1 = variables.Variable(
+ biases1 = variables.VariableV1(
control_flow_ops.cond(
math_ops.less(random.random(), 0.5),
lambda: array_ops.ones([128]), lambda: array_ops.zeros([128])),
@@ -2494,7 +2497,7 @@
# Hidden 2
with ops_lib.name_scope("hidden2"):
- weights2 = variables.Variable(
+ weights2 = variables.VariableV1(
random_ops.truncated_normal(
[128, 32], stddev=1.0 / math.sqrt(float(128))),
name="weights")
@@ -2511,16 +2514,16 @@
return it + 1, biases2
_, biases2 = control_flow_ops.while_loop(loop_cond, loop_body, [
- constant_op.constant(0), variables.Variable(array_ops.zeros([32]))
+ constant_op.constant(0), variables.VariableV1(array_ops.zeros([32]))
])
hidden2 = nn_ops.relu(math_ops.matmul(hidden1, weights2) + biases2)
# Linear
with ops_lib.name_scope("softmax_linear"):
- weights3 = variables.Variable(
+ weights3 = variables.VariableV1(
random_ops.truncated_normal(
[32, 10], stddev=1.0 / math.sqrt(float(32))),
name="weights")
- biases3 = variables.Variable(array_ops.zeros([10]), name="biases")
+ biases3 = variables.VariableV1(array_ops.zeros([10]), name="biases")
logits = math_ops.matmul(hidden2, weights3) + biases3
ops_lib.add_to_collection("logits", logits)
@@ -2566,7 +2569,7 @@
with graph.as_default():
# Hidden 2
with ops_lib.name_scope("hidden2"):
- weights = variables.Variable(
+ weights = variables.VariableV1(
random_ops.truncated_normal(
[128, 32], stddev=1.0 / math.sqrt(float(128))),
name="weights")
@@ -2583,16 +2586,16 @@
return it + 1, biases
_, biases = control_flow_ops.while_loop(loop_cond, loop_body, [
- constant_op.constant(0), variables.Variable(array_ops.zeros([32]))
+ constant_op.constant(0), variables.VariableV1(array_ops.zeros([32]))
])
hidden2 = nn_ops.relu(math_ops.matmul(hidden1, weights) + biases)
# Linear
with ops_lib.name_scope("softmax_linear"):
- weights = variables.Variable(
+ weights = variables.VariableV1(
random_ops.truncated_normal(
[32, 10], stddev=1.0 / math.sqrt(float(32))),
name="weights")
- biases = variables.Variable(array_ops.zeros([10]), name="biases")
+ biases = variables.VariableV1(array_ops.zeros([10]), name="biases")
logits = math_ops.matmul(hidden2, weights) + biases
ops_lib.add_to_collection("logits", logits)
@@ -2629,9 +2632,9 @@
with ops_lib.name_scope("hidden1"):
images = constant_op.constant(
1.0, dtypes.float32, shape=[3, 2], name="images")
- weights1 = variables.Variable(
+ weights1 = variables.VariableV1(
[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], name="weights")
- biases1 = variables.Variable([0.1] * 3, name="biases")
+ biases1 = variables.VariableV1([0.1] * 3, name="biases")
nn_ops.relu(math_ops.matmul(images, weights1) + biases1, name="relu")
# Run the graph and save scoped checkpoint.
@@ -2685,9 +2688,9 @@
with ops_lib.name_scope("hidden1"):
images = constant_op.constant(
1.0, dtypes.float32, shape=[3, 2], name="images")
- weights1 = variables.Variable(
+ weights1 = variables.VariableV1(
[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], name="weights")
- biases1 = variables.Variable([0.1] * 3, name="biases")
+ biases1 = variables.VariableV1([0.1] * 3, name="biases")
nn_ops.relu(math_ops.matmul(images, weights1) + biases1, name="relu")
# Run the graph and save scoped checkpoint.
@@ -2720,12 +2723,12 @@
graph = ops_lib.Graph()
with graph.as_default():
with ops_lib.name_scope("hidden1"):
- variable1 = variables.Variable([1.0], name="variable1")
+ variable1 = variables.VariableV1([1.0], name="variable1")
saver1 = saver_module.Saver(var_list=[variable1])
graph.add_to_collection(ops_lib.GraphKeys.SAVERS, saver1)
with ops_lib.name_scope("hidden2"):
- variable2 = variables.Variable([2.0], name="variable2")
+ variable2 = variables.VariableV1([2.0], name="variable2")
saver2 = saver_module.Saver(var_list=[variable2], name="hidden2/")
graph.add_to_collection(ops_lib.GraphKeys.SAVERS, saver2)
@@ -2978,7 +2981,7 @@
checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt")
with ops_lib.Graph().as_default() as g:
- a = variables.Variable(1., name="a")
+ a = variables.VariableV1(1., name="a")
a_saver = saver_module.Saver([a])
with self.session(graph=g) as sess:
@@ -2986,7 +2989,7 @@
save_path = a_saver.save(sess=sess, save_path=checkpoint_prefix)
with ops_lib.Graph().as_default() as g:
- a = variables.Variable([1.], name="a")
+ a = variables.VariableV1([1.], name="a")
a_saver = saver_module.Saver([a])
with self.session(graph=g) as sess:
with self.assertRaisesRegexp(
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 c7e84e9..5aa7f45 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
@@ -37,8 +37,8 @@
server = server_lib.Server.create_local_server()
with session.Session(server.target) as sess_1:
- v0 = variables.Variable([[2, 1]], name="v0")
- v1 = variables.Variable([[1], [2]], name="v1")
+ v0 = variables.VariableV1([[2, 1]], name="v0")
+ v1 = variables.VariableV1([[1], [2]], name="v1")
v2 = math_ops.matmul(v0, v1)
sess_1.run([v0.initializer, v1.initializer])
self.assertAllEqual([[4]], sess_1.run(v2))
diff --git a/tensorflow/python/training/server_lib_test.py b/tensorflow/python/training/server_lib_test.py
index 063044f..cf99570 100644
--- a/tensorflow/python/training/server_lib_test.py
+++ b/tensorflow/python/training/server_lib_test.py
@@ -76,9 +76,9 @@
def testResetFails(self):
# Creates variable with container name.
with ops.container("test0"):
- v0 = variables.Variable(1.0, name="v0")
+ v0 = variables.VariableV1(1.0, name="v0")
# Creates variable with default container.
- v1 = variables.Variable(2.0, name="v1")
+ v1 = variables.VariableV1(2.0, name="v1")
# Verifies resetting the non-existent target returns error.
with self.assertRaises(errors_impl.NotFoundError):
session.Session.reset("nonexistent", ["test0"])
@@ -234,8 +234,8 @@
[0.], dtype=dtypes.float32))
self.assertIsNotNone(input_queue)
- var = variables.Variable(1., dtype=dtypes.float32, trainable=False,
- name="var")
+ var = variables.VariableV1(1., dtype=dtypes.float32, trainable=False,
+ name="var")
sess.run(variables.global_variables_initializer())
queue_runner_impl.start_queue_runners(sess)
@@ -245,7 +245,7 @@
server = self._cached_server
init_value = array_ops.placeholder(dtypes.int32)
- v = variables.Variable(init_value, validate_shape=False, name="v")
+ v = variables.VariableV1(init_value, validate_shape=False, name="v")
sharing_config = config_pb2.ConfigProto(isolate_session_state=False)
sharing_sess_0 = session.Session(server.target, config=sharing_config)
@@ -302,7 +302,7 @@
isolate_config = config_pb2.ConfigProto(isolate_session_state=True)
with ops.Graph().as_default():
- w_vector = variables.Variable([1, 2, 3], name="w")
+ w_vector = variables.VariableV1([1, 2, 3], name="w")
with session.Session(server.target, config=sharing_config) as sess:
with self.assertRaises(errors_impl.FailedPreconditionError):
sess.run(w_vector)
@@ -310,20 +310,20 @@
self.assertAllEqual([1, 2, 3], sess.run(w_vector))
with ops.Graph().as_default():
- w_vector = variables.Variable([4, 5, 6], name="w")
+ w_vector = variables.VariableV1([4, 5, 6], name="w")
with session.Session(server.target, config=sharing_config) as sess:
self.assertAllEqual([1, 2, 3], sess.run(w_vector))
sess.run(w_vector.initializer)
self.assertAllEqual([4, 5, 6], sess.run(w_vector))
with ops.Graph().as_default():
- w_scalar = variables.Variable(86, name="w")
+ w_scalar = variables.VariableV1(86, name="w")
with session.Session(server.target, config=sharing_config) as sess:
with self.assertRaises(errors_impl.InvalidArgumentError):
sess.run(w_scalar.initializer)
with ops.Graph().as_default():
- w_scalar = variables.Variable(37, name="w")
+ w_scalar = variables.VariableV1(37, name="w")
with session.Session(server.target, config=isolate_config) as sess:
with self.assertRaises(errors_impl.FailedPreconditionError):
sess.run(w_scalar)
diff --git a/tensorflow/python/training/session_manager.py b/tensorflow/python/training/session_manager.py
index a2e0645..cd313c2 100644
--- a/tensorflow/python/training/session_manager.py
+++ b/tensorflow/python/training/session_manager.py
@@ -25,6 +25,7 @@
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
@@ -182,6 +183,12 @@
"""
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 f1d18f7..2b5c3b0 100644
--- a/tensorflow/python/training/session_manager_test.py
+++ b/tensorflow/python/training/session_manager_test.py
@@ -40,7 +40,7 @@
def testPrepareSessionSucceeds(self):
with ops.Graph().as_default():
- v = variables.Variable([1.0, 2.0, 3.0], name="v")
+ v = variables.VariableV1([1.0, 2.0, 3.0], name="v")
sm = session_manager.SessionManager(
ready_op=variables.report_uninitialized_variables())
sess = sm.prepare_session(
@@ -50,7 +50,7 @@
def testPrepareSessionSucceedsWithInitFeedDict(self):
with ops.Graph().as_default():
p = array_ops.placeholder(dtypes.float32, shape=(3,))
- v = variables.Variable(p, name="v")
+ v = variables.VariableV1(p, name="v")
sm = session_manager.SessionManager(
ready_op=variables.report_uninitialized_variables())
sess = sm.prepare_session(
@@ -61,7 +61,7 @@
def testPrepareSessionSucceedsWithInitFn(self):
with ops.Graph().as_default():
- v = variables.Variable([125], name="v")
+ v = variables.VariableV1([125], name="v")
sm = session_manager.SessionManager(
ready_op=variables.report_uninitialized_variables())
sess = sm.prepare_session(
@@ -79,7 +79,7 @@
gfile.MakeDirs(checkpoint_dir)
with ops.Graph().as_default():
- v = variables.Variable([1.0, 2.0, 3.0], name="v")
+ v = variables.VariableV1([1.0, 2.0, 3.0], name="v")
sm = session_manager.SessionManager(
ready_op=variables.report_uninitialized_variables())
saver = saver_lib.Saver({"v": v})
@@ -97,7 +97,7 @@
# Renames the checkpoint directory.
os.rename(checkpoint_dir, checkpoint_dir2)
gfile.MakeDirs(checkpoint_dir)
- v = variables.Variable([6.0, 7.0, 8.0], name="v")
+ v = variables.VariableV1([6.0, 7.0, 8.0], name="v")
with self.cached_session():
self.assertEqual(False, variables.is_variable_initialized(v).eval())
session_manager.SessionManager(
@@ -134,7 +134,7 @@
checkpoint_filename_with_path=None):
# Create a new Graph and SessionManager and recover from a checkpoint.
with ops.Graph().as_default():
- v = variables.Variable(2, name="v")
+ v = variables.VariableV1(2, name="v")
with session_lib.Session():
self.assertEqual(False, variables.is_variable_initialized(v).eval())
sm2 = session_manager.SessionManager(
@@ -162,7 +162,7 @@
gfile.MakeDirs(checkpoint_dir)
with ops.Graph().as_default():
- v = variables.Variable(1, name="v")
+ v = variables.VariableV1(1, name="v")
sm = session_manager.SessionManager(
ready_op=variables.report_uninitialized_variables())
saver = saver_lib.Saver({"v": v})
@@ -186,7 +186,7 @@
def testWaitForSessionReturnsNoneAfterTimeout(self):
with ops.Graph().as_default():
- variables.Variable(1, name="v")
+ variables.VariableV1(1, name="v")
sm = session_manager.SessionManager(
ready_op=variables.report_uninitialized_variables(),
recovery_wait_secs=1)
@@ -217,7 +217,7 @@
gfile.MakeDirs(checkpoint_dir)
with ops.Graph().as_default():
- v = variables.Variable(1, name="v")
+ v = variables.VariableV1(1, name="v")
sm = session_manager.SessionManager(
ready_op=variables.report_uninitialized_variables())
saver = saver_lib.Saver({"v": v})
@@ -230,8 +230,8 @@
os.path.join(checkpoint_dir, "recover_session_checkpoint"))
# Create a new Graph and SessionManager and recover.
with ops.Graph().as_default():
- v = variables.Variable(2, name="v")
- w = variables.Variable(
+ v = variables.VariableV1(2, name="v")
+ w = variables.VariableV1(
v,
trainable=False,
collections=[ops.GraphKeys.LOCAL_VARIABLES],
@@ -275,7 +275,7 @@
gfile.MakeDirs(checkpoint_dir)
with ops.Graph().as_default():
- v = variables.Variable(1, name="v")
+ v = variables.VariableV1(1, name="v")
sm = session_manager.SessionManager(
ready_op=variables.report_uninitialized_variables())
saver = saver_lib.Saver({"v": v})
@@ -288,8 +288,8 @@
os.path.join(checkpoint_dir, "recover_session_checkpoint"))
# Create a new Graph and SessionManager and recover.
with ops.Graph().as_default():
- v = variables.Variable(2, name="v")
- w = variables.Variable(
+ v = variables.VariableV1(2, name="v")
+ w = variables.VariableV1(
v,
trainable=False,
collections=[ops.GraphKeys.LOCAL_VARIABLES],
@@ -321,7 +321,7 @@
# local_init_op exactly once, regardless of whether the session was
# successfully recovered.
with ops.Graph().as_default():
- w = variables.Variable(
+ w = variables.VariableV1(
1,
trainable=False,
collections=[ops.GraphKeys.LOCAL_VARIABLES],
@@ -356,8 +356,8 @@
# Create a new Graph and SessionManager and recover.
with ops.Graph().as_default():
- v = variables.Variable(2, name="v")
- w = variables.Variable(
+ v = variables.VariableV1(2, name="v")
+ w = variables.VariableV1(
1,
trainable=False,
collections=[ops.GraphKeys.LOCAL_VARIABLES],
@@ -389,8 +389,8 @@
def testWaitForSessionLocalInit(self):
server = server_lib.Server.create_local_server()
with ops.Graph().as_default() as graph:
- v = variables.Variable(1, name="v")
- w = variables.Variable(
+ v = variables.VariableV1(1, name="v")
+ w = variables.VariableV1(
v,
trainable=False,
collections=[ops.GraphKeys.LOCAL_VARIABLES],
@@ -420,8 +420,8 @@
def testWaitForSessionWithReadyForLocalInitOpFailsToReadyLocal(self):
with ops.Graph().as_default() as graph:
- v = variables.Variable(1, name="v")
- w = variables.Variable(
+ v = variables.VariableV1(1, name="v")
+ w = variables.VariableV1(
v,
trainable=False,
collections=[ops.GraphKeys.LOCAL_VARIABLES],
@@ -439,8 +439,8 @@
def testWaitForSessionInsufficientReadyForLocalInitCheck(self):
with ops.Graph().as_default() as graph:
- v = variables.Variable(1, name="v")
- w = variables.Variable(
+ v = variables.VariableV1(1, name="v")
+ w = variables.VariableV1(
v,
trainable=False,
collections=[ops.GraphKeys.LOCAL_VARIABLES],
@@ -456,13 +456,13 @@
def testPrepareSessionWithReadyForLocalInitOp(self):
with ops.Graph().as_default():
- v = variables.Variable(1, name="v")
- w = variables.Variable(
+ v = variables.VariableV1(1, name="v")
+ w = variables.VariableV1(
v,
trainable=False,
collections=[ops.GraphKeys.LOCAL_VARIABLES],
name="w")
- x = variables.Variable(
+ x = variables.VariableV1(
3 * v,
trainable=False,
collections=[ops.GraphKeys.LOCAL_VARIABLES],
@@ -495,25 +495,25 @@
def testPrepareSessionWithPartialInitOp(self):
with ops.Graph().as_default():
- v = variables.Variable(1, name="v")
- w = variables.Variable(
+ v = variables.VariableV1(1, name="v")
+ w = variables.VariableV1(
v,
trainable=False,
collections=[ops.GraphKeys.LOCAL_VARIABLES],
name="w")
- x = variables.Variable(
+ x = variables.VariableV1(
3 * v,
trainable=False,
collections=[ops.GraphKeys.LOCAL_VARIABLES],
name="x")
# TODO(b/70206927): Use ResourceVariables once they are handled properly.
- v_res = variables.Variable(1, name="v_res")
- w_res = variables.Variable(
+ v_res = variables.VariableV1(1, name="v_res")
+ w_res = variables.VariableV1(
v_res,
trainable=False,
collections=[ops.GraphKeys.LOCAL_VARIABLES],
name="w_res")
- x_res = variables.Variable(
+ x_res = variables.VariableV1(
3 * v_res,
trainable=False,
collections=[ops.GraphKeys.LOCAL_VARIABLES],
@@ -565,7 +565,7 @@
# cyclic dependencies.
with ops.Graph().as_default():
i = control_flow_ops.while_loop(lambda i: i < 1, lambda i: i + 1, [0])
- v = variables.Variable(array_ops.identity(i), name="v")
+ v = variables.VariableV1(array_ops.identity(i), name="v")
with self.cached_session():
self.assertEqual(False, variables.is_variable_initialized(v).eval())
sm = session_manager.SessionManager(
@@ -579,8 +579,8 @@
def testPrepareSessionDidNotInitLocalVariable(self):
with ops.Graph().as_default():
- v = variables.Variable(1, name="v")
- w = variables.Variable(
+ v = variables.VariableV1(1, name="v")
+ w = variables.VariableV1(
v,
trainable=False,
collections=[ops.GraphKeys.LOCAL_VARIABLES],
@@ -596,8 +596,8 @@
def testPrepareSessionDidNotInitLocalVariableList(self):
with ops.Graph().as_default():
- v = variables.Variable(1, name="v")
- w = variables.Variable(
+ v = variables.VariableV1(1, name="v")
+ w = variables.VariableV1(
v,
trainable=False,
collections=[ops.GraphKeys.LOCAL_VARIABLES],
@@ -613,8 +613,8 @@
def testPrepareSessionWithReadyNotReadyForLocal(self):
with ops.Graph().as_default():
- v = variables.Variable(1, name="v")
- w = variables.Variable(
+ v = variables.VariableV1(1, name="v")
+ w = variables.VariableV1(
v,
trainable=False,
collections=[ops.GraphKeys.LOCAL_VARIABLES],
@@ -634,8 +634,8 @@
def testPrepareSessionWithInsufficientReadyForLocalInitCheck(self):
with ops.Graph().as_default():
- v = variables.Variable(1, name="v")
- w = variables.Variable(
+ v = variables.VariableV1(1, name="v")
+ w = variables.VariableV1(
v,
trainable=False,
collections=[ops.GraphKeys.LOCAL_VARIABLES],
@@ -656,7 +656,7 @@
def testPrepareSessionSucceeds(self):
with ops.Graph().as_default():
- v = variables.Variable([1.0, 2.0, 3.0], name="v")
+ v = variables.VariableV1([1.0, 2.0, 3.0], name="v")
sm = session_manager.SessionManager(
ready_op=variables.assert_variables_initialized())
sess = sm.prepare_session(
@@ -666,7 +666,7 @@
def testPrepareSessionSucceedsWithInitFeedDict(self):
with ops.Graph().as_default():
p = array_ops.placeholder(dtypes.float32, shape=(3,))
- v = variables.Variable(p, name="v")
+ v = variables.VariableV1(p, name="v")
sm = session_manager.SessionManager(
ready_op=variables.assert_variables_initialized())
sess = sm.prepare_session(
@@ -677,7 +677,7 @@
def testPrepareSessionSucceedsWithInitFn(self):
with ops.Graph().as_default():
- v = variables.Variable([125], name="v")
+ v = variables.VariableV1([125], name="v")
sm = session_manager.SessionManager(
ready_op=variables.assert_variables_initialized())
sess = sm.prepare_session(
@@ -695,7 +695,7 @@
gfile.MakeDirs(checkpoint_dir)
with ops.Graph().as_default():
- v = variables.Variable([1.0, 2.0, 3.0], name="v")
+ v = variables.VariableV1([1.0, 2.0, 3.0], name="v")
sm = session_manager.SessionManager(
ready_op=variables.assert_variables_initialized())
saver = saver_lib.Saver({"v": v})
@@ -713,7 +713,7 @@
# Renames the checkpoint directory.
os.rename(checkpoint_dir, checkpoint_dir2)
gfile.MakeDirs(checkpoint_dir)
- v = variables.Variable([6.0, 7.0, 8.0], name="v")
+ v = variables.VariableV1([6.0, 7.0, 8.0], name="v")
with self.cached_session():
self.assertEqual(False, variables.is_variable_initialized(v).eval())
session_manager.SessionManager(
@@ -755,7 +755,7 @@
gfile.MakeDirs(checkpoint_dir)
with ops.Graph().as_default():
- v = variables.Variable(1, name="v")
+ v = variables.VariableV1(1, name="v")
sm = session_manager.SessionManager(
ready_op=variables.assert_variables_initialized())
saver = saver_lib.Saver({"v": v})
@@ -768,7 +768,7 @@
os.path.join(checkpoint_dir, "recover_session_checkpoint"))
# Create a new Graph and SessionManager and recover.
with ops.Graph().as_default():
- v = variables.Variable(2, name="v")
+ v = variables.VariableV1(2, name="v")
with self.cached_session():
self.assertEqual(False, variables.is_variable_initialized(v).eval())
sm2 = session_manager.SessionManager(
@@ -785,7 +785,7 @@
def testWaitForSessionReturnsNoneAfterTimeout(self):
with ops.Graph().as_default():
- variables.Variable(1, name="v")
+ variables.VariableV1(1, name="v")
sm = session_manager.SessionManager(
ready_op=variables.assert_variables_initialized(),
recovery_wait_secs=1)
diff --git a/tensorflow/python/training/supervisor.py b/tensorflow/python/training/supervisor.py
index 0755364..a5e626d 100644
--- a/tensorflow/python/training/supervisor.py
+++ b/tensorflow/python/training/supervisor.py
@@ -242,10 +242,9 @@
ready_for_local_init_op: 1-D string `Tensor`. This tensor is evaluated by
supervisors in `prepare_or_wait_for_session()` to check if the model is
ready to run the local_init_op.
- The model is considered ready if it returns an empty array. Defaults to
- the tensor returned from
- `tf.report_uninitialized_variables(tf.global_variables())`. If `None`,
- the model is not checked for readiness before running local_init_op.
+ The model is considered ready if it returns an empty array. Defaults to
+ `None`. If `None`, the model is not checked for readiness before running
+ local_init_op.
is_chief: If True, create a chief supervisor in charge of initializing
and restoring the model. If False, create a supervisor that relies
on a chief supervisor for inits and restore.
diff --git a/tensorflow/python/training/supervisor_test.py b/tensorflow/python/training/supervisor_test.py
index caf6eba..7cd99d8 100644
--- a/tensorflow/python/training/supervisor_test.py
+++ b/tensorflow/python/training/supervisor_test.py
@@ -423,7 +423,7 @@
def testLogdirButExplicitlyNoSummaryWriter(self):
logdir = self._test_dir("explicit_no_summary_writer")
with ops.Graph().as_default():
- variables.Variable([1.0], name="foo")
+ variables.VariableV1([1.0], name="foo")
summary.scalar("c1", constant_op.constant(1))
summary.scalar("c2", constant_op.constant(2))
summary.scalar("c3", constant_op.constant(3))
@@ -491,7 +491,7 @@
def testNoLogdirSucceeds(self):
with ops.Graph().as_default():
- variables.Variable([1.0, 2.0, 3.0])
+ variables.VariableV1([1.0, 2.0, 3.0])
sv = supervisor.Supervisor(logdir="", summary_op=None)
sess = sv.prepare_or_wait_for_session("")
sess.close()
@@ -499,7 +499,7 @@
def testUseSessionManager(self):
with ops.Graph().as_default():
- variables.Variable([1.0, 2.0, 3.0])
+ variables.VariableV1([1.0, 2.0, 3.0])
sm = session_manager_lib.SessionManager()
# Pass in session_manager. The additional init_op is ignored.
sv = supervisor.Supervisor(logdir="", session_manager=sm)
@@ -508,7 +508,7 @@
def testInitOp(self):
logdir = self._test_dir("default_init_op")
with ops.Graph().as_default():
- v = variables.Variable([1.0, 2.0, 3.0])
+ v = variables.VariableV1([1.0, 2.0, 3.0])
sv = supervisor.Supervisor(logdir=logdir)
sess = sv.prepare_or_wait_for_session("")
self.assertAllClose([1.0, 2.0, 3.0], sess.run(v))
@@ -517,7 +517,7 @@
def testInitFn(self):
logdir = self._test_dir("default_init_op")
with ops.Graph().as_default():
- v = variables.Variable([1.0, 2.0, 3.0])
+ v = variables.VariableV1([1.0, 2.0, 3.0])
def _init_fn(sess):
sess.run(v.initializer)
@@ -531,7 +531,7 @@
logdir = self._test_dir("feed_dict_init_op")
with ops.Graph().as_default():
p = array_ops.placeholder(dtypes.float32, shape=(3,))
- v = variables.Variable(p, name="v")
+ v = variables.VariableV1(p, name="v")
sv = supervisor.Supervisor(
logdir=logdir,
init_op=variables.global_variables_initializer(),
@@ -550,10 +550,10 @@
g = ops.Graph()
with g.as_default():
with ops.device("/job:local"):
- v = variables.Variable(
+ v = variables.VariableV1(
1, name="default_ready_for_local_init_op_v_" + str(uid))
vadd = v.assign_add(1)
- w = variables.Variable(
+ w = variables.VariableV1(
v,
trainable=False,
collections=[ops.GraphKeys.LOCAL_VARIABLES],
@@ -590,7 +590,7 @@
# Create a checkpoint.
with ops.Graph().as_default():
- v = variables.Variable(
+ v = variables.VariableV1(
10.0, name="ready_for_local_init_op_restore_v_" + str(uid))
summary.scalar("ready_for_local_init_op_restore_v_" + str(uid), v)
sv = supervisor.Supervisor(logdir=logdir)
@@ -607,10 +607,10 @@
g = ops.Graph()
with g.as_default():
with ops.device("/job:local"):
- v = variables.Variable(
+ v = variables.VariableV1(
1.0, name="ready_for_local_init_op_restore_v_" + str(uid))
vadd = v.assign_add(1)
- w = variables.Variable(
+ w = variables.VariableV1(
v,
trainable=False,
collections=[ops.GraphKeys.LOCAL_VARIABLES],
@@ -642,13 +642,13 @@
logdir = self._test_dir("default_local_init_op")
with ops.Graph().as_default():
# A local variable.
- v = variables.Variable(
+ v = variables.VariableV1(
[1.0, 2.0, 3.0],
trainable=False,
collections=[ops.GraphKeys.LOCAL_VARIABLES])
# An entity which is initialized through a TABLE_INITIALIZER.
- w = variables.Variable([4, 5, 6], trainable=False, collections=[])
+ w = variables.VariableV1([4, 5, 6], trainable=False, collections=[])
ops.add_to_collection(ops.GraphKeys.TABLE_INITIALIZERS, w.initializer)
# This shouldn't add a variable to the VARIABLES collection responsible
@@ -668,7 +668,7 @@
with ops.Graph().as_default():
with ops.device("/job:localhost"):
# A local variable.
- v = variables.Variable(
+ v = variables.VariableV1(
[1.0, 2.0, 3.0],
trainable=False,
collections=[ops.GraphKeys.LOCAL_VARIABLES])
@@ -687,8 +687,8 @@
server = server_lib.Server.create_local_server()
logdir = self._test_dir("default_init_op_fails")
with ops.Graph().as_default():
- v = variables.Variable([1.0, 2.0, 3.0], name="v")
- variables.Variable([4.0, 5.0, 6.0], name="w")
+ v = variables.VariableV1([1.0, 2.0, 3.0], name="v")
+ variables.VariableV1([4.0, 5.0, 6.0], name="w")
# w will not be initialized.
sv = supervisor.Supervisor(logdir=logdir, init_op=v.initializer)
with self.assertRaisesRegexp(RuntimeError,
@@ -699,11 +699,11 @@
server = server_lib.Server.create_local_server()
logdir = self._test_dir("default_init_op_fails_for_local_variable")
with ops.Graph().as_default():
- v = variables.Variable(
+ v = variables.VariableV1(
[1.0, 2.0, 3.0],
name="v",
collections=[ops.GraphKeys.LOCAL_VARIABLES])
- variables.Variable(
+ variables.VariableV1(
[1.0, 2.0, 3.0],
name="w",
collections=[ops.GraphKeys.LOCAL_VARIABLES])
@@ -716,17 +716,17 @@
def testSetupFail(self):
logdir = self._test_dir("setup_fail")
with ops.Graph().as_default():
- variables.Variable([1.0, 2.0, 3.0], name="v")
+ variables.VariableV1([1.0, 2.0, 3.0], name="v")
with self.assertRaisesRegexp(ValueError, "must have their device set"):
supervisor.Supervisor(logdir=logdir, is_chief=False)
with ops.Graph().as_default(), ops.device("/job:ps"):
- variables.Variable([1.0, 2.0, 3.0], name="v")
+ variables.VariableV1([1.0, 2.0, 3.0], name="v")
supervisor.Supervisor(logdir=logdir, is_chief=False)
def testDefaultGlobalStep(self):
logdir = self._test_dir("default_global_step")
with ops.Graph().as_default():
- variables.Variable(287, name="global_step")
+ variables.VariableV1(287, name="global_step")
sv = supervisor.Supervisor(logdir=logdir)
sess = sv.prepare_or_wait_for_session("")
self.assertEquals(287, sess.run(sv.global_step))
@@ -735,7 +735,7 @@
def testRestoreFromMetaGraph(self):
logdir = self._test_dir("restore_from_meta_graph")
with ops.Graph().as_default():
- variables.Variable(1, name="v0")
+ variables.VariableV1(1, name="v0")
sv = supervisor.Supervisor(logdir=logdir)
sess = sv.prepare_or_wait_for_session("")
filename = sv.saver.save(sess, sv.save_path)
@@ -757,7 +757,7 @@
logdir = self._test_dir("standard_services_without_global_step")
# Create a checkpoint.
with ops.Graph().as_default():
- v = variables.Variable([1.0], name="foo")
+ v = variables.VariableV1([1.0], name="foo")
summary.scalar("v", v[0])
sv = supervisor.Supervisor(logdir=logdir)
meta_graph_def = meta_graph.create_meta_graph_def(
@@ -796,7 +796,7 @@
self.assertRaises(StopIteration, lambda: next(rr))
# There should be a checkpoint file with the variable "foo"
with ops.Graph().as_default(), self.cached_session() as sess:
- v = variables.Variable([10.10], name="foo")
+ v = variables.VariableV1([10.10], name="foo")
sav = saver_lib.Saver([v])
sav.restore(sess, save_path)
self.assertEqual(1.0, v.eval()[0])
@@ -807,7 +807,7 @@
logdir = self._test_dir("standard_services_with_global_step")
# Create a checkpoint.
with ops.Graph().as_default():
- v = variables.Variable([123], name="global_step")
+ v = variables.VariableV1([123], name="global_step")
sv = supervisor.Supervisor(logdir=logdir)
meta_graph_def = meta_graph.create_meta_graph_def(
saver_def=sv.saver.saver_def)
@@ -860,7 +860,7 @@
self.assertRaises(StopIteration, lambda: next(rr))
# There should be a checkpoint file with the variable "foo"
with ops.Graph().as_default(), self.cached_session() as sess:
- v = variables.Variable([-12], name="global_step")
+ v = variables.VariableV1([-12], name="global_step")
sav = saver_lib.Saver([v])
sav.restore(sess, save_path)
self.assertEqual(123, v.eval()[0])
diff --git a/tensorflow/python/training/sync_replicas_optimizer_test.py b/tensorflow/python/training/sync_replicas_optimizer_test.py
index fff1740..1ef8756 100644
--- a/tensorflow/python/training/sync_replicas_optimizer_test.py
+++ b/tensorflow/python/training/sync_replicas_optimizer_test.py
@@ -40,11 +40,12 @@
is_chief = (worker_id == 0)
with graph.as_default():
with ops.device("/job:ps/task:0"):
- global_step = variables.Variable(0, name="global_step", trainable=False)
- var_0 = variables.Variable(0.0, name="v0")
+ global_step = variables.VariableV1(
+ 0, name="global_step", trainable=False)
+ var_0 = variables.VariableV1(0.0, name="v0")
with ops.device("/job:ps/task:1"):
- var_1 = variables.Variable(1.0, name="v1")
- var_sparse = variables.Variable([[3.0], [4.0]], name="v_sparse")
+ var_1 = variables.VariableV1(1.0, name="v1")
+ var_sparse = variables.VariableV1([[3.0], [4.0]], name="v_sparse")
with ops.device("/job:worker/task:" + str(worker_id)):
grads_0 = constant_op.constant(0.1 + worker_id * 0.2)
@@ -272,8 +273,8 @@
replicas_to_aggregate=1,
total_num_replicas=1)
hook = opt.make_session_run_hook(True)
- v = variables.Variable([0.])
- global_step = variables.Variable(0, name="global_step", trainable=False)
+ v = variables.VariableV1([0.])
+ global_step = variables.VariableV1(0, name="global_step", trainable=False)
opt.minimize(v, global_step=global_step)
hook.begin()
@@ -282,8 +283,8 @@
opt=adam.AdamOptimizer(0.01),
replicas_to_aggregate=1,
total_num_replicas=1)
- v = variables.Variable([0.], name="fetch_variable_test")
- global_step = variables.Variable(0, name="global_step", trainable=False)
+ v = variables.VariableV1([0.], name="fetch_variable_test")
+ global_step = variables.VariableV1(0, name="global_step", trainable=False)
opt.minimize(v, global_step=global_step)
opt_variables = opt.variables()
beta1_power, beta2_power = opt._opt._get_beta_accumulators()
diff --git a/tensorflow/python/training/training_ops_test.py b/tensorflow/python/training/training_ops_test.py
index d131a11..f410cea 100644
--- a/tensorflow/python/training/training_ops_test.py
+++ b/tensorflow/python/training/training_ops_test.py
@@ -51,7 +51,7 @@
def _testTypes(self, x, alpha, delta, use_gpu=None):
self.setUp()
with self.test_session(use_gpu=use_gpu):
- var = variables.Variable(x)
+ var = variables.VariableV1(x)
variables.global_variables_initializer().run()
self.assertAllCloseAccordingToType(x, var.eval())
apply_sgd = training_ops.apply_gradient_descent(var, alpha, delta)
@@ -70,8 +70,8 @@
def _testTypesForAdagrad(self, x, y, lr, grad, use_gpu=None):
self.setUp()
with self.test_session(use_gpu=use_gpu):
- var = variables.Variable(x)
- accum = variables.Variable(y)
+ var = variables.VariableV1(x)
+ accum = variables.VariableV1(y)
variables.global_variables_initializer().run()
self.assertAllCloseAccordingToType(x, var.eval())
@@ -94,9 +94,9 @@
lr_power=-0.5):
self.setUp()
with self.test_session(use_gpu=use_gpu):
- var = variables.Variable(x)
- accum = variables.Variable(y)
- linear = variables.Variable(z)
+ var = variables.VariableV1(x)
+ accum = variables.VariableV1(y)
+ linear = variables.VariableV1(z)
variables.global_variables_initializer().run()
self.assertAllCloseAccordingToType(x, var.eval())
@@ -148,8 +148,8 @@
def _testTypesForSparseAdagrad(self, x, y, lr, grad, indices):
self.setUp()
with self.test_session(use_gpu=False):
- var = variables.Variable(x)
- accum = variables.Variable(y)
+ var = variables.VariableV1(x)
+ accum = variables.VariableV1(y)
variables.global_variables_initializer().run()
self.assertAllCloseAccordingToType(x, var.eval())
@@ -178,9 +178,9 @@
lr_power=-0.5):
self.setUp()
with self.test_session(use_gpu=False):
- var = variables.Variable(x)
- accum = variables.Variable(y)
- linear = variables.Variable(z)
+ var = variables.VariableV1(x)
+ accum = variables.VariableV1(y)
+ linear = variables.VariableV1(z)
variables.global_variables_initializer().run()
self.assertAllCloseAccordingToType(x, var.eval())
@@ -257,9 +257,9 @@
def _testTypesForAdam(self, var, m, v, grad, use_gpu):
self.setUp()
with self.test_session(use_gpu=use_gpu):
- var_t = variables.Variable(var)
- m_t = variables.Variable(m)
- v_t = variables.Variable(v)
+ var_t = variables.VariableV1(var)
+ m_t = variables.VariableV1(m)
+ v_t = variables.VariableV1(v)
t = 1
beta1 = np.array(0.9, dtype=var.dtype)
@@ -270,8 +270,8 @@
epsilon = np.array(1e-8, dtype=var.dtype)
beta1_t = constant_op.constant(beta1, self._toType(var.dtype), [])
beta2_t = constant_op.constant(beta2, self._toType(var.dtype), [])
- beta1_power_t = variables.Variable(beta1_power)
- beta2_power_t = variables.Variable(beta2_power)
+ beta1_power_t = variables.VariableV1(beta1_power)
+ 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()
diff --git a/tensorflow/python/training/training_util_test.py b/tensorflow/python/training/training_util_test.py
index 6cc177e..ba64e78 100644
--- a/tensorflow/python/training/training_util_test.py
+++ b/tensorflow/python/training/training_util_test.py
@@ -49,7 +49,7 @@
def test_invalid_shape(self):
with ops.Graph().as_default() as g:
self.assertIsNone(training_util.get_global_step())
- variables.Variable(
+ variables.VariableV1(
[0],
trainable=False,
dtype=dtypes.int32,
@@ -73,7 +73,7 @@
def test_get_global_step(self):
with ops.Graph().as_default() as g:
self.assertIsNone(training_util.get_global_step())
- variables.Variable(
+ variables.VariableV1(
0,
trainable=False,
dtype=dtypes.int32,
diff --git a/tensorflow/python/util/function_utils.py b/tensorflow/python/util/function_utils.py
index 4e9b07e..a56dfbf 100644
--- a/tensorflow/python/util/function_utils.py
+++ b/tensorflow/python/util/function_utils.py
@@ -59,6 +59,29 @@
return tuple(args)
+def has_kwargs(fn):
+ """Returns whether the passed callable has **kwargs in its signature.
+
+ Args:
+ fn: Function, or function-like object (e.g., result of `functools.partial`).
+
+ Returns:
+ `bool`: if `fn` has **kwargs in its signature.
+
+ Raises:
+ `TypeError`: If fn is not a Function, or function-like object.
+ """
+ if isinstance(fn, functools.partial):
+ fn = fn.func
+ elif _is_callable_object(fn):
+ fn = fn.__call__
+ elif not callable(fn):
+ raise TypeError(
+ 'fn should be a function-like object, but is of type {}.'.format(
+ type(fn)))
+ return tf_inspect.getfullargspec(fn).varkw is not None
+
+
def get_func_name(func):
"""Returns name of passed callable."""
_, func = tf_decorator.unwrap(func)
diff --git a/tensorflow/python/util/function_utils_test.py b/tensorflow/python/util/function_utils_test.py
index 1588328..e5b0843 100644
--- a/tensorflow/python/util/function_utils_test.py
+++ b/tensorflow/python/util/function_utils_test.py
@@ -135,6 +135,101 @@
self.assertEqual(3, double_wrapped_fn(a=3))
+class HasKwargsTest(test.TestCase):
+
+ def test_simple_function(self):
+
+ fn_has_kwargs = lambda **x: x
+ self.assertTrue(function_utils.has_kwargs(fn_has_kwargs))
+
+ fn_has_no_kwargs = lambda x: x
+ self.assertFalse(function_utils.has_kwargs(fn_has_no_kwargs))
+
+ def test_callable(self):
+
+ class FooHasKwargs(object):
+
+ def __call__(self, **x):
+ del x
+ self.assertTrue(function_utils.has_kwargs(FooHasKwargs()))
+
+ class FooHasNoKwargs(object):
+
+ def __call__(self, x):
+ del x
+ self.assertFalse(function_utils.has_kwargs(FooHasNoKwargs()))
+
+ def test_bounded_method(self):
+
+ class FooHasKwargs(object):
+
+ def fn(self, **x):
+ del x
+ self.assertTrue(function_utils.has_kwargs(FooHasKwargs().fn))
+
+ class FooHasNoKwargs(object):
+
+ def fn(self, x):
+ del x
+ self.assertFalse(function_utils.has_kwargs(FooHasNoKwargs().fn))
+
+ def test_partial_function(self):
+ expected_test_arg = 123
+
+ def fn_has_kwargs(test_arg, **x):
+ if test_arg != expected_test_arg:
+ return ValueError('partial fn does not work correctly')
+ return x
+
+ wrapped_fn = functools.partial(fn_has_kwargs, test_arg=123)
+ self.assertTrue(function_utils.has_kwargs(wrapped_fn))
+ some_kwargs = dict(x=1, y=2, z=3)
+ self.assertEqual(wrapped_fn(**some_kwargs), some_kwargs)
+
+ def fn_has_no_kwargs(x, test_arg):
+ if test_arg != expected_test_arg:
+ return ValueError('partial fn does not work correctly')
+ return x
+
+ wrapped_fn = functools.partial(fn_has_no_kwargs, test_arg=123)
+ self.assertFalse(function_utils.has_kwargs(wrapped_fn))
+ some_arg = 1
+ self.assertEqual(wrapped_fn(some_arg), some_arg)
+
+ def test_double_partial(self):
+ expected_test_arg1 = 123
+ expected_test_arg2 = 456
+
+ def fn_has_kwargs(test_arg1, test_arg2, **x):
+ if test_arg1 != expected_test_arg1 or test_arg2 != expected_test_arg2:
+ return ValueError('partial does not work correctly')
+ return x
+
+ wrapped_fn = functools.partial(fn_has_kwargs, test_arg2=456)
+ double_wrapped_fn = functools.partial(wrapped_fn, test_arg1=123)
+
+ self.assertTrue(function_utils.has_kwargs(double_wrapped_fn))
+ some_kwargs = dict(x=1, y=2, z=3)
+ self.assertEqual(double_wrapped_fn(**some_kwargs), some_kwargs)
+
+ def fn_has_no_kwargs(x, test_arg1, test_arg2):
+ if test_arg1 != expected_test_arg1 or test_arg2 != expected_test_arg2:
+ return ValueError('partial does not work correctly')
+ return x
+
+ wrapped_fn = functools.partial(fn_has_no_kwargs, test_arg2=456)
+ double_wrapped_fn = functools.partial(wrapped_fn, test_arg1=123)
+
+ self.assertFalse(function_utils.has_kwargs(double_wrapped_fn))
+ some_arg = 1
+ self.assertEqual(double_wrapped_fn(some_arg), some_arg)
+
+ def test_raises_type_error(self):
+ with self.assertRaisesRegexp(
+ TypeError, 'fn should be a function-like object'):
+ function_utils.has_kwargs('not a function')
+
+
class GetFuncNameTest(test.TestCase):
def testWithSimpleFunction(self):
diff --git a/tensorflow/python/util/nest.py b/tensorflow/python/util/nest.py
index 653ca52..d67dbde 100644
--- a/tensorflow/python/util/nest.py
+++ b/tensorflow/python/util/nest.py
@@ -19,6 +19,9 @@
Python sequence, tuple (including `namedtuple`), or dict that can contain
further sequences, tuples, and dicts.
+attr.s decorated classes (http://www.attrs.org) are also supported, in the
+same way as `namedtuple`.
+
The utilities here assume (and do not check) that the nested structures form a
'tree', i.e., no references in the structure of the input of these functions
should be recursive.
@@ -38,6 +41,12 @@
from tensorflow.python import pywrap_tensorflow as _pywrap_tensorflow
+def _get_attrs_values(obj):
+ """Returns the list of values from an attrs instance."""
+ attrs = getattr(obj.__class__, "__attrs_attrs__")
+ return [getattr(obj, a.name) for a in attrs]
+
+
def _sorted(dict_):
"""Returns a sorted list of the dict keys, with error if keys not sortable."""
try:
@@ -64,6 +73,7 @@
# See the swig file (util.i) for documentation.
_is_mapping = _pywrap_tensorflow.IsMapping
+_is_attrs = _pywrap_tensorflow.IsAttrs
def _sequence_like(instance, args):
@@ -85,7 +95,7 @@
# corresponding `OrderedDict` to pack it back).
result = dict(zip(_sorted(instance), args))
return type(instance)((key, result[key]) for key in _six.iterkeys(instance))
- elif _is_namedtuple(instance):
+ elif _is_namedtuple(instance) or _is_attrs(instance):
return type(instance)(*args)
else:
# Not a namedtuple
@@ -93,6 +103,7 @@
def _yield_value(iterable):
+ """Yields the next value from the given iterable."""
if _is_mapping(iterable):
# Iterate through dictionaries in a deterministic order by sorting the
# keys. Notice this means that we ignore the original order of `OrderedDict`
@@ -101,6 +112,9 @@
# corresponding `OrderedDict` to pack it back).
for key in _sorted(iterable):
yield iterable[key]
+ elif _is_attrs(iterable):
+ for value in _get_attrs_values(iterable):
+ yield value
else:
for value in iterable:
yield value
@@ -805,5 +819,5 @@
return list(zip(flat_string_paths, flatten(structure)))
-_pywrap_tensorflow.RegisterSequenceClass(_collections.Sequence)
-_pywrap_tensorflow.RegisterMappingClass(_collections.Mapping)
+_pywrap_tensorflow.RegisterType("Mapping", _collections.Mapping)
+_pywrap_tensorflow.RegisterType("Sequence", _collections.Sequence)
diff --git a/tensorflow/python/util/nest_test.py b/tensorflow/python/util/nest_test.py
index bfb4c6f..e03a8da 100644
--- a/tensorflow/python/util/nest_test.py
+++ b/tensorflow/python/util/nest_test.py
@@ -33,6 +33,11 @@
from tensorflow.python.platform import test
from tensorflow.python.util import nest
+try:
+ import attr # pylint:disable=g-import-not-at-top
+except ImportError:
+ attr = None
+
class _CustomMapping(collections.Mapping):
@@ -53,6 +58,35 @@
PointXY = collections.namedtuple("Point", ["x", "y"]) # pylint: disable=invalid-name
+ if attr:
+ class BadAttr(object):
+ """Class that has a non-iterable __attrs_attrs__."""
+ __attrs_attrs__ = None
+
+ @attr.s
+ class SampleAttr(object):
+ field1 = attr.ib()
+ field2 = attr.ib()
+
+ @test_util.assert_no_new_pyobjects_executing_eagerly
+ def testAttrsFlattenAndPack(self):
+ if attr is None:
+ self.skipTest("attr module is unavailable.")
+
+ field_values = [1, 2]
+ sample_attr = NestTest.SampleAttr(*field_values)
+ self.assertFalse(nest._is_attrs(field_values))
+ self.assertTrue(nest._is_attrs(sample_attr))
+ flat = nest.flatten(sample_attr)
+ self.assertEqual(field_values, flat)
+ restructured_from_flat = nest.pack_sequence_as(sample_attr, flat)
+ self.assertIsInstance(restructured_from_flat, NestTest.SampleAttr)
+ self.assertEqual(restructured_from_flat, sample_attr)
+
+ # Check that flatten fails if attributes are not iterable
+ with self.assertRaisesRegexp(TypeError, "object is not iterable"):
+ flat = nest.flatten(NestTest.BadAttr())
+
@test_util.assert_no_new_pyobjects_executing_eagerly
def testFlattenAndPack(self):
structure = ((3, 4), 5, (6, 7, (9, 10), 8))
diff --git a/tensorflow/python/util/util.cc b/tensorflow/python/util/util.cc
index 2087957..7b3e618 100644
--- a/tensorflow/python/util/util.cc
+++ b/tensorflow/python/util/util.cc
@@ -29,14 +29,51 @@
namespace tensorflow {
namespace swig {
+std::unordered_map<string, PyObject*>* PythonTypesMap() {
+ static auto* m = new std::unordered_map<string, PyObject*>();
+ return m;
+}
+
+PyObject* GetRegisteredType(const string& key) {
+ auto* m = PythonTypesMap();
+ auto it = m->find(key);
+ if (it == m->end()) return nullptr;
+ return it->second;
+}
+
+PyObject* RegisterType(PyObject* type_name, PyObject* type) {
+ if (!PyType_Check(type)) {
+ PyErr_SetString(PyExc_TypeError,
+ tensorflow::strings::StrCat("Expecting a type, got ",
+ Py_TYPE(type)->tp_name)
+ .c_str());
+ return nullptr;
+ }
+
+ string key;
+ if (PyBytes_Check(type_name)) {
+ key = PyBytes_AsString(type_name);
+ }
+#if PY_MAJOR_VERSION >= 3
+ if (PyUnicode_Check(type_name)) {
+ key = PyUnicode_AsUTF8(type_name);
+ }
+#endif
+
+ if (PythonTypesMap()->find(key) != PythonTypesMap()->end()) {
+ PyErr_SetString(PyExc_TypeError, tensorflow::strings::StrCat(
+ "Type already registered for ", key)
+ .c_str());
+ return nullptr;
+ }
+
+ Py_INCREF(type);
+ PythonTypesMap()->emplace(key, type);
+
+ Py_RETURN_NONE;
+}
+
namespace {
-
-// Type object for collections.Sequence. This is set by RegisterSequenceClass.
-PyObject* CollectionsSequenceType = nullptr;
-// Type object for collections.Mapping, set by RegisterMappingClass.
-PyObject* CollectionsMappingType = nullptr;
-PyTypeObject* SparseTensorValueType = nullptr;
-
const int kMaxItemsInCache = 1024;
bool WarnedThatSetIsNotSequence = false;
@@ -177,18 +214,75 @@
// Returns -1 if an error occurred.
int IsMappingHelper(PyObject* o) {
static auto* const check_cache = new CachedTypeCheck([](PyObject* to_check) {
- return PyObject_IsInstance(to_check, CollectionsMappingType);
+ PyObject* collections_mapping_type = GetRegisteredType("Mapping");
+ if (TF_PREDICT_FALSE(collections_mapping_type == nullptr)) {
+ PyErr_SetString(PyExc_RuntimeError,
+ tensorflow::strings::StrCat(
+ "collections.Mapping type has not been set. "
+ "Please register the type with the identifier "
+ "\"Mapping\" using RegisterType.")
+ .c_str());
+ return -1;
+ }
+ return PyObject_IsInstance(to_check, collections_mapping_type);
});
if (PyDict_Check(o)) return true;
- if (TF_PREDICT_FALSE(CollectionsMappingType == nullptr)) {
- PyErr_SetString(
- PyExc_RuntimeError,
- tensorflow::strings::StrCat(
- "collections.Mapping type has not been set. "
- "Please call RegisterMappingClass before using this module")
- .c_str());
- return -1;
- }
+ return check_cache->CachedLookup(o);
+}
+
+// Returns 1 if `o` is an instance of attrs-decorated class.
+// Returns 0 otherwise.
+int IsAttrsHelper(PyObject* o) {
+ static auto* const check_cache = new CachedTypeCheck([](PyObject* to_check) {
+ Safe_PyObjectPtr cls(PyObject_GetAttrString(to_check, "__class__"));
+ if (cls) {
+ return PyObject_HasAttrString(cls.get(), "__attrs_attrs__");
+ }
+
+ // PyObject_GetAttrString returns null on error
+ PyErr_Clear();
+ return 0;
+ });
+ return check_cache->CachedLookup(o);
+}
+
+// Returns 1 if `o` is an object of type IndexedSlices.
+// Returns 0 otherwise.
+// Returns -1 if an error occurred.
+int IsIndexedSlicesHelper(PyObject* o) {
+ static auto* const check_cache = new CachedTypeCheck([](PyObject* to_check) {
+ PyObject* indexed_slices_type = GetRegisteredType("IndexedSlices");
+ if (TF_PREDICT_FALSE(indexed_slices_type == nullptr)) {
+ PyErr_SetString(PyExc_RuntimeError,
+ tensorflow::strings::StrCat(
+ "IndexedSlices type has not been set. "
+ "Please register the type with the identifier "
+ "\"IndexedSlices\" using RegisterType.")
+ .c_str());
+ return -1;
+ }
+ return PyObject_IsInstance(to_check, indexed_slices_type);
+ });
+ return check_cache->CachedLookup(o);
+}
+
+// Returns 1 if `o` is a Tensor.
+// Returns 0 otherwise.
+// Returns -1 if an error occurred.
+int IsTensorHelper(PyObject* o) {
+ static auto* const check_cache = new CachedTypeCheck([](PyObject* to_check) {
+ PyObject* tensor_type = GetRegisteredType("Tensor");
+ if (TF_PREDICT_FALSE(tensor_type == nullptr)) {
+ PyErr_SetString(PyExc_RuntimeError,
+ tensorflow::strings::StrCat(
+ "Tensor type has not been set. "
+ "Please register the type with the identifier "
+ "\"Tensor\" using RegisterType.")
+ .c_str());
+ return -1;
+ }
+ return PyObject_IsInstance(to_check, tensor_type);
+ });
return check_cache->CachedLookup(o);
}
@@ -196,31 +290,33 @@
// Returns 0 otherwise.
// Returns -1 if an error occurred.
int IsSequenceHelper(PyObject* o) {
- static auto* const check_cache = new CachedTypeCheck([](PyObject* to_check) {
- int is_instance = PyObject_IsInstance(to_check, CollectionsSequenceType);
-
- // Don't cache a failed is_instance check.
- if (is_instance == -1) return -1;
-
- return static_cast<int>(is_instance != 0 && !IsString(to_check));
- });
// We treat dicts and other mappings as special cases of sequences.
if (IsMappingHelper(o)) return true;
+ if (IsAttrsHelper(o)) return true;
if (PySet_Check(o) && !WarnedThatSetIsNotSequence) {
LOG(WARNING) << "Sets are not currently considered sequences, "
"but this may change in the future, "
"so consider avoiding using them.";
WarnedThatSetIsNotSequence = true;
}
- if (TF_PREDICT_FALSE(CollectionsSequenceType == nullptr)) {
- PyErr_SetString(
- PyExc_RuntimeError,
- tensorflow::strings::StrCat(
- "collections.Sequence type has not been set. "
- "Please call RegisterSequenceClass before using this module")
- .c_str());
- return -1;
- }
+ static auto* const check_cache = new CachedTypeCheck([](PyObject* to_check) {
+ PyObject* collections_sequence_type = GetRegisteredType("Sequence");
+ if (TF_PREDICT_FALSE(collections_sequence_type == nullptr)) {
+ PyErr_SetString(PyExc_RuntimeError,
+ tensorflow::strings::StrCat(
+ "collections.Sequence type has not been set. "
+ "Please register the type with the identifier "
+ "\"Sequence\" using RegisterType.")
+ .c_str());
+ return -1;
+ }
+ int is_instance = PyObject_IsInstance(to_check, collections_sequence_type);
+
+ // Don't cache a failed is_instance check.
+ if (is_instance == -1) return -1;
+
+ return static_cast<int>(is_instance != 0 && !IsString(to_check));
+ });
return check_cache->CachedLookup(o);
}
@@ -354,12 +450,46 @@
Safe_PyObjectPtr tensor_;
};
+class AttrsValueIterator : public ValueIterator {
+ public:
+ explicit AttrsValueIterator(PyObject* nested) : nested_(nested) {
+ Py_INCREF(nested);
+ cls_.reset(PyObject_GetAttrString(nested_.get(), "__class__"));
+ if (cls_) {
+ attrs_.reset(PyObject_GetAttrString(cls_.get(), "__attrs_attrs__"));
+ if (attrs_) {
+ iter_.reset(PyObject_GetIter(attrs_.get()));
+ }
+ }
+ if (!iter_ || PyErr_Occurred()) invalidate();
+ }
+
+ Safe_PyObjectPtr next() override {
+ Safe_PyObjectPtr result;
+ Safe_PyObjectPtr item(PyIter_Next(iter_.get()));
+ if (item) {
+ Safe_PyObjectPtr name(PyObject_GetAttrString(item.get(), "name"));
+ result.reset(PyObject_GetAttr(nested_.get(), name.get()));
+ }
+
+ return result;
+ }
+
+ private:
+ Safe_PyObjectPtr nested_;
+ Safe_PyObjectPtr cls_;
+ Safe_PyObjectPtr attrs_;
+ Safe_PyObjectPtr iter_;
+};
+
bool IsSparseTensorValueType(PyObject* o) {
- if (TF_PREDICT_FALSE(SparseTensorValueType == nullptr)) {
+ PyObject* sparse_tensor_value_type = GetRegisteredType("SparseTensorValue");
+ if (TF_PREDICT_FALSE(sparse_tensor_value_type == nullptr)) {
return false;
}
- return PyObject_TypeCheck(o, SparseTensorValueType) == 1;
+ return PyObject_TypeCheck(
+ o, reinterpret_cast<PyTypeObject*>(sparse_tensor_value_type)) == 1;
}
int IsSequenceForDataHelper(PyObject* o) {
@@ -372,6 +502,8 @@
return absl::make_unique<DictValueIterator>(nested);
} else if (IsMappingHelper(nested)) {
return absl::make_unique<MappingValueIterator>(nested);
+ } else if (IsAttrsHelper(nested)) {
+ return absl::make_unique<AttrsValueIterator>(nested);
} else {
return absl::make_unique<SequenceValueIterator>(nested);
}
@@ -383,6 +515,8 @@
return absl::make_unique<DictValueIterator>(nested);
} else if (IsMappingHelper(nested)) {
return absl::make_unique<MappingValueIterator>(nested);
+ } else if (IsAttrsHelper(nested)) {
+ return absl::make_unique<AttrsValueIterator>(nested);
} else if (IsSparseTensorValueType(nested)) {
return absl::make_unique<SparseTensorValueIterator>(nested);
} else {
@@ -597,48 +731,11 @@
} // namespace
-void RegisterSequenceClass(PyObject* sequence_class) {
- if (!PyType_Check(sequence_class)) {
- PyErr_SetString(
- PyExc_TypeError,
- tensorflow::strings::StrCat(
- "Expecting a class definition for `collections.Sequence`. Got ",
- Py_TYPE(sequence_class)->tp_name)
- .c_str());
- return;
- }
- CollectionsSequenceType = sequence_class;
-}
-
-void RegisterMappingClass(PyObject* mapping_class) {
- if (!PyType_Check(mapping_class)) {
- PyErr_SetString(
- PyExc_TypeError,
- tensorflow::strings::StrCat(
- "Expecting a class definition for `collections.Mapping`. Got ",
- Py_TYPE(mapping_class)->tp_name)
- .c_str());
- return;
- }
- CollectionsMappingType = mapping_class;
-}
-
-void RegisterSparseTensorValueClass(PyObject* sparse_tensor_value_class) {
- if (!PyType_Check(sparse_tensor_value_class)) {
- PyErr_SetString(
- PyExc_TypeError,
- tensorflow::strings::StrCat(
- "Expecting a class definition for `SparseTensorValue`. Got ",
- Py_TYPE(sparse_tensor_value_class)->tp_name)
- .c_str());
- return;
- }
- SparseTensorValueType =
- reinterpret_cast<PyTypeObject*>(sparse_tensor_value_class);
-}
-
bool IsSequence(PyObject* o) { return IsSequenceHelper(o) == 1; }
bool IsMapping(PyObject* o) { return IsMappingHelper(o) == 1; }
+bool IsAttrs(PyObject* o) { return IsAttrsHelper(o) == 1; }
+bool IsTensor(PyObject* o) { return IsTensorHelper(o) == 1; }
+bool IsIndexedSlices(PyObject* o) { return IsIndexedSlicesHelper(o) == 1; }
PyObject* Flatten(PyObject* nested) {
PyObject* list = PyList_New(0);
@@ -686,13 +783,15 @@
}
}
- if (TF_PREDICT_FALSE(CollectionsSequenceType == nullptr)) {
- PyErr_SetString(
- PyExc_RuntimeError,
- tensorflow::strings::StrCat(
- "collections.Sequence type has not been set. "
- "Please call RegisterSequenceClass before using this module")
- .c_str());
+ PyObject* collections_sequence_type = GetRegisteredType("Sequence");
+
+ if (TF_PREDICT_FALSE(collections_sequence_type == nullptr)) {
+ PyErr_SetString(PyExc_RuntimeError,
+ tensorflow::strings::StrCat(
+ "collections.Sequence type has not been set. "
+ "Please register the type with the identifier "
+ "\"Sequence\" using RegisterType.")
+ .c_str());
return nullptr;
}
@@ -704,7 +803,8 @@
}
Safe_PyObjectPtr fields = make_safe(PyObject_GetAttrString(o, "_fields"));
- int is_instance = PyObject_IsInstance(fields.get(), CollectionsSequenceType);
+ int is_instance =
+ PyObject_IsInstance(fields.get(), collections_sequence_type);
if (is_instance == 0) {
Py_RETURN_FALSE;
} else if (is_instance == -1) {
diff --git a/tensorflow/python/util/util.h b/tensorflow/python/util/util.h
index 3436052..f37cd52 100644
--- a/tensorflow/python/util/util.h
+++ b/tensorflow/python/util/util.h
@@ -56,6 +56,33 @@
// True if the sequence subclasses mapping.
bool IsMapping(PyObject* o);
+// Returns a true if its input is an instance of an attr.s decorated class.
+//
+// Args:
+// o: the input to be checked.
+//
+// Returns:
+// True if the object is an instance of an attr.s decorated class.
+bool IsAttrs(PyObject* o);
+
+// Returns a true if its input is an ops.Tensor.
+//
+// Args:
+// seq: the input to be checked.
+//
+// Returns:
+// True if the object is a tensor.
+bool IsTensor(PyObject* o);
+
+// Returns a true if its input is an ops.IndexesSlices.
+//
+// Args:
+// seq: the input to be checked.
+//
+// Returns:
+// True if the object is an ops.IndexedSlices.
+bool IsIndexedSlices(PyObject* o);
+
// Implements the same interface as tensorflow.util.nest._same_namedtuples
// Returns Py_True iff the two namedtuples have the same name and fields.
// Raises RuntimeError if `o1` or `o2` don't look like namedtuples (don't have
@@ -121,18 +148,6 @@
// TypeError: The nest is or contains a dict with non-sortable keys.
PyObject* Flatten(PyObject* nested);
-// RegisterSequenceClass is used to pass PyTypeObject for collections.Sequence
-// (which is defined in python) into the C++ world.
-// Alternative approach could be to import the collections modules and retrieve
-// the type from the module. This approach also requires some trigger from
-// Python so that we know that Python interpreter had been initialzied.
-void RegisterSequenceClass(PyObject* sequence_class);
-// Like RegisterSequenceClass, but for collections.Mapping.
-void RegisterMappingClass(PyObject* mapping_class);
-// Similar to the above functions, except for the
-// sparse_tensor.SparseTensorValue class.
-void RegisterSparseTensorValueClass(PyObject* sparse_tensor_value_class);
-
// The tensorflow.python.data package has its own nest utility that follows very
// slightly different semantics for its functions than the tensorflow.python
// nest utility. Returns a true if its input is a collections.Sequence (except
@@ -158,6 +173,10 @@
PyObject* AssertSameStructureForData(PyObject* o1, PyObject* o2,
bool check_types);
+// RegisterType is used to pass PyTypeObject (which is defined in python) for an
+// arbitrary identifier `type_name` into C++.
+PyObject* RegisterType(PyObject* type_name, PyObject* type);
+
} // namespace swig
} // namespace tensorflow
diff --git a/tensorflow/python/util/util.i b/tensorflow/python/util/util.i
index 104a615..3c0ec87 100644
--- a/tensorflow/python/util/util.i
+++ b/tensorflow/python/util/util.i
@@ -28,14 +28,8 @@
// for functions in this module because they use python methods that need GIL.
// TODO(iga): Find a way not to leak such definitions across files.
-%unignore tensorflow::swig::RegisterSequenceClass;
-%noexception tensorflow::swig::RegisterSequenceClass;
-
-%unignore tensorflow::swig::RegisterMappingClass;
-%noexception tensorflow::swig::RegisterMappingClass;
-
-%unignore tensorflow::swig::RegisterSparseTensorValueClass;
-%noexception tensorflow::swig::RegisterSparseTensorValueClass;
+%unignore tensorflow::swig::RegisterType;
+%noexception tensorflow::swig::RegisterType;
%feature("docstring") tensorflow::swig::IsSequence
"""Returns a true if its input is a collections.Sequence (except strings).
@@ -65,6 +59,18 @@
%unignore tensorflow::swig::IsMapping;
%noexception tensorflow::swig::IsMapping;
+%feature("docstring") tensorflow::swig::IsAttrs
+"""Returns True iff `instance` is an instance of an `attr.s` decorated class.
+
+Args:
+ instance: An instance of a Python object.
+
+Returns:
+ True if `instance` is an instance of an `attr.s` decorated class.
+"""
+%unignore tensorflow::swig::IsAttrs;
+%noexception tensorflow::swig::IsAttrs;
+
%feature("docstring") tensorflow::swig::SameNamedtuples
"Returns True if the two namedtuples have the same name and fields."
%unignore tensorflow::swig::SameNamedtuples;
diff --git a/tensorflow/tensorflow.bzl b/tensorflow/tensorflow.bzl
index d6c75d6..cad5de1 100644
--- a/tensorflow/tensorflow.bzl
+++ b/tensorflow/tensorflow.bzl
@@ -19,6 +19,14 @@
"@local_config_cuda//cuda:build_defs.bzl",
"cuda_default_copts",
"if_cuda",
+ "if_cuda_is_configured",
+)
+load(
+ "@local_config_rocm//rocm:build_defs.bzl",
+ "if_rocm",
+ "if_rocm_is_configured",
+ "rocm_copts",
+ "rocm_default_copts",
)
load(
"//third_party/mkl:build_defs.bzl",
@@ -39,6 +47,8 @@
def register_extension_info(**kwargs):
pass
+# if_cuda_is_configured def placeholder
+
# Given a source file, generate a test name.
# i.e. "common_runtime/direct_session_test.cc" becomes
# "common_runtime_direct_session_test"
@@ -863,12 +873,16 @@
srcs = srcs + tf_binary_additional_srcs(),
size = size,
args = args,
- copts = _cuda_copts() + tf_copts(),
+ copts = _cuda_copts() + rocm_copts() + tf_copts(),
data = data + tf_binary_dynamic_kernel_dsos(kernels),
- deps = deps + tf_binary_dynamic_kernel_deps(kernels) + if_cuda([
- clean_dep("//tensorflow/core:cuda"),
- clean_dep("//tensorflow/core:gpu_lib"),
- ]),
+ deps = deps + tf_binary_dynamic_kernel_deps(kernels) +
+ if_cuda_is_configured([
+ clean_dep("//tensorflow/core:cuda"),
+ clean_dep("//tensorflow/core:gpu_lib"),
+ ]) +
+ if_rocm_is_configured([
+ clean_dep("//tensorflow/core:gpu_lib"),
+ ]),
linkopts = if_not_windows(["-lpthread", "-lm"]) + linkopts + _rpath_linkopts(name),
linkstatic = linkstatic or select({
# cc_tests with ".so"s in srcs incorrectly link on Darwin
@@ -1003,7 +1017,7 @@
label_regex_for_dep = "{extension_name}",
)
-def _cuda_copts():
+def _cuda_copts(opts = []):
"""Gets the appropriate set of copts for (maybe) CUDA compilation.
If we're doing CUDA compilation, returns copts for our particular CUDA
@@ -1019,13 +1033,17 @@
"@local_config_cuda//cuda:using_clang": ([
"-fcuda-flush-denormals-to-zero",
]),
- })
+ }) + if_cuda_is_configured(opts)
# Build defs for TensorFlow kernels
# When this target is built using --config=cuda, a cc_library is built
# that passes -DGOOGLE_CUDA=1 and '-x cuda', linking in additional
# libraries needed by GPU kernels.
+#
+# When this target is built using --config=rocm, a cc_library is built
+# that passes -DTENSORFLOW_USE_ROCM and '-x rocm', linking in additional
+# libraries needed by GPU kernels.
def tf_gpu_kernel_library(
srcs,
copts = [],
@@ -1033,16 +1051,18 @@
deps = [],
hdrs = [],
**kwargs):
- copts = copts + _cuda_copts() + if_cuda(cuda_copts) + tf_copts()
+ copts = copts + tf_copts() + _cuda_copts(opts = cuda_copts) + rocm_copts(opts = cuda_copts)
kwargs["features"] = kwargs.get("features", []) + ["-use_header_modules"]
native.cc_library(
srcs = srcs,
hdrs = hdrs,
copts = copts,
- deps = deps + if_cuda([
+ deps = deps + if_cuda_is_configured([
clean_dep("//tensorflow/core:cuda"),
clean_dep("//tensorflow/core:gpu_lib"),
+ ]) + if_rocm_is_configured([
+ clean_dep("//tensorflow/core:gpu_lib"),
]),
alwayslink = 1,
**kwargs
@@ -1081,8 +1101,10 @@
deps = deps + if_cuda(cuda_deps + [
clean_dep("//tensorflow/core:cuda"),
"@local_config_cuda//cuda:cuda_headers",
+ ]) + if_rocm_is_configured(cuda_deps + [
+ # rocm_header placeholder
]),
- copts = (copts + if_cuda(["-DGOOGLE_CUDA=1"]) + if_mkl(["-DINTEL_MKL=1"]) +
+ copts = (copts + if_cuda(["-DGOOGLE_CUDA=1"]) + if_rocm(["-DTENSORFLOW_USE_ROCM=1"]) + if_mkl(["-DINTEL_MKL=1"]) +
if_mkl_open_source_only(["-DINTEL_MKL_DNN_ONLY"]) +
if_enable_mkl(["-DENABLE_MKL"]) +
if_tensorrt(["-DGOOGLE_TENSORRT=1"])),
@@ -1465,6 +1487,9 @@
"@local_config_cuda//cuda:cuda_headers",
"@local_config_cuda//cuda:cudart_static",
]
+ rocm_deps = [
+ clean_dep("//tensorflow/core:stream_executor_headers_lib"),
+ ]
deps = deps + tf_custom_op_library_additional_deps()
if gpu_srcs:
basename = name.split(".")[0]
@@ -1473,13 +1498,14 @@
srcs = gpu_srcs,
copts = _cuda_copts() + if_tensorrt(["-DGOOGLE_TENSORRT=1"]),
features = if_cuda(["-use_header_modules"]),
- deps = deps + if_cuda(cuda_deps),
+ deps = deps + if_cuda_is_configured(cuda_deps) + if_rocm_is_configured(rocm_deps),
)
cuda_deps.extend([":" + basename + "_gpu"])
+ rocm_deps.extend([":" + basename + "_gpu"])
check_deps(
name = name + "_check_deps",
- deps = deps + if_cuda(cuda_deps),
+ deps = deps + if_cuda_is_configured(cuda_deps) + if_rocm_is_configured(rocm_deps),
disallowed_deps = [
clean_dep("//tensorflow/core:framework"),
clean_dep("//tensorflow/core:lib"),
@@ -1488,7 +1514,7 @@
tf_cc_shared_object(
name = name,
srcs = srcs,
- deps = deps + if_cuda(cuda_deps),
+ deps = deps + if_cuda_is_configured(cuda_deps) + if_rocm_is_configured(rocm_deps),
data = if_static([name + "_check_deps"]),
copts = tf_copts(is_external = True),
features = ["windows_export_all_symbols"],
@@ -1693,6 +1719,29 @@
label_regex_for_dep = "{extension_name}",
)
+# Similar to py_test above, this macro is used to exclude dependencies for some py_binary
+# targets in order to reduce the size of //tensorflow/tools/pip_package:simple_console_windows.
+# See https://github.com/tensorflow/tensorflow/issues/22390
+def py_binary(name, deps = [], **kwargs):
+ # Add an extra target for dependencies to avoid nested select statement.
+ native.py_library(
+ name = name + "_deps",
+ deps = deps,
+ )
+ native.py_binary(
+ name = name,
+ deps = select({
+ "//conditions:default": [":" + name + "_deps"],
+ clean_dep("//tensorflow:no_tensorflow_py_deps"): [],
+ }),
+ **kwargs
+ )
+
+register_extension_info(
+ extension_name = "py_binary",
+ label_regex_for_dep = "{extension_name}",
+)
+
def tf_py_test(
name,
srcs,
diff --git a/tensorflow/tools/api/golden/tensorflow.-config-proto.-experimental.pbtxt b/tensorflow/tools/api/golden/tensorflow.-config-proto.-experimental.pbtxt
deleted file mode 100644
index eb41dee..0000000
--- a/tensorflow/tools/api/golden/tensorflow.-config-proto.-experimental.pbtxt
+++ /dev/null
@@ -1,24 +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: "client_handles_error_formatting"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_BOOL
- }
- field {
- name: "executor_type"
- number: 3
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- }
-}
diff --git a/tensorflow/tools/api/golden/tensorflow.-config-proto.pbtxt b/tensorflow/tools/api/golden/tensorflow.-config-proto.pbtxt
deleted file mode 100644
index e565b90..0000000
--- a/tensorflow/tools/api/golden/tensorflow.-config-proto.pbtxt
+++ /dev/null
@@ -1,148 +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: "client_handles_error_formatting"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_BOOL
- }
- field {
- name: "executor_type"
- number: 3
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- }
- }
-}
diff --git a/tensorflow/tools/api/golden/tensorflow.data.-iterator.pbtxt b/tensorflow/tools/api/golden/tensorflow.data.-iterator.pbtxt
deleted file mode 100644
index 4f0147a..0000000
--- a/tensorflow/tools/api/golden/tensorflow.data.-iterator.pbtxt
+++ /dev/null
@@ -1,46 +0,0 @@
-path: "tensorflow.data.Iterator"
-tf_class {
- is_instance: "<class \'tensorflow.python.data.ops.iterator_ops.Iterator\'>"
- is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
- is_instance: "<type \'object\'>"
- member {
- name: "initializer"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_classes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_shapes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_types"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'iterator_resource\', \'initializer\', \'output_types\', \'output_shapes\', \'output_classes\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "from_string_handle"
- argspec: "args=[\'string_handle\', \'output_types\', \'output_shapes\', \'output_classes\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "from_structure"
- argspec: "args=[\'output_types\', \'output_shapes\', \'shared_name\', \'output_classes\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
- }
- member_method {
- name: "get_next"
- argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "make_initializer"
- argspec: "args=[\'self\', \'dataset\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "string_handle"
- argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
-}
diff --git a/tensorflow/tools/api/golden/tensorflow.estimator.-boosted-trees-classifier.pbtxt b/tensorflow/tools/api/golden/tensorflow.estimator.-boosted-trees-classifier.pbtxt
deleted file mode 100644
index c23b04b..0000000
--- a/tensorflow/tools/api/golden/tensorflow.estimator.-boosted-trees-classifier.pbtxt
+++ /dev/null
@@ -1,58 +0,0 @@
-path: "tensorflow.estimator.BoostedTreesClassifier"
-tf_class {
- is_instance: "<class \'tensorflow.python.estimator.canned.boosted_trees.BoostedTreesClassifier\'>"
- is_instance: "<class \'tensorflow.python.estimator.estimator.Estimator\'>"
- is_instance: "<type \'object\'>"
- member {
- name: "config"
- mtype: "<type \'property\'>"
- }
- member {
- name: "model_dir"
- mtype: "<type \'property\'>"
- }
- member {
- name: "model_fn"
- mtype: "<type \'property\'>"
- }
- member {
- name: "params"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'feature_columns\', \'n_batches_per_layer\', \'model_dir\', \'n_classes\', \'weight_column\', \'label_vocabulary\', \'n_trees\', \'max_depth\', \'learning_rate\', \'l1_regularization\', \'l2_regularization\', \'tree_complexity\', \'min_node_weight\', \'config\', \'center_bias\', \'pruning_mode\'], varargs=None, keywords=None, defaults=[\'None\', \'<object object instance>\', \'None\', \'None\', \'100\', \'6\', \'0.1\', \'0.0\', \'0.0\', \'0.0\', \'0.0\', \'None\', \'False\', \'none\'], "
- }
- member_method {
- name: "eval_dir"
- argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "evaluate"
- argspec: "args=[\'self\', \'input_fn\', \'steps\', \'hooks\', \'checkpoint_path\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
- }
- member_method {
- name: "export_savedmodel"
- argspec: "args=[\'self\', \'export_dir_base\', \'serving_input_receiver_fn\', \'assets_extra\', \'as_text\', \'checkpoint_path\', \'strip_default_attrs\'], varargs=None, keywords=None, defaults=[\'None\', \'False\', \'None\', \'False\'], "
- }
- member_method {
- name: "get_variable_names"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_variable_value"
- argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "latest_checkpoint"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "predict"
- argspec: "args=[\'self\', \'input_fn\', \'predict_keys\', \'hooks\', \'checkpoint_path\', \'yield_single_examples\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\'], "
- }
- member_method {
- name: "train"
- argspec: "args=[\'self\', \'input_fn\', \'hooks\', \'steps\', \'max_steps\', \'saving_listeners\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
- }
-}
diff --git a/tensorflow/tools/api/golden/tensorflow.estimator.-boosted-trees-regressor.pbtxt b/tensorflow/tools/api/golden/tensorflow.estimator.-boosted-trees-regressor.pbtxt
deleted file mode 100644
index 6878d28..0000000
--- a/tensorflow/tools/api/golden/tensorflow.estimator.-boosted-trees-regressor.pbtxt
+++ /dev/null
@@ -1,58 +0,0 @@
-path: "tensorflow.estimator.BoostedTreesRegressor"
-tf_class {
- is_instance: "<class \'tensorflow.python.estimator.canned.boosted_trees.BoostedTreesRegressor\'>"
- is_instance: "<class \'tensorflow.python.estimator.estimator.Estimator\'>"
- is_instance: "<type \'object\'>"
- member {
- name: "config"
- mtype: "<type \'property\'>"
- }
- member {
- name: "model_dir"
- mtype: "<type \'property\'>"
- }
- member {
- name: "model_fn"
- mtype: "<type \'property\'>"
- }
- member {
- name: "params"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'feature_columns\', \'n_batches_per_layer\', \'model_dir\', \'label_dimension\', \'weight_column\', \'n_trees\', \'max_depth\', \'learning_rate\', \'l1_regularization\', \'l2_regularization\', \'tree_complexity\', \'min_node_weight\', \'config\', \'center_bias\', \'pruning_mode\'], varargs=None, keywords=None, defaults=[\'None\', \'<object object instance>\', \'None\', \'100\', \'6\', \'0.1\', \'0.0\', \'0.0\', \'0.0\', \'0.0\', \'None\', \'False\', \'none\'], "
- }
- member_method {
- name: "eval_dir"
- argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "evaluate"
- argspec: "args=[\'self\', \'input_fn\', \'steps\', \'hooks\', \'checkpoint_path\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
- }
- member_method {
- name: "export_savedmodel"
- argspec: "args=[\'self\', \'export_dir_base\', \'serving_input_receiver_fn\', \'assets_extra\', \'as_text\', \'checkpoint_path\', \'strip_default_attrs\'], varargs=None, keywords=None, defaults=[\'None\', \'False\', \'None\', \'False\'], "
- }
- member_method {
- name: "get_variable_names"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_variable_value"
- argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "latest_checkpoint"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "predict"
- argspec: "args=[\'self\', \'input_fn\', \'predict_keys\', \'hooks\', \'checkpoint_path\', \'yield_single_examples\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\'], "
- }
- member_method {
- name: "train"
- argspec: "args=[\'self\', \'input_fn\', \'hooks\', \'steps\', \'max_steps\', \'saving_listeners\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
- }
-}
diff --git a/tensorflow/tools/api/golden/tensorflow.estimator.-run-config.pbtxt b/tensorflow/tools/api/golden/tensorflow.estimator.-run-config.pbtxt
deleted file mode 100644
index bf1f94b..0000000
--- a/tensorflow/tools/api/golden/tensorflow.estimator.-run-config.pbtxt
+++ /dev/null
@@ -1,105 +0,0 @@
-path: "tensorflow.estimator.RunConfig"
-tf_class {
- is_instance: "<class \'tensorflow.python.estimator.run_config.RunConfig\'>"
- is_instance: "<type \'object\'>"
- member {
- name: "cluster_spec"
- mtype: "<type \'property\'>"
- }
- member {
- name: "device_fn"
- mtype: "<type \'property\'>"
- }
- member {
- name: "eval_distribute"
- mtype: "<type \'property\'>"
- }
- member {
- name: "evaluation_master"
- mtype: "<type \'property\'>"
- }
- member {
- name: "global_id_in_cluster"
- mtype: "<type \'property\'>"
- }
- member {
- name: "is_chief"
- mtype: "<type \'property\'>"
- }
- member {
- name: "keep_checkpoint_every_n_hours"
- mtype: "<type \'property\'>"
- }
- member {
- name: "keep_checkpoint_max"
- mtype: "<type \'property\'>"
- }
- member {
- name: "log_step_count_steps"
- mtype: "<type \'property\'>"
- }
- member {
- name: "master"
- mtype: "<type \'property\'>"
- }
- member {
- name: "model_dir"
- mtype: "<type \'property\'>"
- }
- member {
- name: "num_ps_replicas"
- mtype: "<type \'property\'>"
- }
- member {
- name: "num_worker_replicas"
- mtype: "<type \'property\'>"
- }
- member {
- name: "protocol"
- mtype: "<type \'property\'>"
- }
- member {
- name: "save_checkpoints_secs"
- mtype: "<type \'property\'>"
- }
- member {
- name: "save_checkpoints_steps"
- mtype: "<type \'property\'>"
- }
- member {
- name: "save_summary_steps"
- mtype: "<type \'property\'>"
- }
- member {
- name: "service"
- mtype: "<type \'property\'>"
- }
- member {
- name: "session_config"
- mtype: "<type \'property\'>"
- }
- member {
- name: "task_id"
- mtype: "<type \'property\'>"
- }
- member {
- name: "task_type"
- mtype: "<type \'property\'>"
- }
- member {
- name: "tf_random_seed"
- mtype: "<type \'property\'>"
- }
- member {
- name: "train_distribute"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'model_dir\', \'tf_random_seed\', \'save_summary_steps\', \'save_checkpoints_steps\', \'save_checkpoints_secs\', \'session_config\', \'keep_checkpoint_max\', \'keep_checkpoint_every_n_hours\', \'log_step_count_steps\', \'train_distribute\', \'device_fn\', \'protocol\', \'eval_distribute\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'100\', \'<object object instance>\', \'<object object instance>\', \'None\', \'5\', \'10000\', \'100\', \'None\', \'None\', \'None\', \'None\'], "
- }
- member_method {
- name: "replace"
- argspec: "args=[\'self\'], varargs=None, keywords=kwargs, defaults=None"
- }
-}
diff --git a/tensorflow/tools/api/golden/tensorflow.image.pbtxt b/tensorflow/tools/api/golden/tensorflow.image.pbtxt
deleted file mode 100644
index 5c46dc5..0000000
--- a/tensorflow/tools/api/golden/tensorflow.image.pbtxt
+++ /dev/null
@@ -1,251 +0,0 @@
-path: "tensorflow.image"
-tf_module {
- member {
- name: "ResizeMethod"
- mtype: "<type \'type\'>"
- }
- member_method {
- name: "adjust_brightness"
- argspec: "args=[\'image\', \'delta\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "adjust_contrast"
- argspec: "args=[\'images\', \'contrast_factor\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "adjust_gamma"
- argspec: "args=[\'image\', \'gamma\', \'gain\'], varargs=None, keywords=None, defaults=[\'1\', \'1\'], "
- }
- member_method {
- name: "adjust_hue"
- argspec: "args=[\'image\', \'delta\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "adjust_jpeg_quality"
- argspec: "args=[\'image\', \'jpeg_quality\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "adjust_saturation"
- argspec: "args=[\'image\', \'saturation_factor\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "central_crop"
- argspec: "args=[\'image\', \'central_fraction\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "convert_image_dtype"
- argspec: "args=[\'image\', \'dtype\', \'saturate\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], "
- }
- member_method {
- name: "crop_and_resize"
- argspec: "args=[\'image\', \'boxes\', \'box_ind\', \'crop_size\', \'method\', \'extrapolation_value\', \'name\'], varargs=None, keywords=None, defaults=[\'bilinear\', \'0\', \'None\'], "
- }
- member_method {
- name: "crop_to_bounding_box"
- argspec: "args=[\'image\', \'offset_height\', \'offset_width\', \'target_height\', \'target_width\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "decode_and_crop_jpeg"
- argspec: "args=[\'contents\', \'crop_window\', \'channels\', \'ratio\', \'fancy_upscaling\', \'try_recover_truncated\', \'acceptable_fraction\', \'dct_method\', \'name\'], varargs=None, keywords=None, defaults=[\'0\', \'1\', \'True\', \'False\', \'1\', \'\', \'None\'], "
- }
- member_method {
- name: "decode_bmp"
- argspec: "args=[\'contents\', \'channels\', \'name\'], varargs=None, keywords=None, defaults=[\'0\', \'None\'], "
- }
- member_method {
- name: "decode_gif"
- argspec: "args=[\'contents\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "decode_image"
- argspec: "args=[\'contents\', \'channels\', \'dtype\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \"<dtype: \'uint8\'>\", \'None\'], "
- }
- member_method {
- name: "decode_jpeg"
- argspec: "args=[\'contents\', \'channels\', \'ratio\', \'fancy_upscaling\', \'try_recover_truncated\', \'acceptable_fraction\', \'dct_method\', \'name\'], varargs=None, keywords=None, defaults=[\'0\', \'1\', \'True\', \'False\', \'1\', \'\', \'None\'], "
- }
- member_method {
- name: "decode_png"
- argspec: "args=[\'contents\', \'channels\', \'dtype\', \'name\'], varargs=None, keywords=None, defaults=[\'0\', \"<dtype: \'uint8\'>\", \'None\'], "
- }
- member_method {
- name: "draw_bounding_boxes"
- argspec: "args=[\'images\', \'boxes\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "encode_jpeg"
- argspec: "args=[\'image\', \'format\', \'quality\', \'progressive\', \'optimize_size\', \'chroma_downsampling\', \'density_unit\', \'x_density\', \'y_density\', \'xmp_metadata\', \'name\'], varargs=None, keywords=None, defaults=[\'\', \'95\', \'False\', \'False\', \'True\', \'in\', \'300\', \'300\', \'\', \'None\'], "
- }
- member_method {
- name: "encode_png"
- argspec: "args=[\'image\', \'compression\', \'name\'], varargs=None, keywords=None, defaults=[\'-1\', \'None\'], "
- }
- member_method {
- name: "extract_glimpse"
- argspec: "args=[\'input\', \'size\', \'offsets\', \'centered\', \'normalized\', \'uniform_noise\', \'name\'], varargs=None, keywords=None, defaults=[\'True\', \'True\', \'True\', \'None\'], "
- }
- member_method {
- name: "extract_image_patches"
- argspec: "args=[\'images\', \'ksizes\', \'strides\', \'rates\', \'padding\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "extract_jpeg_shape"
- argspec: "args=[\'contents\', \'output_type\', \'name\'], varargs=None, keywords=None, defaults=[\"<dtype: \'int32\'>\", \'None\'], "
- }
- member_method {
- name: "flip_left_right"
- argspec: "args=[\'image\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "flip_up_down"
- argspec: "args=[\'image\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "grayscale_to_rgb"
- argspec: "args=[\'images\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "hsv_to_rgb"
- argspec: "args=[\'images\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "image_gradients"
- argspec: "args=[\'image\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "is_jpeg"
- argspec: "args=[\'contents\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "non_max_suppression"
- argspec: "args=[\'boxes\', \'scores\', \'max_output_size\', \'iou_threshold\', \'score_threshold\', \'name\'], varargs=None, keywords=None, defaults=[\'0.5\', \'-inf\', \'None\'], "
- }
- member_method {
- name: "non_max_suppression_overlaps"
- argspec: "args=[\'overlaps\', \'scores\', \'max_output_size\', \'overlap_threshold\', \'score_threshold\', \'name\'], varargs=None, keywords=None, defaults=[\'0.5\', \'-inf\', \'None\'], "
- }
- member_method {
- name: "non_max_suppression_padded"
- argspec: "args=[\'boxes\', \'scores\', \'max_output_size\', \'iou_threshold\', \'score_threshold\', \'pad_to_max_output_size\', \'name\'], varargs=None, keywords=None, defaults=[\'0.5\', \'-inf\', \'False\', \'None\'], "
- }
- member_method {
- name: "pad_to_bounding_box"
- argspec: "args=[\'image\', \'offset_height\', \'offset_width\', \'target_height\', \'target_width\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "per_image_standardization"
- argspec: "args=[\'image\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "psnr"
- argspec: "args=[\'a\', \'b\', \'max_val\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "random_brightness"
- argspec: "args=[\'image\', \'max_delta\', \'seed\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "random_contrast"
- argspec: "args=[\'image\', \'lower\', \'upper\', \'seed\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "random_flip_left_right"
- argspec: "args=[\'image\', \'seed\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "random_flip_up_down"
- argspec: "args=[\'image\', \'seed\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "random_hue"
- argspec: "args=[\'image\', \'max_delta\', \'seed\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "random_jpeg_quality"
- argspec: "args=[\'image\', \'min_jpeg_quality\', \'max_jpeg_quality\', \'seed\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "random_saturation"
- argspec: "args=[\'image\', \'lower\', \'upper\', \'seed\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "resize_area"
- argspec: "args=[\'images\', \'size\', \'align_corners\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], "
- }
- member_method {
- name: "resize_bicubic"
- argspec: "args=[\'images\', \'size\', \'align_corners\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], "
- }
- member_method {
- name: "resize_bilinear"
- argspec: "args=[\'images\', \'size\', \'align_corners\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], "
- }
- member_method {
- name: "resize_image_with_crop_or_pad"
- argspec: "args=[\'image\', \'target_height\', \'target_width\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "resize_image_with_pad"
- argspec: "args=[\'image\', \'target_height\', \'target_width\', \'method\'], varargs=None, keywords=None, defaults=[\'0\'], "
- }
- member_method {
- name: "resize_images"
- argspec: "args=[\'images\', \'size\', \'method\', \'align_corners\', \'preserve_aspect_ratio\'], varargs=None, keywords=None, defaults=[\'0\', \'False\', \'False\'], "
- }
- member_method {
- name: "resize_nearest_neighbor"
- argspec: "args=[\'images\', \'size\', \'align_corners\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], "
- }
- member_method {
- name: "rgb_to_grayscale"
- argspec: "args=[\'images\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "rgb_to_hsv"
- argspec: "args=[\'images\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "rgb_to_yiq"
- argspec: "args=[\'images\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "rgb_to_yuv"
- argspec: "args=[\'images\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "rot90"
- argspec: "args=[\'image\', \'k\', \'name\'], varargs=None, keywords=None, defaults=[\'1\', \'None\'], "
- }
- member_method {
- name: "sample_distorted_bounding_box"
- argspec: "args=[\'image_size\', \'bounding_boxes\', \'seed\', \'seed2\', \'min_object_covered\', \'aspect_ratio_range\', \'area_range\', \'max_attempts\', \'use_image_if_no_bounding_boxes\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'0.1\', \'None\', \'None\', \'None\', \'None\', \'None\'], "
- }
- member_method {
- name: "sobel_edges"
- argspec: "args=[\'image\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "ssim"
- argspec: "args=[\'img1\', \'img2\', \'max_val\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "ssim_multiscale"
- argspec: "args=[\'img1\', \'img2\', \'max_val\', \'power_factors\'], varargs=None, keywords=None, defaults=[\'(0.0448, 0.2856, 0.3001, 0.2363, 0.1333)\'], "
- }
- member_method {
- name: "total_variation"
- argspec: "args=[\'images\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "transpose_image"
- argspec: "args=[\'image\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "yiq_to_rgb"
- argspec: "args=[\'images\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "yuv_to_rgb"
- argspec: "args=[\'images\'], varargs=None, keywords=None, defaults=None"
- }
-}
diff --git a/tensorflow/tools/api/golden/tensorflow.keras.-model.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.-model.pbtxt
deleted file mode 100644
index e579fe6..0000000
--- a/tensorflow/tools/api/golden/tensorflow.keras.-model.pbtxt
+++ /dev/null
@@ -1,268 +0,0 @@
-path: "tensorflow.keras.Model"
-tf_class {
- is_instance: "<class \'tensorflow.python.keras.engine.training.Model\'>"
- is_instance: "<class \'tensorflow.python.keras.engine.network.Network\'>"
- 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: "input_spec"
- mtype: "<type \'property\'>"
- }
- member {
- name: "layers"
- 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: "state_updates"
- mtype: "<type \'property\'>"
- }
- member {
- name: "stateful"
- 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: "uses_learning_phase"
- mtype: "<type \'property\'>"
- }
- member {
- name: "variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "weights"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "add_loss"
- argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=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\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], "
- }
- member_method {
- name: "add_weight"
- argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'synchronization\', \'aggregation\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'VariableSynchronization.AUTO\', \'VariableAggregation.NONE\', \'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\', \'training\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "compile"
- argspec: "args=[\'self\', \'optimizer\', \'loss\', \'metrics\', \'loss_weights\', \'sample_weight_mode\', \'weighted_metrics\', \'target_tensors\', \'distribute\'], varargs=None, keywords=kwargs, defaults=[\'None\', \'None\', \'None\', \'None\', \'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: "evaluate"
- argspec: "args=[\'self\', \'x\', \'y\', \'batch_size\', \'verbose\', \'sample_weight\', \'steps\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'1\', \'None\', \'None\'], "
- }
- member_method {
- name: "evaluate_generator"
- argspec: "args=[\'self\', \'generator\', \'steps\', \'max_queue_size\', \'workers\', \'use_multiprocessing\', \'verbose\'], varargs=None, keywords=None, defaults=[\'None\', \'10\', \'1\', \'False\', \'0\'], "
- }
- member_method {
- name: "fit"
- argspec: "args=[\'self\', \'x\', \'y\', \'batch_size\', \'epochs\', \'verbose\', \'callbacks\', \'validation_split\', \'validation_data\', \'shuffle\', \'class_weight\', \'sample_weight\', \'initial_epoch\', \'steps_per_epoch\', \'validation_steps\'], varargs=None, keywords=kwargs, defaults=[\'None\', \'None\', \'None\', \'1\', \'1\', \'None\', \'0.0\', \'None\', \'True\', \'None\', \'None\', \'0\', \'None\', \'None\'], "
- }
- member_method {
- name: "fit_generator"
- argspec: "args=[\'self\', \'generator\', \'steps_per_epoch\', \'epochs\', \'verbose\', \'callbacks\', \'validation_data\', \'validation_steps\', \'class_weight\', \'max_queue_size\', \'workers\', \'use_multiprocessing\', \'shuffle\', \'initial_epoch\'], varargs=None, keywords=None, defaults=[\'None\', \'1\', \'1\', \'None\', \'None\', \'None\', \'None\', \'10\', \'1\', \'False\', \'True\', \'0\'], "
- }
- member_method {
- name: "from_config"
- argspec: "args=[\'cls\', \'config\', \'custom_objects\'], 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_layer"
- argspec: "args=[\'self\', \'name\', \'index\'], varargs=None, keywords=None, defaults=[\'None\', \'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: "load_weights"
- argspec: "args=[\'self\', \'filepath\', \'by_name\'], varargs=None, keywords=None, defaults=[\'False\'], "
- }
- member_method {
- name: "predict"
- argspec: "args=[\'self\', \'x\', \'batch_size\', \'verbose\', \'steps\'], varargs=None, keywords=None, defaults=[\'None\', \'0\', \'None\'], "
- }
- member_method {
- name: "predict_generator"
- argspec: "args=[\'self\', \'generator\', \'steps\', \'max_queue_size\', \'workers\', \'use_multiprocessing\', \'verbose\'], varargs=None, keywords=None, defaults=[\'None\', \'10\', \'1\', \'False\', \'0\'], "
- }
- member_method {
- name: "predict_on_batch"
- argspec: "args=[\'self\', \'x\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "reset_states"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "save"
- argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'include_optimizer\'], varargs=None, keywords=None, defaults=[\'True\', \'True\'], "
- }
- member_method {
- name: "save_weights"
- argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'save_format\'], varargs=None, keywords=None, defaults=[\'True\', \'None\'], "
- }
- member_method {
- name: "set_weights"
- argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "summary"
- argspec: "args=[\'self\', \'line_length\', \'positions\', \'print_fn\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
- }
- member_method {
- name: "test_on_batch"
- argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "to_json"
- argspec: "args=[\'self\'], varargs=None, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "to_yaml"
- argspec: "args=[\'self\'], varargs=None, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "train_on_batch"
- argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'class_weight\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
- }
-}
diff --git a/tensorflow/tools/api/golden/tensorflow.keras.-sequential.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.-sequential.pbtxt
deleted file mode 100644
index 6f05cdd..0000000
--- a/tensorflow/tools/api/golden/tensorflow.keras.-sequential.pbtxt
+++ /dev/null
@@ -1,289 +0,0 @@
-path: "tensorflow.keras.Sequential"
-tf_class {
- is_instance: "<class \'tensorflow.python.keras.engine.sequential.Sequential\'>"
- is_instance: "<class \'tensorflow.python.keras.engine.training.Model\'>"
- is_instance: "<class \'tensorflow.python.keras.engine.network.Network\'>"
- 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: "input_spec"
- mtype: "<type \'property\'>"
- }
- member {
- name: "layers"
- 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: "state_updates"
- mtype: "<type \'property\'>"
- }
- member {
- name: "stateful"
- 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: "uses_learning_phase"
- mtype: "<type \'property\'>"
- }
- member {
- name: "variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "weights"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'layers\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "add"
- argspec: "args=[\'self\', \'layer\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "add_loss"
- argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=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\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], "
- }
- member_method {
- name: "add_weight"
- argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'synchronization\', \'aggregation\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'VariableSynchronization.AUTO\', \'VariableAggregation.NONE\', \'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\', \'training\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "compile"
- argspec: "args=[\'self\', \'optimizer\', \'loss\', \'metrics\', \'loss_weights\', \'sample_weight_mode\', \'weighted_metrics\', \'target_tensors\', \'distribute\'], varargs=None, keywords=kwargs, defaults=[\'None\', \'None\', \'None\', \'None\', \'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: "evaluate"
- argspec: "args=[\'self\', \'x\', \'y\', \'batch_size\', \'verbose\', \'sample_weight\', \'steps\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'1\', \'None\', \'None\'], "
- }
- member_method {
- name: "evaluate_generator"
- argspec: "args=[\'self\', \'generator\', \'steps\', \'max_queue_size\', \'workers\', \'use_multiprocessing\', \'verbose\'], varargs=None, keywords=None, defaults=[\'None\', \'10\', \'1\', \'False\', \'0\'], "
- }
- member_method {
- name: "fit"
- argspec: "args=[\'self\', \'x\', \'y\', \'batch_size\', \'epochs\', \'verbose\', \'callbacks\', \'validation_split\', \'validation_data\', \'shuffle\', \'class_weight\', \'sample_weight\', \'initial_epoch\', \'steps_per_epoch\', \'validation_steps\'], varargs=None, keywords=kwargs, defaults=[\'None\', \'None\', \'None\', \'1\', \'1\', \'None\', \'0.0\', \'None\', \'True\', \'None\', \'None\', \'0\', \'None\', \'None\'], "
- }
- member_method {
- name: "fit_generator"
- argspec: "args=[\'self\', \'generator\', \'steps_per_epoch\', \'epochs\', \'verbose\', \'callbacks\', \'validation_data\', \'validation_steps\', \'class_weight\', \'max_queue_size\', \'workers\', \'use_multiprocessing\', \'shuffle\', \'initial_epoch\'], varargs=None, keywords=None, defaults=[\'None\', \'1\', \'1\', \'None\', \'None\', \'None\', \'None\', \'10\', \'1\', \'False\', \'True\', \'0\'], "
- }
- member_method {
- name: "from_config"
- argspec: "args=[\'cls\', \'config\', \'custom_objects\'], 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_layer"
- argspec: "args=[\'self\', \'name\', \'index\'], varargs=None, keywords=None, defaults=[\'None\', \'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: "load_weights"
- argspec: "args=[\'self\', \'filepath\', \'by_name\'], varargs=None, keywords=None, defaults=[\'False\'], "
- }
- member_method {
- name: "pop"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "predict"
- argspec: "args=[\'self\', \'x\', \'batch_size\', \'verbose\', \'steps\'], varargs=None, keywords=None, defaults=[\'None\', \'0\', \'None\'], "
- }
- member_method {
- name: "predict_classes"
- argspec: "args=[\'self\', \'x\', \'batch_size\', \'verbose\'], varargs=None, keywords=None, defaults=[\'32\', \'0\'], "
- }
- member_method {
- name: "predict_generator"
- argspec: "args=[\'self\', \'generator\', \'steps\', \'max_queue_size\', \'workers\', \'use_multiprocessing\', \'verbose\'], varargs=None, keywords=None, defaults=[\'None\', \'10\', \'1\', \'False\', \'0\'], "
- }
- member_method {
- name: "predict_on_batch"
- argspec: "args=[\'self\', \'x\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "predict_proba"
- argspec: "args=[\'self\', \'x\', \'batch_size\', \'verbose\'], varargs=None, keywords=None, defaults=[\'32\', \'0\'], "
- }
- member_method {
- name: "reset_states"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "save"
- argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'include_optimizer\'], varargs=None, keywords=None, defaults=[\'True\', \'True\'], "
- }
- member_method {
- name: "save_weights"
- argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'save_format\'], varargs=None, keywords=None, defaults=[\'True\', \'None\'], "
- }
- member_method {
- name: "set_weights"
- argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "summary"
- argspec: "args=[\'self\', \'line_length\', \'positions\', \'print_fn\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
- }
- member_method {
- name: "symbolic_set_inputs"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "test_on_batch"
- argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "to_json"
- argspec: "args=[\'self\'], varargs=None, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "to_yaml"
- argspec: "args=[\'self\'], varargs=None, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "train_on_batch"
- argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'class_weight\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
- }
-}
diff --git a/tensorflow/tools/api/golden/tensorflow.keras.activations.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.activations.pbtxt
deleted file mode 100644
index 2e9de9e..0000000
--- a/tensorflow/tools/api/golden/tensorflow.keras.activations.pbtxt
+++ /dev/null
@@ -1,55 +0,0 @@
-path: "tensorflow.keras.activations"
-tf_module {
- member_method {
- name: "deserialize"
- argspec: "args=[\'name\', \'custom_objects\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "elu"
- argspec: "args=[\'x\', \'alpha\'], varargs=None, keywords=None, defaults=[\'1.0\'], "
- }
- member_method {
- name: "get"
- argspec: "args=[\'identifier\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "hard_sigmoid"
- argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "linear"
- argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "relu"
- argspec: "args=[\'x\', \'alpha\', \'max_value\', \'threshold\'], varargs=None, keywords=None, defaults=[\'0.0\', \'None\', \'0\'], "
- }
- member_method {
- name: "selu"
- argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "serialize"
- argspec: "args=[\'activation\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "sigmoid"
- argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "softmax"
- argspec: "args=[\'x\', \'axis\'], varargs=None, keywords=None, defaults=[\'-1\'], "
- }
- member_method {
- name: "softplus"
- argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "softsign"
- argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "tanh"
- argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None"
- }
-}
diff --git a/tensorflow/tools/api/golden/tensorflow.keras.models.-model.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.models.-model.pbtxt
deleted file mode 100644
index 56914e1..0000000
--- a/tensorflow/tools/api/golden/tensorflow.keras.models.-model.pbtxt
+++ /dev/null
@@ -1,268 +0,0 @@
-path: "tensorflow.keras.models.Model"
-tf_class {
- is_instance: "<class \'tensorflow.python.keras.engine.training.Model\'>"
- is_instance: "<class \'tensorflow.python.keras.engine.network.Network\'>"
- 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: "input_spec"
- mtype: "<type \'property\'>"
- }
- member {
- name: "layers"
- 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: "state_updates"
- mtype: "<type \'property\'>"
- }
- member {
- name: "stateful"
- 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: "uses_learning_phase"
- mtype: "<type \'property\'>"
- }
- member {
- name: "variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "weights"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "add_loss"
- argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=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\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], "
- }
- member_method {
- name: "add_weight"
- argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'synchronization\', \'aggregation\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'VariableSynchronization.AUTO\', \'VariableAggregation.NONE\', \'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\', \'training\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "compile"
- argspec: "args=[\'self\', \'optimizer\', \'loss\', \'metrics\', \'loss_weights\', \'sample_weight_mode\', \'weighted_metrics\', \'target_tensors\', \'distribute\'], varargs=None, keywords=kwargs, defaults=[\'None\', \'None\', \'None\', \'None\', \'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: "evaluate"
- argspec: "args=[\'self\', \'x\', \'y\', \'batch_size\', \'verbose\', \'sample_weight\', \'steps\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'1\', \'None\', \'None\'], "
- }
- member_method {
- name: "evaluate_generator"
- argspec: "args=[\'self\', \'generator\', \'steps\', \'max_queue_size\', \'workers\', \'use_multiprocessing\', \'verbose\'], varargs=None, keywords=None, defaults=[\'None\', \'10\', \'1\', \'False\', \'0\'], "
- }
- member_method {
- name: "fit"
- argspec: "args=[\'self\', \'x\', \'y\', \'batch_size\', \'epochs\', \'verbose\', \'callbacks\', \'validation_split\', \'validation_data\', \'shuffle\', \'class_weight\', \'sample_weight\', \'initial_epoch\', \'steps_per_epoch\', \'validation_steps\'], varargs=None, keywords=kwargs, defaults=[\'None\', \'None\', \'None\', \'1\', \'1\', \'None\', \'0.0\', \'None\', \'True\', \'None\', \'None\', \'0\', \'None\', \'None\'], "
- }
- member_method {
- name: "fit_generator"
- argspec: "args=[\'self\', \'generator\', \'steps_per_epoch\', \'epochs\', \'verbose\', \'callbacks\', \'validation_data\', \'validation_steps\', \'class_weight\', \'max_queue_size\', \'workers\', \'use_multiprocessing\', \'shuffle\', \'initial_epoch\'], varargs=None, keywords=None, defaults=[\'None\', \'1\', \'1\', \'None\', \'None\', \'None\', \'None\', \'10\', \'1\', \'False\', \'True\', \'0\'], "
- }
- member_method {
- name: "from_config"
- argspec: "args=[\'cls\', \'config\', \'custom_objects\'], 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_layer"
- argspec: "args=[\'self\', \'name\', \'index\'], varargs=None, keywords=None, defaults=[\'None\', \'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: "load_weights"
- argspec: "args=[\'self\', \'filepath\', \'by_name\'], varargs=None, keywords=None, defaults=[\'False\'], "
- }
- member_method {
- name: "predict"
- argspec: "args=[\'self\', \'x\', \'batch_size\', \'verbose\', \'steps\'], varargs=None, keywords=None, defaults=[\'None\', \'0\', \'None\'], "
- }
- member_method {
- name: "predict_generator"
- argspec: "args=[\'self\', \'generator\', \'steps\', \'max_queue_size\', \'workers\', \'use_multiprocessing\', \'verbose\'], varargs=None, keywords=None, defaults=[\'None\', \'10\', \'1\', \'False\', \'0\'], "
- }
- member_method {
- name: "predict_on_batch"
- argspec: "args=[\'self\', \'x\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "reset_states"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "save"
- argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'include_optimizer\'], varargs=None, keywords=None, defaults=[\'True\', \'True\'], "
- }
- member_method {
- name: "save_weights"
- argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'save_format\'], varargs=None, keywords=None, defaults=[\'True\', \'None\'], "
- }
- member_method {
- name: "set_weights"
- argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "summary"
- argspec: "args=[\'self\', \'line_length\', \'positions\', \'print_fn\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
- }
- member_method {
- name: "test_on_batch"
- argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "to_json"
- argspec: "args=[\'self\'], varargs=None, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "to_yaml"
- argspec: "args=[\'self\'], varargs=None, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "train_on_batch"
- argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'class_weight\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
- }
-}
diff --git a/tensorflow/tools/api/golden/tensorflow.keras.models.-sequential.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.models.-sequential.pbtxt
deleted file mode 100644
index 4c1c540..0000000
--- a/tensorflow/tools/api/golden/tensorflow.keras.models.-sequential.pbtxt
+++ /dev/null
@@ -1,289 +0,0 @@
-path: "tensorflow.keras.models.Sequential"
-tf_class {
- is_instance: "<class \'tensorflow.python.keras.engine.sequential.Sequential\'>"
- is_instance: "<class \'tensorflow.python.keras.engine.training.Model\'>"
- is_instance: "<class \'tensorflow.python.keras.engine.network.Network\'>"
- 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: "input_spec"
- mtype: "<type \'property\'>"
- }
- member {
- name: "layers"
- 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: "state_updates"
- mtype: "<type \'property\'>"
- }
- member {
- name: "stateful"
- 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: "uses_learning_phase"
- mtype: "<type \'property\'>"
- }
- member {
- name: "variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "weights"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'layers\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "add"
- argspec: "args=[\'self\', \'layer\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "add_loss"
- argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=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\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], "
- }
- member_method {
- name: "add_weight"
- argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'synchronization\', \'aggregation\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'VariableSynchronization.AUTO\', \'VariableAggregation.NONE\', \'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\', \'training\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "compile"
- argspec: "args=[\'self\', \'optimizer\', \'loss\', \'metrics\', \'loss_weights\', \'sample_weight_mode\', \'weighted_metrics\', \'target_tensors\', \'distribute\'], varargs=None, keywords=kwargs, defaults=[\'None\', \'None\', \'None\', \'None\', \'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: "evaluate"
- argspec: "args=[\'self\', \'x\', \'y\', \'batch_size\', \'verbose\', \'sample_weight\', \'steps\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'1\', \'None\', \'None\'], "
- }
- member_method {
- name: "evaluate_generator"
- argspec: "args=[\'self\', \'generator\', \'steps\', \'max_queue_size\', \'workers\', \'use_multiprocessing\', \'verbose\'], varargs=None, keywords=None, defaults=[\'None\', \'10\', \'1\', \'False\', \'0\'], "
- }
- member_method {
- name: "fit"
- argspec: "args=[\'self\', \'x\', \'y\', \'batch_size\', \'epochs\', \'verbose\', \'callbacks\', \'validation_split\', \'validation_data\', \'shuffle\', \'class_weight\', \'sample_weight\', \'initial_epoch\', \'steps_per_epoch\', \'validation_steps\'], varargs=None, keywords=kwargs, defaults=[\'None\', \'None\', \'None\', \'1\', \'1\', \'None\', \'0.0\', \'None\', \'True\', \'None\', \'None\', \'0\', \'None\', \'None\'], "
- }
- member_method {
- name: "fit_generator"
- argspec: "args=[\'self\', \'generator\', \'steps_per_epoch\', \'epochs\', \'verbose\', \'callbacks\', \'validation_data\', \'validation_steps\', \'class_weight\', \'max_queue_size\', \'workers\', \'use_multiprocessing\', \'shuffle\', \'initial_epoch\'], varargs=None, keywords=None, defaults=[\'None\', \'1\', \'1\', \'None\', \'None\', \'None\', \'None\', \'10\', \'1\', \'False\', \'True\', \'0\'], "
- }
- member_method {
- name: "from_config"
- argspec: "args=[\'cls\', \'config\', \'custom_objects\'], 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_layer"
- argspec: "args=[\'self\', \'name\', \'index\'], varargs=None, keywords=None, defaults=[\'None\', \'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: "load_weights"
- argspec: "args=[\'self\', \'filepath\', \'by_name\'], varargs=None, keywords=None, defaults=[\'False\'], "
- }
- member_method {
- name: "pop"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "predict"
- argspec: "args=[\'self\', \'x\', \'batch_size\', \'verbose\', \'steps\'], varargs=None, keywords=None, defaults=[\'None\', \'0\', \'None\'], "
- }
- member_method {
- name: "predict_classes"
- argspec: "args=[\'self\', \'x\', \'batch_size\', \'verbose\'], varargs=None, keywords=None, defaults=[\'32\', \'0\'], "
- }
- member_method {
- name: "predict_generator"
- argspec: "args=[\'self\', \'generator\', \'steps\', \'max_queue_size\', \'workers\', \'use_multiprocessing\', \'verbose\'], varargs=None, keywords=None, defaults=[\'None\', \'10\', \'1\', \'False\', \'0\'], "
- }
- member_method {
- name: "predict_on_batch"
- argspec: "args=[\'self\', \'x\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "predict_proba"
- argspec: "args=[\'self\', \'x\', \'batch_size\', \'verbose\'], varargs=None, keywords=None, defaults=[\'32\', \'0\'], "
- }
- member_method {
- name: "reset_states"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "save"
- argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'include_optimizer\'], varargs=None, keywords=None, defaults=[\'True\', \'True\'], "
- }
- member_method {
- name: "save_weights"
- argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'save_format\'], varargs=None, keywords=None, defaults=[\'True\', \'None\'], "
- }
- member_method {
- name: "set_weights"
- argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "summary"
- argspec: "args=[\'self\', \'line_length\', \'positions\', \'print_fn\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
- }
- member_method {
- name: "symbolic_set_inputs"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "test_on_batch"
- argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "to_json"
- argspec: "args=[\'self\'], varargs=None, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "to_yaml"
- argspec: "args=[\'self\'], varargs=None, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "train_on_batch"
- argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'class_weight\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
- }
-}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.-run-options.-experimental.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.-run-options.-experimental.pbtxt
index 537e73a..47b5b56 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.-run-options.-experimental.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.-run-options.-experimental.pbtxt
@@ -8,5 +8,11 @@
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/v1/tensorflow.-run-options.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.-run-options.pbtxt
index cec04a2..c0c2e7b 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.-run-options.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.-run-options.pbtxt
@@ -55,6 +55,12 @@
label: LABEL_OPTIONAL
type: TYPE_INT64
}
+ field {
+ name: "use_run_handler_pool"
+ number: 2
+ label: LABEL_OPTIONAL
+ type: TYPE_BOOL
+ }
}
enum_type {
name: "TraceLevel"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.-variable.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.-variable.pbtxt
index 05698b0..af7fc9d 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.-variable.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.-variable.pbtxt
@@ -1,5 +1,6 @@
path: "tensorflow.Variable"
tf_class {
+ is_instance: "<class \'tensorflow.python.ops.variables.VariableV1\'>"
is_instance: "<class \'tensorflow.python.ops.variables.Variable\'>"
is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
is_instance: "<type \'object\'>"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.data.-dataset.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.data.-dataset.pbtxt
index c3ba2db..825afb6 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.data.-dataset.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.data.-dataset.pbtxt
@@ -91,6 +91,10 @@
argspec: "args=[], varargs=args, keywords=None, defaults=None"
}
member_method {
+ name: "reduce"
+ argspec: "args=[\'self\', \'initial_state\', \'reduce_func\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
name: "repeat"
argspec: "args=[\'self\', \'count\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.data.-fixed-length-record-dataset.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.data.-fixed-length-record-dataset.pbtxt
index 3541671..cdad5f6 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.data.-fixed-length-record-dataset.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.data.-fixed-length-record-dataset.pbtxt
@@ -92,6 +92,10 @@
argspec: "args=[], varargs=args, keywords=None, defaults=None"
}
member_method {
+ name: "reduce"
+ argspec: "args=[\'self\', \'initial_state\', \'reduce_func\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
name: "repeat"
argspec: "args=[\'self\', \'count\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.data.-t-f-record-dataset.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.data.-t-f-record-dataset.pbtxt
index b113c18..df41bff 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.data.-t-f-record-dataset.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.data.-t-f-record-dataset.pbtxt
@@ -92,6 +92,10 @@
argspec: "args=[], varargs=args, keywords=None, defaults=None"
}
member_method {
+ name: "reduce"
+ argspec: "args=[\'self\', \'initial_state\', \'reduce_func\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
name: "repeat"
argspec: "args=[\'self\', \'count\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.data.-text-line-dataset.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.data.-text-line-dataset.pbtxt
index 7210bf5..028bcc2 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.data.-text-line-dataset.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.data.-text-line-dataset.pbtxt
@@ -92,6 +92,10 @@
argspec: "args=[], varargs=args, keywords=None, defaults=None"
}
member_method {
+ name: "reduce"
+ argspec: "args=[\'self\', \'initial_state\', \'reduce_func\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
name: "repeat"
argspec: "args=[\'self\', \'count\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-checkpoint-input-pipeline-hook.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-checkpoint-input-pipeline-hook.pbtxt
new file mode 100644
index 0000000..03c16cd
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-checkpoint-input-pipeline-hook.pbtxt
@@ -0,0 +1,30 @@
+path: "tensorflow.data.experimental.CheckpointInputPipelineHook"
+tf_class {
+ is_instance: "<class \'tensorflow.python.data.experimental.ops.iterator_ops.CheckpointInputPipelineHook\'>"
+ is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunHook\'>"
+ is_instance: "<type \'object\'>"
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'estimator\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "after_create_session"
+ argspec: "args=[\'self\', \'session\', \'coord\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "after_run"
+ argspec: "args=[\'self\', \'run_context\', \'run_values\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "before_run"
+ argspec: "args=[\'self\', \'run_context\'], 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\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-csv-dataset.__metaclass__.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-csv-dataset.__metaclass__.pbtxt
new file mode 100644
index 0000000..3eeaa1b
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-csv-dataset.__metaclass__.pbtxt
@@ -0,0 +1,14 @@
+path: "tensorflow.data.experimental.CsvDataset.__metaclass__"
+tf_class {
+ is_instance: "<class \'abc.ABCMeta\'>"
+ member_method {
+ name: "__init__"
+ }
+ member_method {
+ name: "mro"
+ }
+ member_method {
+ name: "register"
+ argspec: "args=[\'cls\', \'subclass\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-csv-dataset.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-csv-dataset.pbtxt
new file mode 100644
index 0000000..0c0405e
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-csv-dataset.pbtxt
@@ -0,0 +1,127 @@
+path: "tensorflow.data.experimental.CsvDataset"
+tf_class {
+ is_instance: "<class \'tensorflow.python.data.experimental.ops.readers.CsvDataset\'>"
+ is_instance: "<class \'tensorflow.python.data.ops.dataset_ops.DatasetSource\'>"
+ is_instance: "<class \'tensorflow.python.data.ops.dataset_ops.Dataset\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "output_classes"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "output_shapes"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "output_types"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'filenames\', \'record_defaults\', \'compression_type\', \'buffer_size\', \'header\', \'field_delim\', \'use_quote_delim\', \'na_value\', \'select_cols\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'False\', \',\', \'True\', \'\', \'None\'], "
+ }
+ member_method {
+ name: "apply"
+ argspec: "args=[\'self\', \'transformation_func\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "batch"
+ argspec: "args=[\'self\', \'batch_size\', \'drop_remainder\'], varargs=None, keywords=None, defaults=[\'False\'], "
+ }
+ member_method {
+ name: "cache"
+ argspec: "args=[\'self\', \'filename\'], varargs=None, keywords=None, defaults=[\'\'], "
+ }
+ member_method {
+ name: "concatenate"
+ argspec: "args=[\'self\', \'dataset\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "filter"
+ argspec: "args=[\'self\', \'predicate\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "flat_map"
+ argspec: "args=[\'self\', \'map_func\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "from_generator"
+ argspec: "args=[\'generator\', \'output_types\', \'output_shapes\', \'args\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "from_sparse_tensor_slices"
+ argspec: "args=[\'sparse_tensor\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "from_tensor_slices"
+ argspec: "args=[\'tensors\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "from_tensors"
+ argspec: "args=[\'tensors\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "interleave"
+ argspec: "args=[\'self\', \'map_func\', \'cycle_length\', \'block_length\', \'num_parallel_calls\'], varargs=None, keywords=None, defaults=[\'1\', \'None\'], "
+ }
+ member_method {
+ name: "list_files"
+ argspec: "args=[\'file_pattern\', \'shuffle\', \'seed\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "make_initializable_iterator"
+ argspec: "args=[\'self\', \'shared_name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "make_one_shot_iterator"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "map"
+ argspec: "args=[\'self\', \'map_func\', \'num_parallel_calls\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "padded_batch"
+ argspec: "args=[\'self\', \'batch_size\', \'padded_shapes\', \'padding_values\', \'drop_remainder\'], varargs=None, keywords=None, defaults=[\'None\', \'False\'], "
+ }
+ member_method {
+ name: "prefetch"
+ argspec: "args=[\'self\', \'buffer_size\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "range"
+ argspec: "args=[], varargs=args, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "reduce"
+ argspec: "args=[\'self\', \'initial_state\', \'reduce_func\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "repeat"
+ argspec: "args=[\'self\', \'count\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "shard"
+ argspec: "args=[\'self\', \'num_shards\', \'index\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "shuffle"
+ argspec: "args=[\'self\', \'buffer_size\', \'seed\', \'reshuffle_each_iteration\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "skip"
+ argspec: "args=[\'self\', \'count\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "take"
+ argspec: "args=[\'self\', \'count\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "window"
+ argspec: "args=[\'self\', \'size\', \'shift\', \'stride\', \'drop_remainder\'], varargs=None, keywords=None, defaults=[\'None\', \'1\', \'False\'], "
+ }
+ member_method {
+ name: "zip"
+ argspec: "args=[\'datasets\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-optional.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-optional.pbtxt
new file mode 100644
index 0000000..b4c9459
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-optional.pbtxt
@@ -0,0 +1,28 @@
+path: "tensorflow.data.experimental.Optional"
+tf_class {
+ is_instance: "<class \'tensorflow.python.data.ops.optional_ops.Optional\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "value_structure"
+ mtype: "<class \'abc.abstractproperty\'>"
+ }
+ member_method {
+ name: "__init__"
+ }
+ member_method {
+ name: "from_value"
+ argspec: "args=[\'value\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "get_value"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "has_value"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "none_from_structure"
+ argspec: "args=[\'value_structure\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-random-dataset.__metaclass__.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-random-dataset.__metaclass__.pbtxt
new file mode 100644
index 0000000..2991b12
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-random-dataset.__metaclass__.pbtxt
@@ -0,0 +1,14 @@
+path: "tensorflow.data.experimental.RandomDataset.__metaclass__"
+tf_class {
+ is_instance: "<class \'abc.ABCMeta\'>"
+ member_method {
+ name: "__init__"
+ }
+ member_method {
+ name: "mro"
+ }
+ member_method {
+ name: "register"
+ argspec: "args=[\'cls\', \'subclass\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-random-dataset.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-random-dataset.pbtxt
new file mode 100644
index 0000000..bce0be4
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-random-dataset.pbtxt
@@ -0,0 +1,127 @@
+path: "tensorflow.data.experimental.RandomDataset"
+tf_class {
+ is_instance: "<class \'tensorflow.python.data.experimental.ops.random_ops.RandomDataset\'>"
+ is_instance: "<class \'tensorflow.python.data.ops.dataset_ops.DatasetSource\'>"
+ is_instance: "<class \'tensorflow.python.data.ops.dataset_ops.Dataset\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "output_classes"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "output_shapes"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "output_types"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'seed\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "apply"
+ argspec: "args=[\'self\', \'transformation_func\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "batch"
+ argspec: "args=[\'self\', \'batch_size\', \'drop_remainder\'], varargs=None, keywords=None, defaults=[\'False\'], "
+ }
+ member_method {
+ name: "cache"
+ argspec: "args=[\'self\', \'filename\'], varargs=None, keywords=None, defaults=[\'\'], "
+ }
+ member_method {
+ name: "concatenate"
+ argspec: "args=[\'self\', \'dataset\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "filter"
+ argspec: "args=[\'self\', \'predicate\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "flat_map"
+ argspec: "args=[\'self\', \'map_func\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "from_generator"
+ argspec: "args=[\'generator\', \'output_types\', \'output_shapes\', \'args\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "from_sparse_tensor_slices"
+ argspec: "args=[\'sparse_tensor\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "from_tensor_slices"
+ argspec: "args=[\'tensors\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "from_tensors"
+ argspec: "args=[\'tensors\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "interleave"
+ argspec: "args=[\'self\', \'map_func\', \'cycle_length\', \'block_length\', \'num_parallel_calls\'], varargs=None, keywords=None, defaults=[\'1\', \'None\'], "
+ }
+ member_method {
+ name: "list_files"
+ argspec: "args=[\'file_pattern\', \'shuffle\', \'seed\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "make_initializable_iterator"
+ argspec: "args=[\'self\', \'shared_name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "make_one_shot_iterator"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "map"
+ argspec: "args=[\'self\', \'map_func\', \'num_parallel_calls\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "padded_batch"
+ argspec: "args=[\'self\', \'batch_size\', \'padded_shapes\', \'padding_values\', \'drop_remainder\'], varargs=None, keywords=None, defaults=[\'None\', \'False\'], "
+ }
+ member_method {
+ name: "prefetch"
+ argspec: "args=[\'self\', \'buffer_size\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "range"
+ argspec: "args=[], varargs=args, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "reduce"
+ argspec: "args=[\'self\', \'initial_state\', \'reduce_func\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "repeat"
+ argspec: "args=[\'self\', \'count\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "shard"
+ argspec: "args=[\'self\', \'num_shards\', \'index\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "shuffle"
+ argspec: "args=[\'self\', \'buffer_size\', \'seed\', \'reshuffle_each_iteration\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "skip"
+ argspec: "args=[\'self\', \'count\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "take"
+ argspec: "args=[\'self\', \'count\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "window"
+ argspec: "args=[\'self\', \'size\', \'shift\', \'stride\', \'drop_remainder\'], varargs=None, keywords=None, defaults=[\'None\', \'1\', \'False\'], "
+ }
+ member_method {
+ name: "zip"
+ argspec: "args=[\'datasets\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-reducer.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-reducer.pbtxt
new file mode 100644
index 0000000..6b477a8
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-reducer.pbtxt
@@ -0,0 +1,21 @@
+path: "tensorflow.data.experimental.Reducer"
+tf_class {
+ is_instance: "<class \'tensorflow.python.data.experimental.ops.grouping.Reducer\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "finalize_func"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "init_func"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "reduce_func"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'init_func\', \'reduce_func\', \'finalize_func\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-sql-dataset.__metaclass__.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-sql-dataset.__metaclass__.pbtxt
new file mode 100644
index 0000000..948e99e
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-sql-dataset.__metaclass__.pbtxt
@@ -0,0 +1,14 @@
+path: "tensorflow.data.experimental.SqlDataset.__metaclass__"
+tf_class {
+ is_instance: "<class \'abc.ABCMeta\'>"
+ member_method {
+ name: "__init__"
+ }
+ member_method {
+ name: "mro"
+ }
+ member_method {
+ name: "register"
+ argspec: "args=[\'cls\', \'subclass\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-sql-dataset.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-sql-dataset.pbtxt
new file mode 100644
index 0000000..8aeae92
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-sql-dataset.pbtxt
@@ -0,0 +1,127 @@
+path: "tensorflow.data.experimental.SqlDataset"
+tf_class {
+ is_instance: "<class \'tensorflow.python.data.experimental.ops.readers.SqlDataset\'>"
+ is_instance: "<class \'tensorflow.python.data.ops.dataset_ops.DatasetSource\'>"
+ is_instance: "<class \'tensorflow.python.data.ops.dataset_ops.Dataset\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "output_classes"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "output_shapes"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "output_types"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'driver_name\', \'data_source_name\', \'query\', \'output_types\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "apply"
+ argspec: "args=[\'self\', \'transformation_func\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "batch"
+ argspec: "args=[\'self\', \'batch_size\', \'drop_remainder\'], varargs=None, keywords=None, defaults=[\'False\'], "
+ }
+ member_method {
+ name: "cache"
+ argspec: "args=[\'self\', \'filename\'], varargs=None, keywords=None, defaults=[\'\'], "
+ }
+ member_method {
+ name: "concatenate"
+ argspec: "args=[\'self\', \'dataset\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "filter"
+ argspec: "args=[\'self\', \'predicate\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "flat_map"
+ argspec: "args=[\'self\', \'map_func\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "from_generator"
+ argspec: "args=[\'generator\', \'output_types\', \'output_shapes\', \'args\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "from_sparse_tensor_slices"
+ argspec: "args=[\'sparse_tensor\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "from_tensor_slices"
+ argspec: "args=[\'tensors\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "from_tensors"
+ argspec: "args=[\'tensors\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "interleave"
+ argspec: "args=[\'self\', \'map_func\', \'cycle_length\', \'block_length\', \'num_parallel_calls\'], varargs=None, keywords=None, defaults=[\'1\', \'None\'], "
+ }
+ member_method {
+ name: "list_files"
+ argspec: "args=[\'file_pattern\', \'shuffle\', \'seed\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "make_initializable_iterator"
+ argspec: "args=[\'self\', \'shared_name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "make_one_shot_iterator"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "map"
+ argspec: "args=[\'self\', \'map_func\', \'num_parallel_calls\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "padded_batch"
+ argspec: "args=[\'self\', \'batch_size\', \'padded_shapes\', \'padding_values\', \'drop_remainder\'], varargs=None, keywords=None, defaults=[\'None\', \'False\'], "
+ }
+ member_method {
+ name: "prefetch"
+ argspec: "args=[\'self\', \'buffer_size\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "range"
+ argspec: "args=[], varargs=args, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "reduce"
+ argspec: "args=[\'self\', \'initial_state\', \'reduce_func\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "repeat"
+ argspec: "args=[\'self\', \'count\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "shard"
+ argspec: "args=[\'self\', \'num_shards\', \'index\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "shuffle"
+ argspec: "args=[\'self\', \'buffer_size\', \'seed\', \'reshuffle_each_iteration\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "skip"
+ argspec: "args=[\'self\', \'count\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "take"
+ argspec: "args=[\'self\', \'count\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "window"
+ argspec: "args=[\'self\', \'size\', \'shift\', \'stride\', \'drop_remainder\'], varargs=None, keywords=None, defaults=[\'None\', \'1\', \'False\'], "
+ }
+ member_method {
+ name: "zip"
+ argspec: "args=[\'datasets\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-stats-aggregator.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-stats-aggregator.pbtxt
new file mode 100644
index 0000000..0bcc8cf
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-stats-aggregator.pbtxt
@@ -0,0 +1,13 @@
+path: "tensorflow.data.experimental.StatsAggregator"
+tf_class {
+ is_instance: "<class \'tensorflow.python.data.experimental.ops.stats_ops.StatsAggregator\'>"
+ is_instance: "<type \'object\'>"
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "get_summary"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-t-f-record-writer.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-t-f-record-writer.pbtxt
new file mode 100644
index 0000000..6f9d18a
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-t-f-record-writer.pbtxt
@@ -0,0 +1,13 @@
+path: "tensorflow.data.experimental.TFRecordWriter"
+tf_class {
+ is_instance: "<class \'tensorflow.python.data.experimental.ops.writers.TFRecordWriter\'>"
+ is_instance: "<type \'object\'>"
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'filename\', \'compression_type\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "write"
+ argspec: "args=[\'self\', \'dataset\'], 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
new file mode 100644
index 0000000..b14585f
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.pbtxt
@@ -0,0 +1,139 @@
+path: "tensorflow.data.experimental"
+tf_module {
+ member {
+ name: "CheckpointInputPipelineHook"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "CsvDataset"
+ mtype: "<class \'abc.ABCMeta\'>"
+ }
+ member {
+ name: "Optional"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "RandomDataset"
+ mtype: "<class \'abc.ABCMeta\'>"
+ }
+ member {
+ name: "Reducer"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "SqlDataset"
+ mtype: "<class \'abc.ABCMeta\'>"
+ }
+ member {
+ name: "StatsAggregator"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "TFRecordWriter"
+ mtype: "<type \'type\'>"
+ }
+ member_method {
+ name: "Counter"
+ argspec: "args=[\'start\', \'step\', \'dtype\'], varargs=None, keywords=None, defaults=[\'0\', \'1\', \"<dtype: \'int64\'>\"], "
+ }
+ member_method {
+ name: "bucket_by_sequence_length"
+ argspec: "args=[\'element_length_func\', \'bucket_boundaries\', \'bucket_batch_sizes\', \'padded_shapes\', \'padding_values\', \'pad_to_bucket_boundary\', \'no_padding\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'False\', \'False\'], "
+ }
+ member_method {
+ name: "choose_from_datasets"
+ argspec: "args=[\'datasets\', \'choice_dataset\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "copy_to_device"
+ argspec: "args=[\'target_device\', \'source_device\'], varargs=None, keywords=None, defaults=[\'/cpu:0\'], "
+ }
+ member_method {
+ name: "dense_to_sparse_batch"
+ argspec: "args=[\'batch_size\', \'row_shape\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "enumerate_dataset"
+ argspec: "args=[\'start\'], varargs=None, keywords=None, defaults=[\'0\'], "
+ }
+ member_method {
+ name: "get_next_as_optional"
+ argspec: "args=[\'iterator\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "get_single_element"
+ argspec: "args=[\'dataset\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "group_by_reducer"
+ argspec: "args=[\'key_func\', \'reducer\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "group_by_window"
+ argspec: "args=[\'key_func\', \'reduce_func\', \'window_size\', \'window_size_func\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "ignore_errors"
+ argspec: "args=[], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "latency_stats"
+ argspec: "args=[\'tag\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "make_batched_features_dataset"
+ argspec: "args=[\'file_pattern\', \'batch_size\', \'features\', \'reader\', \'label_key\', \'reader_args\', \'num_epochs\', \'shuffle\', \'shuffle_buffer_size\', \'shuffle_seed\', \'prefetch_buffer_size\', \'reader_num_threads\', \'parser_num_threads\', \'sloppy_ordering\', \'drop_final_batch\'], varargs=None, keywords=None, defaults=[\"<class \'tensorflow.python.data.ops.readers.TFRecordDataset\'>\", \'None\', \'None\', \'None\', \'True\', \'10000\', \'None\', \'-1\', \'1\', \'2\', \'False\', \'False\'], "
+ }
+ member_method {
+ name: "make_csv_dataset"
+ argspec: "args=[\'file_pattern\', \'batch_size\', \'column_names\', \'column_defaults\', \'label_name\', \'select_columns\', \'field_delim\', \'use_quote_delim\', \'na_value\', \'header\', \'num_epochs\', \'shuffle\', \'shuffle_buffer_size\', \'shuffle_seed\', \'prefetch_buffer_size\', \'num_parallel_reads\', \'sloppy\', \'num_rows_for_inference\', \'compression_type\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \',\', \'True\', \'\', \'True\', \'None\', \'True\', \'10000\', \'None\', \'-1\', \'1\', \'False\', \'100\', \'None\'], "
+ }
+ member_method {
+ name: "make_saveable_from_iterator"
+ argspec: "args=[\'iterator\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "map_and_batch"
+ argspec: "args=[\'map_func\', \'batch_size\', \'num_parallel_batches\', \'drop_remainder\', \'num_parallel_calls\'], varargs=None, keywords=None, defaults=[\'None\', \'False\', \'None\'], "
+ }
+ member_method {
+ name: "parallel_interleave"
+ argspec: "args=[\'map_func\', \'cycle_length\', \'block_length\', \'sloppy\', \'buffer_output_elements\', \'prefetch_input_elements\'], varargs=None, keywords=None, defaults=[\'1\', \'False\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "parse_example_dataset"
+ argspec: "args=[\'features\', \'num_parallel_calls\'], varargs=None, keywords=None, defaults=[\'1\'], "
+ }
+ member_method {
+ name: "prefetch_to_device"
+ argspec: "args=[\'device\', \'buffer_size\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "rejection_resample"
+ argspec: "args=[\'class_func\', \'target_dist\', \'initial_dist\', \'seed\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "sample_from_datasets"
+ argspec: "args=[\'datasets\', \'weights\', \'seed\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "scan"
+ argspec: "args=[\'initial_state\', \'scan_func\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "set_stats_aggregator"
+ argspec: "args=[\'stats_aggregator\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "shuffle_and_repeat"
+ argspec: "args=[\'buffer_size\', \'count\', \'seed\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "unbatch"
+ argspec: "args=[], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "unique"
+ argspec: "args=[], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.data.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.data.pbtxt
index 56fb270..e205157 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.data.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.data.pbtxt
@@ -20,4 +20,8 @@
name: "TextLineDataset"
mtype: "<class \'abc.ABCMeta\'>"
}
+ member {
+ name: "experimental"
+ mtype: "<type \'module\'>"
+ }
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.debugging.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.debugging.pbtxt
index d9efe97..ab6287f 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.debugging.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.debugging.pbtxt
@@ -1,6 +1,90 @@
path: "tensorflow.debugging"
tf_module {
member_method {
+ name: "Assert"
+ argspec: "args=[\'condition\', \'data\', \'summarize\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_all_finite"
+ argspec: "args=[\'t\', \'msg\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "assert_equal"
+ argspec: "args=[\'x\', \'y\', \'data\', \'summarize\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_greater"
+ argspec: "args=[\'x\', \'y\', \'data\', \'summarize\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_greater_equal"
+ argspec: "args=[\'x\', \'y\', \'data\', \'summarize\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_integer"
+ argspec: "args=[\'x\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_less"
+ argspec: "args=[\'x\', \'y\', \'data\', \'summarize\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_less_equal"
+ argspec: "args=[\'x\', \'y\', \'data\', \'summarize\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_near"
+ argspec: "args=[\'x\', \'y\', \'rtol\', \'atol\', \'data\', \'summarize\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_negative"
+ argspec: "args=[\'x\', \'data\', \'summarize\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_non_negative"
+ argspec: "args=[\'x\', \'data\', \'summarize\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_non_positive"
+ argspec: "args=[\'x\', \'data\', \'summarize\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_none_equal"
+ argspec: "args=[\'x\', \'y\', \'data\', \'summarize\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_positive"
+ argspec: "args=[\'x\', \'data\', \'summarize\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_proper_iterable"
+ argspec: "args=[\'values\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "assert_rank"
+ argspec: "args=[\'x\', \'rank\', \'data\', \'summarize\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_rank_at_least"
+ argspec: "args=[\'x\', \'rank\', \'data\', \'summarize\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_rank_in"
+ argspec: "args=[\'x\', \'ranks\', \'data\', \'summarize\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_same_float_dtype"
+ argspec: "args=[\'tensors\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_scalar"
+ argspec: "args=[\'tensor\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "assert_type"
+ argspec: "args=[\'tensor\', \'tf_type\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
name: "check_numerics"
argspec: "args=[\'tensor\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -16,4 +100,16 @@
name: "is_nan"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
+ member_method {
+ name: "is_non_decreasing"
+ argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "is_numeric_tensor"
+ argspec: "args=[\'tensor\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "is_strictly_increasing"
+ argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.dtypes.-d-type.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.dtypes.-d-type.pbtxt
new file mode 100644
index 0000000..423eca3
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.dtypes.-d-type.pbtxt
@@ -0,0 +1,77 @@
+path: "tensorflow.dtypes.DType"
+tf_class {
+ is_instance: "<class \'tensorflow.python.framework.dtypes.DType\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "as_datatype_enum"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "as_numpy_dtype"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "base_dtype"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "is_bool"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "is_complex"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "is_floating"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "is_integer"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "is_numpy_compatible"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "is_quantized"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "is_unsigned"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "limits"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "max"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "min"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "name"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "real_dtype"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "size"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'type_enum\'], 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.dtypes.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.dtypes.pbtxt
index 98e1fee..ea23fec 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.dtypes.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.dtypes.pbtxt
@@ -1,7 +1,27 @@
path: "tensorflow.dtypes"
tf_module {
+ member {
+ name: "DType"
+ mtype: "<type \'type\'>"
+ }
+ member_method {
+ name: "as_dtype"
+ argspec: "args=[\'type_value\'], varargs=None, keywords=None, defaults=None"
+ }
member_method {
name: "as_string"
argspec: "args=[\'input\', \'precision\', \'scientific\', \'shortest\', \'width\', \'fill\', \'name\'], varargs=None, keywords=None, defaults=[\'-1\', \'False\', \'False\', \'-1\', \'\', \'None\'], "
}
+ member_method {
+ name: "cast"
+ argspec: "args=[\'x\', \'dtype\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "complex"
+ argspec: "args=[\'real\', \'imag\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "saturate_cast"
+ argspec: "args=[\'value\', \'dtype\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.graph_util.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.graph_util.pbtxt
index eeabf84..162ee76 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.graph_util.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.graph_util.pbtxt
@@ -9,6 +9,10 @@
argspec: "args=[\'graph_def\', \'dest_nodes\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "import_graph_def"
+ argspec: "args=[\'graph_def\', \'input_map\', \'return_elements\', \'name\', \'op_dict\', \'producer_op_list\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
name: "must_run_on_cpu"
argspec: "args=[\'node\', \'pin_variables_on_cpu\'], varargs=None, keywords=None, defaults=[\'False\'], "
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.image.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.image.pbtxt
index 5c46dc5..0a231f1 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.image.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.image.pbtxt
@@ -149,6 +149,10 @@
argspec: "args=[\'image\', \'lower\', \'upper\', \'seed\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "random_crop"
+ argspec: "args=[\'value\', \'size\', \'seed\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
name: "random_flip_left_right"
argspec: "args=[\'image\', \'seed\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.initializers.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.initializers.pbtxt
index d499c67..19ca621 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.initializers.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.initializers.pbtxt
@@ -73,6 +73,10 @@
argspec: "args=[], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "tables_initializer"
+ argspec: "args=[\'name\'], varargs=None, keywords=None, defaults=[\'init_all_tables\'], "
+ }
+ member_method {
name: "variables"
argspec: "args=[\'var_list\', \'name\'], varargs=None, keywords=None, defaults=[\'init\'], "
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.io.-fixed-len-feature.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.io.-fixed-len-feature.pbtxt
new file mode 100644
index 0000000..cd0e51c
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.io.-fixed-len-feature.pbtxt
@@ -0,0 +1,27 @@
+path: "tensorflow.io.FixedLenFeature"
+tf_class {
+ is_instance: "<class \'tensorflow.python.ops.parsing_ops.FixedLenFeature\'>"
+ is_instance: "<class \'tensorflow.python.ops.parsing_ops.FixedLenFeature\'>"
+ is_instance: "<type \'tuple\'>"
+ member {
+ name: "default_value"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "dtype"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "shape"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ }
+ member_method {
+ name: "count"
+ }
+ member_method {
+ name: "index"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.io.-fixed-len-sequence-feature.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.io.-fixed-len-sequence-feature.pbtxt
new file mode 100644
index 0000000..8a38f25
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.io.-fixed-len-sequence-feature.pbtxt
@@ -0,0 +1,31 @@
+path: "tensorflow.io.FixedLenSequenceFeature"
+tf_class {
+ is_instance: "<class \'tensorflow.python.ops.parsing_ops.FixedLenSequenceFeature\'>"
+ is_instance: "<class \'tensorflow.python.ops.parsing_ops.FixedLenSequenceFeature\'>"
+ is_instance: "<type \'tuple\'>"
+ member {
+ name: "allow_missing"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "default_value"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "dtype"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "shape"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ }
+ member_method {
+ name: "count"
+ }
+ member_method {
+ name: "index"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.io.-padding-f-i-f-o-queue.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.io.-padding-f-i-f-o-queue.pbtxt
new file mode 100644
index 0000000..85306fd
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.io.-padding-f-i-f-o-queue.pbtxt
@@ -0,0 +1,66 @@
+path: "tensorflow.io.PaddingFIFOQueue"
+tf_class {
+ is_instance: "<class \'tensorflow.python.ops.data_flow_ops.PaddingFIFOQueue\'>"
+ is_instance: "<class \'tensorflow.python.ops.data_flow_ops.QueueBase\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "dtypes"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "name"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "names"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "queue_ref"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "shapes"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'capacity\', \'dtypes\', \'shapes\', \'names\', \'shared_name\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'padding_fifo_queue\'], "
+ }
+ member_method {
+ name: "close"
+ argspec: "args=[\'self\', \'cancel_pending_enqueues\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], "
+ }
+ member_method {
+ name: "dequeue"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "dequeue_many"
+ argspec: "args=[\'self\', \'n\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "dequeue_up_to"
+ argspec: "args=[\'self\', \'n\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "enqueue"
+ argspec: "args=[\'self\', \'vals\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "enqueue_many"
+ argspec: "args=[\'self\', \'vals\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "from_list"
+ argspec: "args=[\'index\', \'queues\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "is_closed"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "size"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.io.-priority-queue.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.io.-priority-queue.pbtxt
new file mode 100644
index 0000000..02d8037
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.io.-priority-queue.pbtxt
@@ -0,0 +1,66 @@
+path: "tensorflow.io.PriorityQueue"
+tf_class {
+ is_instance: "<class \'tensorflow.python.ops.data_flow_ops.PriorityQueue\'>"
+ is_instance: "<class \'tensorflow.python.ops.data_flow_ops.QueueBase\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "dtypes"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "name"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "names"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "queue_ref"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "shapes"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'capacity\', \'types\', \'shapes\', \'names\', \'shared_name\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'priority_queue\'], "
+ }
+ member_method {
+ name: "close"
+ argspec: "args=[\'self\', \'cancel_pending_enqueues\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], "
+ }
+ member_method {
+ name: "dequeue"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "dequeue_many"
+ argspec: "args=[\'self\', \'n\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "dequeue_up_to"
+ argspec: "args=[\'self\', \'n\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "enqueue"
+ argspec: "args=[\'self\', \'vals\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "enqueue_many"
+ argspec: "args=[\'self\', \'vals\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "from_list"
+ argspec: "args=[\'index\', \'queues\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "is_closed"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "size"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.io.-queue-base.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.io.-queue-base.pbtxt
new file mode 100644
index 0000000..a30481a
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.io.-queue-base.pbtxt
@@ -0,0 +1,65 @@
+path: "tensorflow.io.QueueBase"
+tf_class {
+ is_instance: "<class \'tensorflow.python.ops.data_flow_ops.QueueBase\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "dtypes"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "name"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "names"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "queue_ref"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "shapes"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'dtypes\', \'shapes\', \'names\', \'queue_ref\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "close"
+ argspec: "args=[\'self\', \'cancel_pending_enqueues\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], "
+ }
+ member_method {
+ name: "dequeue"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "dequeue_many"
+ argspec: "args=[\'self\', \'n\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "dequeue_up_to"
+ argspec: "args=[\'self\', \'n\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "enqueue"
+ argspec: "args=[\'self\', \'vals\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "enqueue_many"
+ argspec: "args=[\'self\', \'vals\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "from_list"
+ argspec: "args=[\'index\', \'queues\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "is_closed"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "size"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.io.-random-shuffle-queue.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.io.-random-shuffle-queue.pbtxt
new file mode 100644
index 0000000..82cbf98
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.io.-random-shuffle-queue.pbtxt
@@ -0,0 +1,66 @@
+path: "tensorflow.io.RandomShuffleQueue"
+tf_class {
+ is_instance: "<class \'tensorflow.python.ops.data_flow_ops.RandomShuffleQueue\'>"
+ is_instance: "<class \'tensorflow.python.ops.data_flow_ops.QueueBase\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "dtypes"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "name"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "names"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "queue_ref"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "shapes"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'capacity\', \'min_after_dequeue\', \'dtypes\', \'shapes\', \'names\', \'seed\', \'shared_name\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'random_shuffle_queue\'], "
+ }
+ member_method {
+ name: "close"
+ argspec: "args=[\'self\', \'cancel_pending_enqueues\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], "
+ }
+ member_method {
+ name: "dequeue"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "dequeue_many"
+ argspec: "args=[\'self\', \'n\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "dequeue_up_to"
+ argspec: "args=[\'self\', \'n\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "enqueue"
+ argspec: "args=[\'self\', \'vals\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "enqueue_many"
+ argspec: "args=[\'self\', \'vals\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "from_list"
+ argspec: "args=[\'index\', \'queues\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "is_closed"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "size"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.io.-sparse-feature.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.io.-sparse-feature.pbtxt
new file mode 100644
index 0000000..216947b
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.io.-sparse-feature.pbtxt
@@ -0,0 +1,35 @@
+path: "tensorflow.io.SparseFeature"
+tf_class {
+ is_instance: "<class \'tensorflow.python.ops.parsing_ops.SparseFeature\'>"
+ is_instance: "<class \'tensorflow.python.ops.parsing_ops.SparseFeature\'>"
+ is_instance: "<type \'tuple\'>"
+ member {
+ name: "already_sorted"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "dtype"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "index_key"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "size"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "value_key"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ }
+ member_method {
+ name: "count"
+ }
+ member_method {
+ name: "index"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.io.-t-f-record-compression-type.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.io.-t-f-record-compression-type.pbtxt
new file mode 100644
index 0000000..b598f73
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.io.-t-f-record-compression-type.pbtxt
@@ -0,0 +1,20 @@
+path: "tensorflow.io.TFRecordCompressionType"
+tf_class {
+ is_instance: "<class \'tensorflow.python.lib.io.tf_record.TFRecordCompressionType\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "GZIP"
+ mtype: "<type \'int\'>"
+ }
+ member {
+ name: "NONE"
+ mtype: "<type \'int\'>"
+ }
+ member {
+ name: "ZLIB"
+ mtype: "<type \'int\'>"
+ }
+ member_method {
+ name: "__init__"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.io.-t-f-record-options.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.io.-t-f-record-options.pbtxt
new file mode 100644
index 0000000..bfbf37c
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.io.-t-f-record-options.pbtxt
@@ -0,0 +1,17 @@
+path: "tensorflow.io.TFRecordOptions"
+tf_class {
+ is_instance: "<class \'tensorflow.python.lib.io.tf_record.TFRecordOptions\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "compression_type_map"
+ mtype: "<type \'dict\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'compression_type\', \'flush_mode\', \'input_buffer_size\', \'output_buffer_size\', \'window_bits\', \'compression_level\', \'compression_method\', \'mem_level\', \'compression_strategy\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "get_compression_type_string"
+ argspec: "args=[\'cls\', \'options\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.io.-t-f-record-writer.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.io.-t-f-record-writer.pbtxt
new file mode 100644
index 0000000..6fd443f
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.io.-t-f-record-writer.pbtxt
@@ -0,0 +1,21 @@
+path: "tensorflow.io.TFRecordWriter"
+tf_class {
+ is_instance: "<class \'tensorflow.python.lib.io.tf_record.TFRecordWriter\'>"
+ is_instance: "<type \'object\'>"
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'path\', \'options\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "close"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "flush"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "write"
+ argspec: "args=[\'self\', \'record\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.io.-var-len-feature.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.io.-var-len-feature.pbtxt
new file mode 100644
index 0000000..fd835db
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.io.-var-len-feature.pbtxt
@@ -0,0 +1,19 @@
+path: "tensorflow.io.VarLenFeature"
+tf_class {
+ is_instance: "<class \'tensorflow.python.ops.parsing_ops.VarLenFeature\'>"
+ is_instance: "<class \'tensorflow.python.ops.parsing_ops.VarLenFeature\'>"
+ is_instance: "<type \'tuple\'>"
+ member {
+ name: "dtype"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ }
+ member_method {
+ name: "count"
+ }
+ member_method {
+ name: "index"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.io.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.io.pbtxt
index 8938cf2..dccf136 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.io.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.io.pbtxt
@@ -1,5 +1,49 @@
path: "tensorflow.io"
tf_module {
+ member {
+ name: "FixedLenFeature"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "FixedLenSequenceFeature"
+ 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\'>"
+ }
+ member {
+ name: "TFRecordCompressionType"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "TFRecordOptions"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "TFRecordWriter"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "VarLenFeature"
+ mtype: "<type \'type\'>"
+ }
member_method {
name: "decode_base64"
argspec: "args=[\'input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
@@ -9,6 +53,10 @@
argspec: "args=[\'bytes\', \'compression_type\', \'name\'], varargs=None, keywords=None, defaults=[\'\', \'None\'], "
}
member_method {
+ name: "decode_csv"
+ argspec: "args=[\'records\', \'record_defaults\', \'field_delim\', \'use_quote_delim\', \'name\', \'na_value\', \'select_cols\'], varargs=None, keywords=None, defaults=[\',\', \'True\', \'None\', \'\', \'None\'], "
+ }
+ member_method {
name: "decode_json_example"
argspec: "args=[\'json_examples\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -17,18 +65,38 @@
argspec: "args=[\'bytes\', \'out_type\', \'little_endian\', \'name\'], varargs=None, keywords=None, defaults=[\'True\', \'None\'], "
}
member_method {
+ name: "deserialize_many_sparse"
+ argspec: "args=[\'serialized_sparse\', \'dtype\', \'rank\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
name: "encode_base64"
argspec: "args=[\'input\', \'pad\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], "
}
member_method {
+ name: "match_filenames_once"
+ argspec: "args=[\'pattern\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "matching_files"
argspec: "args=[\'pattern\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "parse_example"
+ argspec: "args=[\'serialized\', \'features\', \'name\', \'example_names\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
name: "parse_sequence_example"
argspec: "args=[\'serialized\', \'context_features\', \'sequence_features\', \'example_names\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
}
member_method {
+ name: "parse_single_example"
+ argspec: "args=[\'serialized\', \'features\', \'name\', \'example_names\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "parse_single_sequence_example"
+ argspec: "args=[\'serialized\', \'context_features\', \'sequence_features\', \'example_name\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
name: "parse_tensor"
argspec: "args=[\'serialized\', \'out_type\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -37,7 +105,23 @@
argspec: "args=[\'filename\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "serialize_many_sparse"
+ argspec: "args=[\'sp_input\', \'name\', \'out_type\'], varargs=None, keywords=None, defaults=[\'None\', \"<dtype: \'string\'>\"], "
+ }
+ member_method {
+ name: "serialize_sparse"
+ argspec: "args=[\'sp_input\', \'name\', \'out_type\'], varargs=None, keywords=None, defaults=[\'None\', \"<dtype: \'string\'>\"], "
+ }
+ member_method {
+ name: "tf_record_iterator"
+ argspec: "args=[\'path\', \'options\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "write_file"
argspec: "args=[\'filename\', \'contents\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
+ member_method {
+ name: "write_graph"
+ argspec: "args=[\'graph_or_graph_def\', \'logdir\', \'name\', \'as_text\'], varargs=None, keywords=None, defaults=[\'True\'], "
+ }
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.linalg.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.linalg.pbtxt
index d979116..6ac95d9 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.linalg.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.linalg.pbtxt
@@ -109,10 +109,18 @@
argspec: "args=[\'num_rows\', \'num_columns\', \'batch_shape\', \'dtype\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \"<dtype: \'float32\'>\", \'None\'], "
}
member_method {
+ name: "global_norm"
+ argspec: "args=[\'t_list\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "inv"
argspec: "args=[\'input\', \'adjoint\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], "
}
member_method {
+ name: "l2_normalize"
+ argspec: "args=[\'x\', \'axis\', \'epsilon\', \'name\', \'dim\'], varargs=None, keywords=None, defaults=[\'None\', \'1e-12\', \'None\', \'None\'], "
+ }
+ member_method {
name: "logdet"
argspec: "args=[\'matrix\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -125,6 +133,10 @@
argspec: "args=[\'matrix\', \'rhs\', \'l2_regularizer\', \'fast\', \'name\'], varargs=None, keywords=None, defaults=[\'0.0\', \'True\', \'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\'], "
+ }
+ member_method {
name: "norm"
argspec: "args=[\'tensor\', \'ord\', \'axis\', \'keepdims\', \'name\', \'keep_dims\'], varargs=None, keywords=None, defaults=[\'euclidean\', \'None\', \'None\', \'None\', \'None\'], "
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.math.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.math.pbtxt
index a308c76..459b9e3 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.math.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.math.pbtxt
@@ -1,6 +1,14 @@
path: "tensorflow.math"
tf_module {
member_method {
+ name: "abs"
+ argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "accumulate_n"
+ argspec: "args=[\'inputs\', \'shape\', \'tensor_dtype\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
+ }
+ member_method {
name: "acos"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -13,6 +21,22 @@
argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "add_n"
+ argspec: "args=[\'inputs\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "angle"
+ argspec: "args=[\'input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "argmax"
+ argspec: "args=[\'input\', \'axis\', \'name\', \'dimension\', \'output_type\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \"<dtype: \'int64\'>\"], "
+ }
+ member_method {
+ name: "argmin"
+ argspec: "args=[\'input\', \'axis\', \'name\', \'dimension\', \'output_type\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \"<dtype: \'int64\'>\"], "
+ }
+ member_method {
name: "asin"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -53,10 +77,18 @@
argspec: "args=[\'a\', \'b\', \'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "bincount"
+ argspec: "args=[\'arr\', \'weights\', \'minlength\', \'maxlength\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \"<dtype: \'int32\'>\"], "
+ }
+ member_method {
name: "ceil"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "conj"
+ argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "cos"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -65,14 +97,34 @@
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "count_nonzero"
+ argspec: "args=[\'input_tensor\', \'axis\', \'keepdims\', \'dtype\', \'name\', \'reduction_indices\', \'keep_dims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \"<dtype: \'int64\'>\", \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "cumprod"
+ argspec: "args=[\'x\', \'axis\', \'exclusive\', \'reverse\', \'name\'], varargs=None, keywords=None, defaults=[\'0\', \'False\', \'False\', \'None\'], "
+ }
+ member_method {
+ name: "cumsum"
+ argspec: "args=[\'x\', \'axis\', \'exclusive\', \'reverse\', \'name\'], varargs=None, keywords=None, defaults=[\'0\', \'False\', \'False\', \'None\'], "
+ }
+ member_method {
name: "digamma"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "divide"
+ argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "equal"
argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "erf"
+ argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "erfc"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -89,6 +141,10 @@
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "floordiv"
+ argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "greater"
argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -105,10 +161,26 @@
argspec: "args=[\'a\', \'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "imag"
+ argspec: "args=[\'input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "in_top_k"
+ argspec: "args=[\'predictions\', \'targets\', \'k\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "invert_permutation"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "l2_normalize"
+ argspec: "args=[\'x\', \'axis\', \'epsilon\', \'name\', \'dim\'], varargs=None, keywords=None, defaults=[\'None\', \'1e-12\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "lbeta"
+ argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "less"
argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -129,6 +201,14 @@
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "log_sigmoid"
+ argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "log_softmax"
+ argspec: "args=[\'logits\', \'axis\', \'name\', \'dim\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
+ }
+ member_method {
name: "logical_and"
argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -141,6 +221,10 @@
argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "logical_xor"
+ argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'LogicalXor\'], "
+ }
+ member_method {
name: "maximum"
argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -149,6 +233,14 @@
argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "multiply"
+ argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "negative"
+ argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "not_equal"
argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -161,18 +253,66 @@
argspec: "args=[\'coeffs\', \'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "pow"
+ argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "real"
+ argspec: "args=[\'input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "reciprocal"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "reduce_all"
+ argspec: "args=[\'input_tensor\', \'axis\', \'keepdims\', \'name\', \'reduction_indices\', \'keep_dims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "reduce_any"
+ argspec: "args=[\'input_tensor\', \'axis\', \'keepdims\', \'name\', \'reduction_indices\', \'keep_dims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "reduce_logsumexp"
+ argspec: "args=[\'input_tensor\', \'axis\', \'keepdims\', \'name\', \'reduction_indices\', \'keep_dims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "reduce_max"
+ argspec: "args=[\'input_tensor\', \'axis\', \'keepdims\', \'name\', \'reduction_indices\', \'keep_dims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "reduce_mean"
+ argspec: "args=[\'input_tensor\', \'axis\', \'keepdims\', \'name\', \'reduction_indices\', \'keep_dims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "reduce_min"
+ argspec: "args=[\'input_tensor\', \'axis\', \'keepdims\', \'name\', \'reduction_indices\', \'keep_dims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "reduce_prod"
+ argspec: "args=[\'input_tensor\', \'axis\', \'keepdims\', \'name\', \'reduction_indices\', \'keep_dims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "reduce_sum"
+ argspec: "args=[\'input_tensor\', \'axis\', \'keepdims\', \'name\', \'reduction_indices\', \'keep_dims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
name: "rint"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "round"
+ argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "rsqrt"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "scalar_mul"
+ argspec: "args=[\'scalar\', \'x\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
name: "segment_max"
argspec: "args=[\'data\', \'segment_ids\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -193,6 +333,14 @@
argspec: "args=[\'data\', \'segment_ids\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "sigmoid"
+ argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "sign"
+ argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "sin"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -201,6 +349,10 @@
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "softmax"
+ argspec: "args=[\'logits\', \'axis\', \'name\', \'dim\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
+ }
+ member_method {
name: "softplus"
argspec: "args=[\'features\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -209,18 +361,46 @@
argspec: "args=[\'features\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "sqrt"
+ argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "square"
+ argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "squared_difference"
argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "subtract"
+ argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "tan"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "tanh"
+ argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "top_k"
+ argspec: "args=[\'input\', \'k\', \'sorted\', \'name\'], varargs=None, keywords=None, defaults=[\'1\', \'True\', \'None\'], "
+ }
+ member_method {
+ name: "truediv"
+ argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "unsorted_segment_max"
argspec: "args=[\'data\', \'segment_ids\', \'num_segments\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "unsorted_segment_mean"
+ argspec: "args=[\'data\', \'segment_ids\', \'num_segments\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "unsorted_segment_min"
argspec: "args=[\'data\', \'segment_ids\', \'num_segments\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -229,10 +409,26 @@
argspec: "args=[\'data\', \'segment_ids\', \'num_segments\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "unsorted_segment_sqrt_n"
+ argspec: "args=[\'data\', \'segment_ids\', \'num_segments\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "unsorted_segment_sum"
argspec: "args=[\'data\', \'segment_ids\', \'num_segments\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "xdivy"
+ argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "xlogy"
+ argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "zero_fraction"
+ argspec: "args=[\'value\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "zeta"
argspec: "args=[\'x\', \'q\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.nn.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.nn.pbtxt
index d9e5b0d..9b28ce5 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.nn.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.nn.pbtxt
@@ -101,6 +101,10 @@
argspec: "args=[\'labels\', \'inputs\', \'sequence_length\', \'preprocess_collapse_repeated\', \'ctc_merge_repeated\', \'ignore_longer_outputs_than_inputs\', \'time_major\'], varargs=None, keywords=None, defaults=[\'False\', \'True\', \'False\', \'True\'], "
}
member_method {
+ name: "depth_to_space"
+ argspec: "args=[\'input\', \'block_size\', \'name\', \'data_format\'], varargs=None, keywords=None, defaults=[\'None\', \'NHWC\'], "
+ }
+ member_method {
name: "depthwise_conv2d"
argspec: "args=[\'input\', \'filter\', \'strides\', \'padding\', \'rate\', \'name\', \'data_format\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
}
@@ -305,6 +309,14 @@
argspec: "args=[\'features\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "space_to_batch"
+ argspec: "args=[\'input\', \'paddings\', \'block_size\', \'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\'], "
+ }
+ member_method {
name: "sparse_softmax_cross_entropy_with_logits"
argspec: "args=[\'_sentinel\', \'labels\', \'logits\', \'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 503e145..a268529 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.pbtxt
@@ -497,6 +497,10 @@
mtype: "<class \'tensorflow.python.framework.dtypes.DType\'>"
}
member {
+ name: "random"
+ mtype: "<type \'module\'>"
+ }
+ member {
name: "random_normal_initializer"
mtype: "<type \'type\'>"
}
@@ -1745,6 +1749,10 @@
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "roll"
+ argspec: "args=[\'input\', \'shift\', \'axis\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
name: "round"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -2221,6 +2229,10 @@
argspec: "args=[\'max_shard_bytes\', \'axis\', \'bytes_per_string_element\', \'max_shards\'], varargs=None, keywords=None, defaults=[\'0\', \'16\', \'None\'], "
}
member_method {
+ name: "variable_creator_scope"
+ argspec: "args=[\'variable_creator\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
name: "variable_op_scope"
argspec: "args=[\'values\', \'name_or_scope\', \'default_name\', \'initializer\', \'regularizer\', \'caching_device\', \'partitioner\', \'custom_getter\', \'reuse\', \'dtype\', \'use_resource\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\'], "
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.quantization.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.quantization.pbtxt
index 6d865ef..77c92ae 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.quantization.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.quantization.pbtxt
@@ -29,6 +29,10 @@
argspec: "args=[\'gradients\', \'inputs\', \'min\', \'max\', \'num_bits\', \'narrow_range\', \'name\'], varargs=None, keywords=None, defaults=[\'8\', \'False\', \'None\'], "
}
member_method {
+ name: "quantize"
+ argspec: "args=[\'input\', \'min_range\', \'max_range\', \'T\', \'mode\', \'round_mode\', \'name\'], varargs=None, keywords=None, defaults=[\'MIN_COMBINED\', \'HALF_AWAY_FROM_ZERO\', \'None\'], "
+ }
+ member_method {
name: "quantized_concat"
argspec: "args=[\'concat_dim\', \'values\', \'input_mins\', \'input_maxes\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.random.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.random.pbtxt
new file mode 100644
index 0000000..a568dd4
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.random.pbtxt
@@ -0,0 +1,47 @@
+path: "tensorflow.random"
+tf_module {
+ member_method {
+ name: "gamma"
+ argspec: "args=[\'shape\', \'alpha\', \'beta\', \'dtype\', \'seed\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \"<dtype: \'float32\'>\", \'None\', \'None\'], "
+ }
+ member_method {
+ name: "get_seed"
+ argspec: "args=[\'op_seed\'], varargs=None, keywords=None, defaults=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\'], "
+ }
+ member_method {
+ name: "multinomial"
+ argspec: "args=[\'logits\', \'num_samples\', \'seed\', \'name\', \'output_dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "normal"
+ argspec: "args=[\'shape\', \'mean\', \'stddev\', \'dtype\', \'seed\', \'name\'], varargs=None, keywords=None, defaults=[\'0.0\', \'1.0\', \"<dtype: \'float32\'>\", \'None\', \'None\'], "
+ }
+ member_method {
+ name: "poisson"
+ argspec: "args=[\'lam\', \'shape\', \'dtype\', \'seed\', \'name\'], varargs=None, keywords=None, defaults=[\"<dtype: \'float32\'>\", \'None\', \'None\'], "
+ }
+ member_method {
+ name: "set_random_seed"
+ argspec: "args=[\'seed\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "shuffle"
+ argspec: "args=[\'value\', \'seed\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "truncated_normal"
+ argspec: "args=[\'shape\', \'mean\', \'stddev\', \'dtype\', \'seed\', \'name\'], varargs=None, keywords=None, defaults=[\'0.0\', \'1.0\', \"<dtype: \'float32\'>\", \'None\', \'None\'], "
+ }
+ member_method {
+ name: "uniform"
+ argspec: "args=[\'shape\', \'minval\', \'maxval\', \'dtype\', \'seed\', \'name\'], varargs=None, keywords=None, defaults=[\'0\', \'None\', \"<dtype: \'float32\'>\", \'None\', \'None\'], "
+ }
+ member_method {
+ name: "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.saved_model.-builder.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.saved_model.-builder.pbtxt
new file mode 100644
index 0000000..67457de
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.saved_model.-builder.pbtxt
@@ -0,0 +1,21 @@
+path: "tensorflow.saved_model.Builder"
+tf_class {
+ is_instance: "<class \'tensorflow.python.saved_model.builder_impl.SavedModelBuilder\'>"
+ is_instance: "<type \'object\'>"
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'export_dir\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "add_meta_graph"
+ argspec: "args=[\'self\', \'tags\', \'signature_def_map\', \'assets_collection\', \'legacy_init_op\', \'clear_devices\', \'main_op\', \'strip_default_attrs\', \'saver\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'False\', \'None\', \'False\', \'None\'], "
+ }
+ member_method {
+ name: "add_meta_graph_and_variables"
+ argspec: "args=[\'self\', \'sess\', \'tags\', \'signature_def_map\', \'assets_collection\', \'legacy_init_op\', \'clear_devices\', \'main_op\', \'strip_default_attrs\', \'saver\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'False\', \'None\', \'False\', \'None\'], "
+ }
+ member_method {
+ name: "save"
+ argspec: "args=[\'self\', \'as_text\'], varargs=None, keywords=None, defaults=[\'False\'], "
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.saved_model.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.saved_model.pbtxt
index e1a0385..3f4965f 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.saved_model.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.saved_model.pbtxt
@@ -1,6 +1,10 @@
path: "tensorflow.saved_model"
tf_module {
member {
+ name: "Builder"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "builder"
mtype: "<type \'module\'>"
}
@@ -33,6 +37,46 @@
mtype: "<type \'module\'>"
}
member_method {
+ name: "build_signature_def"
+ argspec: "args=[\'inputs\', \'outputs\', \'method_name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "build_tensor_info"
+ argspec: "args=[\'tensor\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "classification_signature_def"
+ argspec: "args=[\'examples\', \'classes\', \'scores\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "get_tensor_from_tensor_info"
+ argspec: "args=[\'tensor_info\', \'graph\', \'import_scope\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "is_valid_signature"
+ argspec: "args=[\'signature_def\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "load"
+ argspec: "args=[\'sess\', \'tags\', \'export_dir\', \'import_scope\'], varargs=None, keywords=saver_kwargs, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "main_op_with_restore"
+ argspec: "args=[\'restore_op_name\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "maybe_saved_model_directory"
+ argspec: "args=[\'export_dir\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "predict_signature_def"
+ argspec: "args=[\'inputs\', \'outputs\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "regression_signature_def"
+ argspec: "args=[\'examples\', \'predictions\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
name: "simple_save"
argspec: "args=[\'session\', \'export_dir\', \'inputs\', \'outputs\', \'legacy_init_op\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.sparse.-sparse-conditional-accumulator.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.sparse.-sparse-conditional-accumulator.pbtxt
new file mode 100644
index 0000000..cd97716
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.sparse.-sparse-conditional-accumulator.pbtxt
@@ -0,0 +1,46 @@
+path: "tensorflow.sparse.SparseConditionalAccumulator"
+tf_class {
+ is_instance: "<class \'tensorflow.python.ops.data_flow_ops.SparseConditionalAccumulator\'>"
+ is_instance: "<class \'tensorflow.python.ops.data_flow_ops.ConditionalAccumulatorBase\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "accumulator_ref"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "dtype"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "name"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'dtype\', \'shape\', \'shared_name\', \'name\', \'reduction_type\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'sparse_conditional_accumulator\', \'MEAN\'], "
+ }
+ member_method {
+ name: "apply_grad"
+ argspec: "args=[\'self\', \'grad_indices\', \'grad_values\', \'grad_shape\', \'local_step\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'0\', \'None\'], "
+ }
+ member_method {
+ name: "apply_indexed_slices_grad"
+ argspec: "args=[\'self\', \'grad\', \'local_step\', \'name\'], varargs=None, keywords=None, defaults=[\'0\', \'None\'], "
+ }
+ member_method {
+ name: "num_accumulated"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "set_global_step"
+ argspec: "args=[\'self\', \'new_global_step\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "take_grad"
+ argspec: "args=[\'self\', \'num_required\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "take_indexed_slices_grad"
+ argspec: "args=[\'self\', \'num_required\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.sparse.-sparse-tensor.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.sparse.-sparse-tensor.pbtxt
new file mode 100644
index 0000000..02e59a6
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.sparse.-sparse-tensor.pbtxt
@@ -0,0 +1,54 @@
+path: "tensorflow.sparse.SparseTensor"
+tf_class {
+ is_instance: "<class \'tensorflow.python.framework.sparse_tensor.SparseTensor\'>"
+ is_instance: "<class \'tensorflow.python.framework.ops._TensorLike\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "dense_shape"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "dtype"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "graph"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "indices"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "op"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "shape"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "values"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'indices\', \'values\', \'dense_shape\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "consumers"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "eval"
+ argspec: "args=[\'self\', \'feed_dict\', \'session\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "from_value"
+ argspec: "args=[\'cls\', \'sparse_tensor_value\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "get_shape"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.sparse.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.sparse.pbtxt
index ba9e651..32bd8d5f 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.sparse.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.sparse.pbtxt
@@ -1,5 +1,21 @@
path: "tensorflow.sparse"
tf_module {
+ member {
+ name: "SparseConditionalAccumulator"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "SparseTensor"
+ mtype: "<type \'type\'>"
+ }
+ member_method {
+ name: "add"
+ argspec: "args=[\'a\', \'b\', \'thresh\'], varargs=None, keywords=None, defaults=[\'0\'], "
+ }
+ member_method {
+ name: "concat"
+ argspec: "args=[\'axis\', \'sp_inputs\', \'name\', \'expand_nonconcat_dim\', \'concat_dim\'], varargs=None, keywords=None, defaults=[\'None\', \'False\', \'None\'], "
+ }
member_method {
name: "cross"
argspec: "args=[\'inputs\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
@@ -16,4 +32,100 @@
name: "eye"
argspec: "args=[\'num_rows\', \'num_columns\', \'dtype\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \"<dtype: \'float32\'>\", \'None\'], "
}
+ member_method {
+ name: "fill_empty_rows"
+ argspec: "args=[\'sp_input\', \'default_value\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "mask"
+ argspec: "args=[\'a\', \'mask_indices\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "matmul"
+ argspec: "args=[\'sp_a\', \'b\', \'adjoint_a\', \'adjoint_b\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'False\', \'None\'], "
+ }
+ member_method {
+ name: "maximum"
+ argspec: "args=[\'sp_a\', \'sp_b\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "merge"
+ argspec: "args=[\'sp_ids\', \'sp_values\', \'vocab_size\', \'name\', \'already_sorted\'], varargs=None, keywords=None, defaults=[\'None\', \'False\'], "
+ }
+ member_method {
+ name: "minimum"
+ argspec: "args=[\'sp_a\', \'sp_b\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "placeholder"
+ argspec: "args=[\'dtype\', \'shape\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "reduce_max"
+ argspec: "args=[\'sp_input\', \'axis\', \'keepdims\', \'reduction_axes\', \'keep_dims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "reduce_max_sparse"
+ argspec: "args=[\'sp_input\', \'axis\', \'keepdims\', \'reduction_axes\', \'keep_dims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "reduce_sum"
+ argspec: "args=[\'sp_input\', \'axis\', \'keepdims\', \'reduction_axes\', \'keep_dims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "reduce_sum_sparse"
+ argspec: "args=[\'sp_input\', \'axis\', \'keepdims\', \'reduction_axes\', \'keep_dims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "reorder"
+ argspec: "args=[\'sp_input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "reset_shape"
+ argspec: "args=[\'sp_input\', \'new_shape\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "reshape"
+ argspec: "args=[\'sp_input\', \'shape\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "retain"
+ argspec: "args=[\'sp_input\', \'to_retain\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "segment_mean"
+ argspec: "args=[\'data\', \'indices\', \'segment_ids\', \'name\', \'num_segments\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "segment_sqrt_n"
+ argspec: "args=[\'data\', \'indices\', \'segment_ids\', \'name\', \'num_segments\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "segment_sum"
+ argspec: "args=[\'data\', \'indices\', \'segment_ids\', \'name\', \'num_segments\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "slice"
+ argspec: "args=[\'sp_input\', \'start\', \'size\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "softmax"
+ argspec: "args=[\'sp_input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "split"
+ argspec: "args=[\'keyword_required\', \'sp_input\', \'num_split\', \'axis\', \'name\', \'split_dim\'], varargs=None, keywords=None, defaults=[\'KeywordRequired()\', \'None\', \'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "to_dense"
+ argspec: "args=[\'sp_input\', \'default_value\', \'validate_indices\', \'name\'], varargs=None, keywords=None, defaults=[\'0\', \'True\', \'None\'], "
+ }
+ member_method {
+ name: "to_indicator"
+ argspec: "args=[\'sp_input\', \'vocab_size\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "transpose"
+ argspec: "args=[\'sp_input\', \'perm\', \'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 c52581d..ebdaf57 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.strings.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.strings.pbtxt
@@ -13,6 +13,10 @@
argspec: "args=[\'input\', \'name\', \'unit\'], varargs=None, keywords=None, defaults=[\'None\', \'BYTE\'], "
}
member_method {
+ name: "reduce_join"
+ argspec: "args=[\'inputs\', \'axis\', \'keep_dims\', \'separator\', \'name\', \'reduction_indices\'], varargs=None, keywords=None, defaults=[\'None\', \'False\', \'\', \'None\', \'None\'], "
+ }
+ member_method {
name: "regex_full_match"
argspec: "args=[\'input\', \'pattern\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -48,4 +52,8 @@
name: "to_number"
argspec: "args=[\'string_tensor\', \'out_type\', \'name\'], varargs=None, keywords=None, defaults=[\"<dtype: \'float32\'>\", \'None\'], "
}
+ member_method {
+ name: "unicode_script"
+ argspec: "args=[\'input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.train.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.train.pbtxt
index 9f35395..45c81fd 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.train.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.train.pbtxt
@@ -273,6 +273,10 @@
argspec: "args=[\'checkpoint_prefix\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "confusion_matrix"
+ argspec: "args=[\'labels\', \'predictions\', \'num_classes\', \'dtype\', \'name\', \'weights\'], varargs=None, keywords=None, defaults=[\'None\', \"<dtype: \'int32\'>\", \'None\', \'None\'], "
+ }
+ 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/golden/v2/tensorflow.-run-options.-experimental.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.-run-options.-experimental.pbtxt
index 537e73a..47b5b56 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.-run-options.-experimental.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.-run-options.-experimental.pbtxt
@@ -8,5 +8,11 @@
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
index cec04a2..c0c2e7b 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.-run-options.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.-run-options.pbtxt
@@ -55,6 +55,12 @@
label: LABEL_OPTIONAL
type: TYPE_INT64
}
+ field {
+ name: "use_run_handler_pool"
+ number: 2
+ label: LABEL_OPTIONAL
+ type: TYPE_BOOL
+ }
}
enum_type {
name: "TraceLevel"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-variable-scope.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.-variable-scope.pbtxt
deleted file mode 100644
index c13eb7b..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.-variable-scope.pbtxt
+++ /dev/null
@@ -1,105 +0,0 @@
-path: "tensorflow.VariableScope"
-tf_class {
- is_instance: "<class \'tensorflow.python.ops.variable_scope.VariableScope\'>"
- is_instance: "<type \'object\'>"
- member {
- name: "caching_device"
- mtype: "<type \'property\'>"
- }
- member {
- name: "constraint"
- mtype: "<type \'property\'>"
- }
- member {
- name: "custom_getter"
- mtype: "<type \'property\'>"
- }
- member {
- name: "dtype"
- mtype: "<type \'property\'>"
- }
- member {
- name: "initializer"
- mtype: "<type \'property\'>"
- }
- member {
- name: "name"
- mtype: "<type \'property\'>"
- }
- member {
- name: "original_name_scope"
- mtype: "<type \'property\'>"
- }
- member {
- name: "partitioner"
- mtype: "<type \'property\'>"
- }
- member {
- name: "regularizer"
- mtype: "<type \'property\'>"
- }
- member {
- name: "reuse"
- mtype: "<type \'property\'>"
- }
- member {
- name: "use_resource"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'reuse\', \'name\', \'initializer\', \'regularizer\', \'caching_device\', \'partitioner\', \'custom_getter\', \'name_scope\', \'dtype\', \'use_resource\', \'constraint\'], varargs=None, keywords=None, defaults=[\'\', \'None\', \'None\', \'None\', \'None\', \'None\', \'\', \"<dtype: \'float32\'>\", \'None\', \'None\'], "
- }
- member_method {
- name: "get_collection"
- argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_variable"
- argspec: "args=[\'self\', \'var_store\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'reuse\', \'trainable\', \'collections\', \'caching_device\', \'partitioner\', \'validate_shape\', \'use_resource\', \'custom_getter\', \'constraint\', \'synchronization\', \'aggregation\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'VariableSynchronization.AUTO\', \'VariableAggregation.NONE\'], "
- }
- member_method {
- name: "global_variables"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "local_variables"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "reuse_variables"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "set_caching_device"
- argspec: "args=[\'self\', \'caching_device\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "set_custom_getter"
- argspec: "args=[\'self\', \'custom_getter\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "set_dtype"
- argspec: "args=[\'self\', \'dtype\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "set_initializer"
- argspec: "args=[\'self\', \'initializer\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "set_partitioner"
- argspec: "args=[\'self\', \'partitioner\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "set_regularizer"
- argspec: "args=[\'self\', \'regularizer\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "set_use_resource"
- argspec: "args=[\'self\', \'use_resource\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "trainable_variables"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-variable.-save-slice-info.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.-variable.-save-slice-info.pbtxt
deleted file mode 100644
index ac3ccd4..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.-variable.-save-slice-info.pbtxt
+++ /dev/null
@@ -1,17 +0,0 @@
-path: "tensorflow.Variable.SaveSliceInfo"
-tf_class {
- is_instance: "<class \'tensorflow.python.ops.variables.SaveSliceInfo\'>"
- is_instance: "<type \'object\'>"
- member {
- name: "spec"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'full_name\', \'full_shape\', \'var_offset\', \'var_shape\', \'save_slice_info_def\', \'import_scope\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\', \'None\'], "
- }
- member_method {
- name: "to_proto"
- argspec: "args=[\'self\', \'export_scope\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-variable.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.-variable.pbtxt
deleted file mode 100644
index 05698b0..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.-variable.pbtxt
+++ /dev/null
@@ -1,130 +0,0 @@
-path: "tensorflow.Variable"
-tf_class {
- is_instance: "<class \'tensorflow.python.ops.variables.Variable\'>"
- is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
- is_instance: "<type \'object\'>"
- member {
- name: "SaveSliceInfo"
- mtype: "<type \'type\'>"
- }
- member {
- name: "constraint"
- mtype: "<type \'property\'>"
- }
- member {
- name: "device"
- mtype: "<type \'property\'>"
- }
- member {
- name: "dtype"
- mtype: "<type \'property\'>"
- }
- member {
- name: "graph"
- mtype: "<type \'property\'>"
- }
- member {
- name: "initial_value"
- mtype: "<type \'property\'>"
- }
- member {
- name: "initializer"
- mtype: "<type \'property\'>"
- }
- member {
- name: "name"
- mtype: "<type \'property\'>"
- }
- member {
- name: "op"
- mtype: "<type \'property\'>"
- }
- member {
- name: "shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'initial_value\', \'trainable\', \'collections\', \'validate_shape\', \'caching_device\', \'name\', \'variable_def\', \'dtype\', \'expected_shape\', \'import_scope\', \'constraint\', \'use_resource\', \'synchronization\', \'aggregation\'], varargs=None, keywords=None, defaults=[\'None\', \'True\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'VariableSynchronization.AUTO\', \'VariableAggregation.NONE\'], "
- }
- member_method {
- name: "assign"
- argspec: "args=[\'self\', \'value\', \'use_locking\', \'name\', \'read_value\'], varargs=None, keywords=None, defaults=[\'False\', \'None\', \'True\'], "
- }
- member_method {
- name: "assign_add"
- argspec: "args=[\'self\', \'delta\', \'use_locking\', \'name\', \'read_value\'], varargs=None, keywords=None, defaults=[\'False\', \'None\', \'True\'], "
- }
- member_method {
- name: "assign_sub"
- argspec: "args=[\'self\', \'delta\', \'use_locking\', \'name\', \'read_value\'], varargs=None, keywords=None, defaults=[\'False\', \'None\', \'True\'], "
- }
- member_method {
- name: "count_up_to"
- argspec: "args=[\'self\', \'limit\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "eval"
- argspec: "args=[\'self\', \'session\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "from_proto"
- argspec: "args=[\'variable_def\', \'import_scope\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "get_shape"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "initialized_value"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "load"
- argspec: "args=[\'self\', \'value\', \'session\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "read_value"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "scatter_add"
- argspec: "args=[\'self\', \'sparse_delta\', \'use_locking\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], "
- }
- member_method {
- name: "scatter_nd_add"
- argspec: "args=[\'self\', \'indices\', \'updates\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "scatter_nd_sub"
- argspec: "args=[\'self\', \'indices\', \'updates\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "scatter_nd_update"
- argspec: "args=[\'self\', \'indices\', \'updates\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "scatter_sub"
- argspec: "args=[\'self\', \'sparse_delta\', \'use_locking\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], "
- }
- member_method {
- name: "scatter_update"
- argspec: "args=[\'self\', \'sparse_delta\', \'use_locking\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], "
- }
- member_method {
- name: "set_shape"
- argspec: "args=[\'self\', \'shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "to_proto"
- argspec: "args=[\'self\', \'export_scope\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "value"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.-dataset.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.-dataset.pbtxt
index c3ba2db..825afb6 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.data.-dataset.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.data.-dataset.pbtxt
@@ -91,6 +91,10 @@
argspec: "args=[], varargs=args, keywords=None, defaults=None"
}
member_method {
+ name: "reduce"
+ argspec: "args=[\'self\', \'initial_state\', \'reduce_func\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
name: "repeat"
argspec: "args=[\'self\', \'count\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.-fixed-length-record-dataset.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.-fixed-length-record-dataset.pbtxt
index 3541671..cdad5f6 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.data.-fixed-length-record-dataset.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.data.-fixed-length-record-dataset.pbtxt
@@ -92,6 +92,10 @@
argspec: "args=[], varargs=args, keywords=None, defaults=None"
}
member_method {
+ name: "reduce"
+ argspec: "args=[\'self\', \'initial_state\', \'reduce_func\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
name: "repeat"
argspec: "args=[\'self\', \'count\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.-t-f-record-dataset.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.-t-f-record-dataset.pbtxt
index b113c18..df41bff 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.data.-t-f-record-dataset.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.data.-t-f-record-dataset.pbtxt
@@ -92,6 +92,10 @@
argspec: "args=[], varargs=args, keywords=None, defaults=None"
}
member_method {
+ name: "reduce"
+ argspec: "args=[\'self\', \'initial_state\', \'reduce_func\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
name: "repeat"
argspec: "args=[\'self\', \'count\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.-text-line-dataset.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.-text-line-dataset.pbtxt
index 7210bf5..028bcc2 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.data.-text-line-dataset.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.data.-text-line-dataset.pbtxt
@@ -92,6 +92,10 @@
argspec: "args=[], varargs=args, keywords=None, defaults=None"
}
member_method {
+ name: "reduce"
+ argspec: "args=[\'self\', \'initial_state\', \'reduce_func\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
name: "repeat"
argspec: "args=[\'self\', \'count\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-checkpoint-input-pipeline-hook.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-checkpoint-input-pipeline-hook.pbtxt
new file mode 100644
index 0000000..03c16cd
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-checkpoint-input-pipeline-hook.pbtxt
@@ -0,0 +1,30 @@
+path: "tensorflow.data.experimental.CheckpointInputPipelineHook"
+tf_class {
+ is_instance: "<class \'tensorflow.python.data.experimental.ops.iterator_ops.CheckpointInputPipelineHook\'>"
+ is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunHook\'>"
+ is_instance: "<type \'object\'>"
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'estimator\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "after_create_session"
+ argspec: "args=[\'self\', \'session\', \'coord\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "after_run"
+ argspec: "args=[\'self\', \'run_context\', \'run_values\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "before_run"
+ argspec: "args=[\'self\', \'run_context\'], 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\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-csv-dataset.__metaclass__.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-csv-dataset.__metaclass__.pbtxt
new file mode 100644
index 0000000..3eeaa1b
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-csv-dataset.__metaclass__.pbtxt
@@ -0,0 +1,14 @@
+path: "tensorflow.data.experimental.CsvDataset.__metaclass__"
+tf_class {
+ is_instance: "<class \'abc.ABCMeta\'>"
+ member_method {
+ name: "__init__"
+ }
+ member_method {
+ name: "mro"
+ }
+ member_method {
+ name: "register"
+ argspec: "args=[\'cls\', \'subclass\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-csv-dataset.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-csv-dataset.pbtxt
new file mode 100644
index 0000000..0c0405e
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-csv-dataset.pbtxt
@@ -0,0 +1,127 @@
+path: "tensorflow.data.experimental.CsvDataset"
+tf_class {
+ is_instance: "<class \'tensorflow.python.data.experimental.ops.readers.CsvDataset\'>"
+ is_instance: "<class \'tensorflow.python.data.ops.dataset_ops.DatasetSource\'>"
+ is_instance: "<class \'tensorflow.python.data.ops.dataset_ops.Dataset\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "output_classes"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "output_shapes"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "output_types"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'filenames\', \'record_defaults\', \'compression_type\', \'buffer_size\', \'header\', \'field_delim\', \'use_quote_delim\', \'na_value\', \'select_cols\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'False\', \',\', \'True\', \'\', \'None\'], "
+ }
+ member_method {
+ name: "apply"
+ argspec: "args=[\'self\', \'transformation_func\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "batch"
+ argspec: "args=[\'self\', \'batch_size\', \'drop_remainder\'], varargs=None, keywords=None, defaults=[\'False\'], "
+ }
+ member_method {
+ name: "cache"
+ argspec: "args=[\'self\', \'filename\'], varargs=None, keywords=None, defaults=[\'\'], "
+ }
+ member_method {
+ name: "concatenate"
+ argspec: "args=[\'self\', \'dataset\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "filter"
+ argspec: "args=[\'self\', \'predicate\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "flat_map"
+ argspec: "args=[\'self\', \'map_func\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "from_generator"
+ argspec: "args=[\'generator\', \'output_types\', \'output_shapes\', \'args\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "from_sparse_tensor_slices"
+ argspec: "args=[\'sparse_tensor\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "from_tensor_slices"
+ argspec: "args=[\'tensors\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "from_tensors"
+ argspec: "args=[\'tensors\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "interleave"
+ argspec: "args=[\'self\', \'map_func\', \'cycle_length\', \'block_length\', \'num_parallel_calls\'], varargs=None, keywords=None, defaults=[\'1\', \'None\'], "
+ }
+ member_method {
+ name: "list_files"
+ argspec: "args=[\'file_pattern\', \'shuffle\', \'seed\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "make_initializable_iterator"
+ argspec: "args=[\'self\', \'shared_name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "make_one_shot_iterator"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "map"
+ argspec: "args=[\'self\', \'map_func\', \'num_parallel_calls\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "padded_batch"
+ argspec: "args=[\'self\', \'batch_size\', \'padded_shapes\', \'padding_values\', \'drop_remainder\'], varargs=None, keywords=None, defaults=[\'None\', \'False\'], "
+ }
+ member_method {
+ name: "prefetch"
+ argspec: "args=[\'self\', \'buffer_size\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "range"
+ argspec: "args=[], varargs=args, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "reduce"
+ argspec: "args=[\'self\', \'initial_state\', \'reduce_func\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "repeat"
+ argspec: "args=[\'self\', \'count\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "shard"
+ argspec: "args=[\'self\', \'num_shards\', \'index\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "shuffle"
+ argspec: "args=[\'self\', \'buffer_size\', \'seed\', \'reshuffle_each_iteration\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "skip"
+ argspec: "args=[\'self\', \'count\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "take"
+ argspec: "args=[\'self\', \'count\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "window"
+ argspec: "args=[\'self\', \'size\', \'shift\', \'stride\', \'drop_remainder\'], varargs=None, keywords=None, defaults=[\'None\', \'1\', \'False\'], "
+ }
+ member_method {
+ name: "zip"
+ argspec: "args=[\'datasets\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-optional.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-optional.pbtxt
new file mode 100644
index 0000000..b4c9459
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-optional.pbtxt
@@ -0,0 +1,28 @@
+path: "tensorflow.data.experimental.Optional"
+tf_class {
+ is_instance: "<class \'tensorflow.python.data.ops.optional_ops.Optional\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "value_structure"
+ mtype: "<class \'abc.abstractproperty\'>"
+ }
+ member_method {
+ name: "__init__"
+ }
+ member_method {
+ name: "from_value"
+ argspec: "args=[\'value\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "get_value"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "has_value"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "none_from_structure"
+ argspec: "args=[\'value_structure\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-random-dataset.__metaclass__.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-random-dataset.__metaclass__.pbtxt
new file mode 100644
index 0000000..2991b12
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-random-dataset.__metaclass__.pbtxt
@@ -0,0 +1,14 @@
+path: "tensorflow.data.experimental.RandomDataset.__metaclass__"
+tf_class {
+ is_instance: "<class \'abc.ABCMeta\'>"
+ member_method {
+ name: "__init__"
+ }
+ member_method {
+ name: "mro"
+ }
+ member_method {
+ name: "register"
+ argspec: "args=[\'cls\', \'subclass\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-random-dataset.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-random-dataset.pbtxt
new file mode 100644
index 0000000..bce0be4
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-random-dataset.pbtxt
@@ -0,0 +1,127 @@
+path: "tensorflow.data.experimental.RandomDataset"
+tf_class {
+ is_instance: "<class \'tensorflow.python.data.experimental.ops.random_ops.RandomDataset\'>"
+ is_instance: "<class \'tensorflow.python.data.ops.dataset_ops.DatasetSource\'>"
+ is_instance: "<class \'tensorflow.python.data.ops.dataset_ops.Dataset\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "output_classes"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "output_shapes"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "output_types"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'seed\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "apply"
+ argspec: "args=[\'self\', \'transformation_func\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "batch"
+ argspec: "args=[\'self\', \'batch_size\', \'drop_remainder\'], varargs=None, keywords=None, defaults=[\'False\'], "
+ }
+ member_method {
+ name: "cache"
+ argspec: "args=[\'self\', \'filename\'], varargs=None, keywords=None, defaults=[\'\'], "
+ }
+ member_method {
+ name: "concatenate"
+ argspec: "args=[\'self\', \'dataset\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "filter"
+ argspec: "args=[\'self\', \'predicate\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "flat_map"
+ argspec: "args=[\'self\', \'map_func\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "from_generator"
+ argspec: "args=[\'generator\', \'output_types\', \'output_shapes\', \'args\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "from_sparse_tensor_slices"
+ argspec: "args=[\'sparse_tensor\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "from_tensor_slices"
+ argspec: "args=[\'tensors\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "from_tensors"
+ argspec: "args=[\'tensors\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "interleave"
+ argspec: "args=[\'self\', \'map_func\', \'cycle_length\', \'block_length\', \'num_parallel_calls\'], varargs=None, keywords=None, defaults=[\'1\', \'None\'], "
+ }
+ member_method {
+ name: "list_files"
+ argspec: "args=[\'file_pattern\', \'shuffle\', \'seed\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "make_initializable_iterator"
+ argspec: "args=[\'self\', \'shared_name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "make_one_shot_iterator"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "map"
+ argspec: "args=[\'self\', \'map_func\', \'num_parallel_calls\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "padded_batch"
+ argspec: "args=[\'self\', \'batch_size\', \'padded_shapes\', \'padding_values\', \'drop_remainder\'], varargs=None, keywords=None, defaults=[\'None\', \'False\'], "
+ }
+ member_method {
+ name: "prefetch"
+ argspec: "args=[\'self\', \'buffer_size\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "range"
+ argspec: "args=[], varargs=args, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "reduce"
+ argspec: "args=[\'self\', \'initial_state\', \'reduce_func\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "repeat"
+ argspec: "args=[\'self\', \'count\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "shard"
+ argspec: "args=[\'self\', \'num_shards\', \'index\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "shuffle"
+ argspec: "args=[\'self\', \'buffer_size\', \'seed\', \'reshuffle_each_iteration\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "skip"
+ argspec: "args=[\'self\', \'count\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "take"
+ argspec: "args=[\'self\', \'count\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "window"
+ argspec: "args=[\'self\', \'size\', \'shift\', \'stride\', \'drop_remainder\'], varargs=None, keywords=None, defaults=[\'None\', \'1\', \'False\'], "
+ }
+ member_method {
+ name: "zip"
+ argspec: "args=[\'datasets\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-reducer.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-reducer.pbtxt
new file mode 100644
index 0000000..6b477a8
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-reducer.pbtxt
@@ -0,0 +1,21 @@
+path: "tensorflow.data.experimental.Reducer"
+tf_class {
+ is_instance: "<class \'tensorflow.python.data.experimental.ops.grouping.Reducer\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "finalize_func"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "init_func"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "reduce_func"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'init_func\', \'reduce_func\', \'finalize_func\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-sql-dataset.__metaclass__.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-sql-dataset.__metaclass__.pbtxt
new file mode 100644
index 0000000..948e99e
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-sql-dataset.__metaclass__.pbtxt
@@ -0,0 +1,14 @@
+path: "tensorflow.data.experimental.SqlDataset.__metaclass__"
+tf_class {
+ is_instance: "<class \'abc.ABCMeta\'>"
+ member_method {
+ name: "__init__"
+ }
+ member_method {
+ name: "mro"
+ }
+ member_method {
+ name: "register"
+ argspec: "args=[\'cls\', \'subclass\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-sql-dataset.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-sql-dataset.pbtxt
new file mode 100644
index 0000000..8aeae92
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-sql-dataset.pbtxt
@@ -0,0 +1,127 @@
+path: "tensorflow.data.experimental.SqlDataset"
+tf_class {
+ is_instance: "<class \'tensorflow.python.data.experimental.ops.readers.SqlDataset\'>"
+ is_instance: "<class \'tensorflow.python.data.ops.dataset_ops.DatasetSource\'>"
+ is_instance: "<class \'tensorflow.python.data.ops.dataset_ops.Dataset\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "output_classes"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "output_shapes"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "output_types"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'driver_name\', \'data_source_name\', \'query\', \'output_types\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "apply"
+ argspec: "args=[\'self\', \'transformation_func\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "batch"
+ argspec: "args=[\'self\', \'batch_size\', \'drop_remainder\'], varargs=None, keywords=None, defaults=[\'False\'], "
+ }
+ member_method {
+ name: "cache"
+ argspec: "args=[\'self\', \'filename\'], varargs=None, keywords=None, defaults=[\'\'], "
+ }
+ member_method {
+ name: "concatenate"
+ argspec: "args=[\'self\', \'dataset\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "filter"
+ argspec: "args=[\'self\', \'predicate\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "flat_map"
+ argspec: "args=[\'self\', \'map_func\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "from_generator"
+ argspec: "args=[\'generator\', \'output_types\', \'output_shapes\', \'args\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "from_sparse_tensor_slices"
+ argspec: "args=[\'sparse_tensor\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "from_tensor_slices"
+ argspec: "args=[\'tensors\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "from_tensors"
+ argspec: "args=[\'tensors\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "interleave"
+ argspec: "args=[\'self\', \'map_func\', \'cycle_length\', \'block_length\', \'num_parallel_calls\'], varargs=None, keywords=None, defaults=[\'1\', \'None\'], "
+ }
+ member_method {
+ name: "list_files"
+ argspec: "args=[\'file_pattern\', \'shuffle\', \'seed\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "make_initializable_iterator"
+ argspec: "args=[\'self\', \'shared_name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "make_one_shot_iterator"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "map"
+ argspec: "args=[\'self\', \'map_func\', \'num_parallel_calls\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "padded_batch"
+ argspec: "args=[\'self\', \'batch_size\', \'padded_shapes\', \'padding_values\', \'drop_remainder\'], varargs=None, keywords=None, defaults=[\'None\', \'False\'], "
+ }
+ member_method {
+ name: "prefetch"
+ argspec: "args=[\'self\', \'buffer_size\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "range"
+ argspec: "args=[], varargs=args, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "reduce"
+ argspec: "args=[\'self\', \'initial_state\', \'reduce_func\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "repeat"
+ argspec: "args=[\'self\', \'count\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "shard"
+ argspec: "args=[\'self\', \'num_shards\', \'index\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "shuffle"
+ argspec: "args=[\'self\', \'buffer_size\', \'seed\', \'reshuffle_each_iteration\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "skip"
+ argspec: "args=[\'self\', \'count\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "take"
+ argspec: "args=[\'self\', \'count\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "window"
+ argspec: "args=[\'self\', \'size\', \'shift\', \'stride\', \'drop_remainder\'], varargs=None, keywords=None, defaults=[\'None\', \'1\', \'False\'], "
+ }
+ member_method {
+ name: "zip"
+ argspec: "args=[\'datasets\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-stats-aggregator.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-stats-aggregator.pbtxt
new file mode 100644
index 0000000..0bcc8cf
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-stats-aggregator.pbtxt
@@ -0,0 +1,13 @@
+path: "tensorflow.data.experimental.StatsAggregator"
+tf_class {
+ is_instance: "<class \'tensorflow.python.data.experimental.ops.stats_ops.StatsAggregator\'>"
+ is_instance: "<type \'object\'>"
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "get_summary"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-t-f-record-writer.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-t-f-record-writer.pbtxt
new file mode 100644
index 0000000..6f9d18a
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-t-f-record-writer.pbtxt
@@ -0,0 +1,13 @@
+path: "tensorflow.data.experimental.TFRecordWriter"
+tf_class {
+ is_instance: "<class \'tensorflow.python.data.experimental.ops.writers.TFRecordWriter\'>"
+ is_instance: "<type \'object\'>"
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'filename\', \'compression_type\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "write"
+ argspec: "args=[\'self\', \'dataset\'], 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
new file mode 100644
index 0000000..b14585f
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.pbtxt
@@ -0,0 +1,139 @@
+path: "tensorflow.data.experimental"
+tf_module {
+ member {
+ name: "CheckpointInputPipelineHook"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "CsvDataset"
+ mtype: "<class \'abc.ABCMeta\'>"
+ }
+ member {
+ name: "Optional"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "RandomDataset"
+ mtype: "<class \'abc.ABCMeta\'>"
+ }
+ member {
+ name: "Reducer"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "SqlDataset"
+ mtype: "<class \'abc.ABCMeta\'>"
+ }
+ member {
+ name: "StatsAggregator"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "TFRecordWriter"
+ mtype: "<type \'type\'>"
+ }
+ member_method {
+ name: "Counter"
+ argspec: "args=[\'start\', \'step\', \'dtype\'], varargs=None, keywords=None, defaults=[\'0\', \'1\', \"<dtype: \'int64\'>\"], "
+ }
+ member_method {
+ name: "bucket_by_sequence_length"
+ argspec: "args=[\'element_length_func\', \'bucket_boundaries\', \'bucket_batch_sizes\', \'padded_shapes\', \'padding_values\', \'pad_to_bucket_boundary\', \'no_padding\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'False\', \'False\'], "
+ }
+ member_method {
+ name: "choose_from_datasets"
+ argspec: "args=[\'datasets\', \'choice_dataset\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "copy_to_device"
+ argspec: "args=[\'target_device\', \'source_device\'], varargs=None, keywords=None, defaults=[\'/cpu:0\'], "
+ }
+ member_method {
+ name: "dense_to_sparse_batch"
+ argspec: "args=[\'batch_size\', \'row_shape\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "enumerate_dataset"
+ argspec: "args=[\'start\'], varargs=None, keywords=None, defaults=[\'0\'], "
+ }
+ member_method {
+ name: "get_next_as_optional"
+ argspec: "args=[\'iterator\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "get_single_element"
+ argspec: "args=[\'dataset\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "group_by_reducer"
+ argspec: "args=[\'key_func\', \'reducer\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "group_by_window"
+ argspec: "args=[\'key_func\', \'reduce_func\', \'window_size\', \'window_size_func\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "ignore_errors"
+ argspec: "args=[], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "latency_stats"
+ argspec: "args=[\'tag\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "make_batched_features_dataset"
+ argspec: "args=[\'file_pattern\', \'batch_size\', \'features\', \'reader\', \'label_key\', \'reader_args\', \'num_epochs\', \'shuffle\', \'shuffle_buffer_size\', \'shuffle_seed\', \'prefetch_buffer_size\', \'reader_num_threads\', \'parser_num_threads\', \'sloppy_ordering\', \'drop_final_batch\'], varargs=None, keywords=None, defaults=[\"<class \'tensorflow.python.data.ops.readers.TFRecordDataset\'>\", \'None\', \'None\', \'None\', \'True\', \'10000\', \'None\', \'-1\', \'1\', \'2\', \'False\', \'False\'], "
+ }
+ member_method {
+ name: "make_csv_dataset"
+ argspec: "args=[\'file_pattern\', \'batch_size\', \'column_names\', \'column_defaults\', \'label_name\', \'select_columns\', \'field_delim\', \'use_quote_delim\', \'na_value\', \'header\', \'num_epochs\', \'shuffle\', \'shuffle_buffer_size\', \'shuffle_seed\', \'prefetch_buffer_size\', \'num_parallel_reads\', \'sloppy\', \'num_rows_for_inference\', \'compression_type\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \',\', \'True\', \'\', \'True\', \'None\', \'True\', \'10000\', \'None\', \'-1\', \'1\', \'False\', \'100\', \'None\'], "
+ }
+ member_method {
+ name: "make_saveable_from_iterator"
+ argspec: "args=[\'iterator\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "map_and_batch"
+ argspec: "args=[\'map_func\', \'batch_size\', \'num_parallel_batches\', \'drop_remainder\', \'num_parallel_calls\'], varargs=None, keywords=None, defaults=[\'None\', \'False\', \'None\'], "
+ }
+ member_method {
+ name: "parallel_interleave"
+ argspec: "args=[\'map_func\', \'cycle_length\', \'block_length\', \'sloppy\', \'buffer_output_elements\', \'prefetch_input_elements\'], varargs=None, keywords=None, defaults=[\'1\', \'False\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "parse_example_dataset"
+ argspec: "args=[\'features\', \'num_parallel_calls\'], varargs=None, keywords=None, defaults=[\'1\'], "
+ }
+ member_method {
+ name: "prefetch_to_device"
+ argspec: "args=[\'device\', \'buffer_size\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "rejection_resample"
+ argspec: "args=[\'class_func\', \'target_dist\', \'initial_dist\', \'seed\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "sample_from_datasets"
+ argspec: "args=[\'datasets\', \'weights\', \'seed\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "scan"
+ argspec: "args=[\'initial_state\', \'scan_func\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "set_stats_aggregator"
+ argspec: "args=[\'stats_aggregator\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "shuffle_and_repeat"
+ argspec: "args=[\'buffer_size\', \'count\', \'seed\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "unbatch"
+ argspec: "args=[], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "unique"
+ argspec: "args=[], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.pbtxt
index 56fb270..e205157 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.data.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.data.pbtxt
@@ -20,4 +20,8 @@
name: "TextLineDataset"
mtype: "<class \'abc.ABCMeta\'>"
}
+ member {
+ name: "experimental"
+ mtype: "<type \'module\'>"
+ }
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.debugging.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.debugging.pbtxt
index d9efe97..ab6287f 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.debugging.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.debugging.pbtxt
@@ -1,6 +1,90 @@
path: "tensorflow.debugging"
tf_module {
member_method {
+ name: "Assert"
+ argspec: "args=[\'condition\', \'data\', \'summarize\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_all_finite"
+ argspec: "args=[\'t\', \'msg\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "assert_equal"
+ argspec: "args=[\'x\', \'y\', \'data\', \'summarize\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_greater"
+ argspec: "args=[\'x\', \'y\', \'data\', \'summarize\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_greater_equal"
+ argspec: "args=[\'x\', \'y\', \'data\', \'summarize\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_integer"
+ argspec: "args=[\'x\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_less"
+ argspec: "args=[\'x\', \'y\', \'data\', \'summarize\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_less_equal"
+ argspec: "args=[\'x\', \'y\', \'data\', \'summarize\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_near"
+ argspec: "args=[\'x\', \'y\', \'rtol\', \'atol\', \'data\', \'summarize\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_negative"
+ argspec: "args=[\'x\', \'data\', \'summarize\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_non_negative"
+ argspec: "args=[\'x\', \'data\', \'summarize\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_non_positive"
+ argspec: "args=[\'x\', \'data\', \'summarize\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_none_equal"
+ argspec: "args=[\'x\', \'y\', \'data\', \'summarize\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_positive"
+ argspec: "args=[\'x\', \'data\', \'summarize\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_proper_iterable"
+ argspec: "args=[\'values\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "assert_rank"
+ argspec: "args=[\'x\', \'rank\', \'data\', \'summarize\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_rank_at_least"
+ argspec: "args=[\'x\', \'rank\', \'data\', \'summarize\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_rank_in"
+ argspec: "args=[\'x\', \'ranks\', \'data\', \'summarize\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_same_float_dtype"
+ argspec: "args=[\'tensors\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "assert_scalar"
+ argspec: "args=[\'tensor\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "assert_type"
+ argspec: "args=[\'tensor\', \'tf_type\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
name: "check_numerics"
argspec: "args=[\'tensor\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -16,4 +100,16 @@
name: "is_nan"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
+ member_method {
+ name: "is_non_decreasing"
+ argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "is_numeric_tensor"
+ argspec: "args=[\'tensor\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "is_strictly_increasing"
+ argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.dtypes.-d-type.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.dtypes.-d-type.pbtxt
new file mode 100644
index 0000000..423eca3
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.dtypes.-d-type.pbtxt
@@ -0,0 +1,77 @@
+path: "tensorflow.dtypes.DType"
+tf_class {
+ is_instance: "<class \'tensorflow.python.framework.dtypes.DType\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "as_datatype_enum"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "as_numpy_dtype"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "base_dtype"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "is_bool"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "is_complex"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "is_floating"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "is_integer"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "is_numpy_compatible"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "is_quantized"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "is_unsigned"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "limits"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "max"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "min"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "name"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "real_dtype"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "size"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'type_enum\'], 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.dtypes.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.dtypes.pbtxt
index 98e1fee..ea23fec 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.dtypes.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.dtypes.pbtxt
@@ -1,7 +1,27 @@
path: "tensorflow.dtypes"
tf_module {
+ member {
+ name: "DType"
+ mtype: "<type \'type\'>"
+ }
+ member_method {
+ name: "as_dtype"
+ argspec: "args=[\'type_value\'], varargs=None, keywords=None, defaults=None"
+ }
member_method {
name: "as_string"
argspec: "args=[\'input\', \'precision\', \'scientific\', \'shortest\', \'width\', \'fill\', \'name\'], varargs=None, keywords=None, defaults=[\'-1\', \'False\', \'False\', \'-1\', \'\', \'None\'], "
}
+ member_method {
+ name: "cast"
+ argspec: "args=[\'x\', \'dtype\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "complex"
+ argspec: "args=[\'real\', \'imag\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "saturate_cast"
+ argspec: "args=[\'value\', \'dtype\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.graph_util.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.graph_util.pbtxt
index eeabf84..162ee76 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.graph_util.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.graph_util.pbtxt
@@ -9,6 +9,10 @@
argspec: "args=[\'graph_def\', \'dest_nodes\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "import_graph_def"
+ argspec: "args=[\'graph_def\', \'input_map\', \'return_elements\', \'name\', \'op_dict\', \'producer_op_list\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
name: "must_run_on_cpu"
argspec: "args=[\'node\', \'pin_variables_on_cpu\'], varargs=None, keywords=None, defaults=[\'False\'], "
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.image.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.image.pbtxt
index 5c46dc5..0a231f1 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.image.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.image.pbtxt
@@ -149,6 +149,10 @@
argspec: "args=[\'image\', \'lower\', \'upper\', \'seed\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "random_crop"
+ argspec: "args=[\'value\', \'size\', \'seed\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
name: "random_flip_left_right"
argspec: "args=[\'image\', \'seed\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.initializers.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.initializers.pbtxt
index d499c67..d491817 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.initializers.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.initializers.pbtxt
@@ -49,10 +49,6 @@
mtype: "<type \'type\'>"
}
member_method {
- name: "global_variables"
- argspec: "args=[], varargs=None, keywords=None, defaults=None"
- }
- member_method {
name: "he_normal"
argspec: "args=[\'seed\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -69,11 +65,7 @@
argspec: "args=[\'seed\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
- name: "local_variables"
- argspec: "args=[], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "variables"
- argspec: "args=[\'var_list\', \'name\'], varargs=None, keywords=None, defaults=[\'init\'], "
+ name: "tables_initializer"
+ argspec: "args=[\'name\'], varargs=None, keywords=None, defaults=[\'init_all_tables\'], "
}
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.io.-fixed-len-feature.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.io.-fixed-len-feature.pbtxt
new file mode 100644
index 0000000..cd0e51c
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.io.-fixed-len-feature.pbtxt
@@ -0,0 +1,27 @@
+path: "tensorflow.io.FixedLenFeature"
+tf_class {
+ is_instance: "<class \'tensorflow.python.ops.parsing_ops.FixedLenFeature\'>"
+ is_instance: "<class \'tensorflow.python.ops.parsing_ops.FixedLenFeature\'>"
+ is_instance: "<type \'tuple\'>"
+ member {
+ name: "default_value"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "dtype"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "shape"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ }
+ member_method {
+ name: "count"
+ }
+ member_method {
+ name: "index"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.io.-fixed-len-sequence-feature.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.io.-fixed-len-sequence-feature.pbtxt
new file mode 100644
index 0000000..8a38f25
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.io.-fixed-len-sequence-feature.pbtxt
@@ -0,0 +1,31 @@
+path: "tensorflow.io.FixedLenSequenceFeature"
+tf_class {
+ is_instance: "<class \'tensorflow.python.ops.parsing_ops.FixedLenSequenceFeature\'>"
+ is_instance: "<class \'tensorflow.python.ops.parsing_ops.FixedLenSequenceFeature\'>"
+ is_instance: "<type \'tuple\'>"
+ member {
+ name: "allow_missing"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "default_value"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "dtype"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "shape"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ }
+ member_method {
+ name: "count"
+ }
+ member_method {
+ name: "index"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.io.-padding-f-i-f-o-queue.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.io.-padding-f-i-f-o-queue.pbtxt
new file mode 100644
index 0000000..85306fd
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.io.-padding-f-i-f-o-queue.pbtxt
@@ -0,0 +1,66 @@
+path: "tensorflow.io.PaddingFIFOQueue"
+tf_class {
+ is_instance: "<class \'tensorflow.python.ops.data_flow_ops.PaddingFIFOQueue\'>"
+ is_instance: "<class \'tensorflow.python.ops.data_flow_ops.QueueBase\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "dtypes"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "name"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "names"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "queue_ref"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "shapes"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'capacity\', \'dtypes\', \'shapes\', \'names\', \'shared_name\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'padding_fifo_queue\'], "
+ }
+ member_method {
+ name: "close"
+ argspec: "args=[\'self\', \'cancel_pending_enqueues\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], "
+ }
+ member_method {
+ name: "dequeue"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "dequeue_many"
+ argspec: "args=[\'self\', \'n\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "dequeue_up_to"
+ argspec: "args=[\'self\', \'n\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "enqueue"
+ argspec: "args=[\'self\', \'vals\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "enqueue_many"
+ argspec: "args=[\'self\', \'vals\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "from_list"
+ argspec: "args=[\'index\', \'queues\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "is_closed"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "size"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.io.-priority-queue.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.io.-priority-queue.pbtxt
new file mode 100644
index 0000000..02d8037
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.io.-priority-queue.pbtxt
@@ -0,0 +1,66 @@
+path: "tensorflow.io.PriorityQueue"
+tf_class {
+ is_instance: "<class \'tensorflow.python.ops.data_flow_ops.PriorityQueue\'>"
+ is_instance: "<class \'tensorflow.python.ops.data_flow_ops.QueueBase\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "dtypes"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "name"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "names"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "queue_ref"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "shapes"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'capacity\', \'types\', \'shapes\', \'names\', \'shared_name\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'priority_queue\'], "
+ }
+ member_method {
+ name: "close"
+ argspec: "args=[\'self\', \'cancel_pending_enqueues\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], "
+ }
+ member_method {
+ name: "dequeue"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "dequeue_many"
+ argspec: "args=[\'self\', \'n\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "dequeue_up_to"
+ argspec: "args=[\'self\', \'n\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "enqueue"
+ argspec: "args=[\'self\', \'vals\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "enqueue_many"
+ argspec: "args=[\'self\', \'vals\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "from_list"
+ argspec: "args=[\'index\', \'queues\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "is_closed"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "size"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.io.-queue-base.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.io.-queue-base.pbtxt
new file mode 100644
index 0000000..a30481a
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.io.-queue-base.pbtxt
@@ -0,0 +1,65 @@
+path: "tensorflow.io.QueueBase"
+tf_class {
+ is_instance: "<class \'tensorflow.python.ops.data_flow_ops.QueueBase\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "dtypes"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "name"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "names"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "queue_ref"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "shapes"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'dtypes\', \'shapes\', \'names\', \'queue_ref\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "close"
+ argspec: "args=[\'self\', \'cancel_pending_enqueues\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], "
+ }
+ member_method {
+ name: "dequeue"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "dequeue_many"
+ argspec: "args=[\'self\', \'n\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "dequeue_up_to"
+ argspec: "args=[\'self\', \'n\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "enqueue"
+ argspec: "args=[\'self\', \'vals\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "enqueue_many"
+ argspec: "args=[\'self\', \'vals\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "from_list"
+ argspec: "args=[\'index\', \'queues\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "is_closed"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "size"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.io.-random-shuffle-queue.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.io.-random-shuffle-queue.pbtxt
new file mode 100644
index 0000000..82cbf98
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.io.-random-shuffle-queue.pbtxt
@@ -0,0 +1,66 @@
+path: "tensorflow.io.RandomShuffleQueue"
+tf_class {
+ is_instance: "<class \'tensorflow.python.ops.data_flow_ops.RandomShuffleQueue\'>"
+ is_instance: "<class \'tensorflow.python.ops.data_flow_ops.QueueBase\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "dtypes"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "name"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "names"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "queue_ref"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "shapes"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'capacity\', \'min_after_dequeue\', \'dtypes\', \'shapes\', \'names\', \'seed\', \'shared_name\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'random_shuffle_queue\'], "
+ }
+ member_method {
+ name: "close"
+ argspec: "args=[\'self\', \'cancel_pending_enqueues\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], "
+ }
+ member_method {
+ name: "dequeue"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "dequeue_many"
+ argspec: "args=[\'self\', \'n\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "dequeue_up_to"
+ argspec: "args=[\'self\', \'n\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "enqueue"
+ argspec: "args=[\'self\', \'vals\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "enqueue_many"
+ argspec: "args=[\'self\', \'vals\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "from_list"
+ argspec: "args=[\'index\', \'queues\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "is_closed"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "size"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.io.-sparse-feature.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.io.-sparse-feature.pbtxt
new file mode 100644
index 0000000..216947b
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.io.-sparse-feature.pbtxt
@@ -0,0 +1,35 @@
+path: "tensorflow.io.SparseFeature"
+tf_class {
+ is_instance: "<class \'tensorflow.python.ops.parsing_ops.SparseFeature\'>"
+ is_instance: "<class \'tensorflow.python.ops.parsing_ops.SparseFeature\'>"
+ is_instance: "<type \'tuple\'>"
+ member {
+ name: "already_sorted"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "dtype"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "index_key"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "size"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "value_key"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ }
+ member_method {
+ name: "count"
+ }
+ member_method {
+ name: "index"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.io.-t-f-record-compression-type.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.io.-t-f-record-compression-type.pbtxt
new file mode 100644
index 0000000..b598f73
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.io.-t-f-record-compression-type.pbtxt
@@ -0,0 +1,20 @@
+path: "tensorflow.io.TFRecordCompressionType"
+tf_class {
+ is_instance: "<class \'tensorflow.python.lib.io.tf_record.TFRecordCompressionType\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "GZIP"
+ mtype: "<type \'int\'>"
+ }
+ member {
+ name: "NONE"
+ mtype: "<type \'int\'>"
+ }
+ member {
+ name: "ZLIB"
+ mtype: "<type \'int\'>"
+ }
+ member_method {
+ name: "__init__"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.io.-t-f-record-options.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.io.-t-f-record-options.pbtxt
new file mode 100644
index 0000000..bfbf37c
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.io.-t-f-record-options.pbtxt
@@ -0,0 +1,17 @@
+path: "tensorflow.io.TFRecordOptions"
+tf_class {
+ is_instance: "<class \'tensorflow.python.lib.io.tf_record.TFRecordOptions\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "compression_type_map"
+ mtype: "<type \'dict\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'compression_type\', \'flush_mode\', \'input_buffer_size\', \'output_buffer_size\', \'window_bits\', \'compression_level\', \'compression_method\', \'mem_level\', \'compression_strategy\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "get_compression_type_string"
+ argspec: "args=[\'cls\', \'options\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.io.-t-f-record-writer.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.io.-t-f-record-writer.pbtxt
new file mode 100644
index 0000000..6fd443f
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.io.-t-f-record-writer.pbtxt
@@ -0,0 +1,21 @@
+path: "tensorflow.io.TFRecordWriter"
+tf_class {
+ is_instance: "<class \'tensorflow.python.lib.io.tf_record.TFRecordWriter\'>"
+ is_instance: "<type \'object\'>"
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'path\', \'options\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "close"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "flush"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "write"
+ argspec: "args=[\'self\', \'record\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.io.-var-len-feature.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.io.-var-len-feature.pbtxt
new file mode 100644
index 0000000..fd835db
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.io.-var-len-feature.pbtxt
@@ -0,0 +1,19 @@
+path: "tensorflow.io.VarLenFeature"
+tf_class {
+ is_instance: "<class \'tensorflow.python.ops.parsing_ops.VarLenFeature\'>"
+ is_instance: "<class \'tensorflow.python.ops.parsing_ops.VarLenFeature\'>"
+ is_instance: "<type \'tuple\'>"
+ member {
+ name: "dtype"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ }
+ member_method {
+ name: "count"
+ }
+ member_method {
+ name: "index"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.io.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.io.pbtxt
index 8938cf2..dccf136 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.io.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.io.pbtxt
@@ -1,5 +1,49 @@
path: "tensorflow.io"
tf_module {
+ member {
+ name: "FixedLenFeature"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "FixedLenSequenceFeature"
+ 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\'>"
+ }
+ member {
+ name: "TFRecordCompressionType"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "TFRecordOptions"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "TFRecordWriter"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "VarLenFeature"
+ mtype: "<type \'type\'>"
+ }
member_method {
name: "decode_base64"
argspec: "args=[\'input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
@@ -9,6 +53,10 @@
argspec: "args=[\'bytes\', \'compression_type\', \'name\'], varargs=None, keywords=None, defaults=[\'\', \'None\'], "
}
member_method {
+ name: "decode_csv"
+ argspec: "args=[\'records\', \'record_defaults\', \'field_delim\', \'use_quote_delim\', \'name\', \'na_value\', \'select_cols\'], varargs=None, keywords=None, defaults=[\',\', \'True\', \'None\', \'\', \'None\'], "
+ }
+ member_method {
name: "decode_json_example"
argspec: "args=[\'json_examples\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -17,18 +65,38 @@
argspec: "args=[\'bytes\', \'out_type\', \'little_endian\', \'name\'], varargs=None, keywords=None, defaults=[\'True\', \'None\'], "
}
member_method {
+ name: "deserialize_many_sparse"
+ argspec: "args=[\'serialized_sparse\', \'dtype\', \'rank\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
name: "encode_base64"
argspec: "args=[\'input\', \'pad\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], "
}
member_method {
+ name: "match_filenames_once"
+ argspec: "args=[\'pattern\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "matching_files"
argspec: "args=[\'pattern\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "parse_example"
+ argspec: "args=[\'serialized\', \'features\', \'name\', \'example_names\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
name: "parse_sequence_example"
argspec: "args=[\'serialized\', \'context_features\', \'sequence_features\', \'example_names\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
}
member_method {
+ name: "parse_single_example"
+ argspec: "args=[\'serialized\', \'features\', \'name\', \'example_names\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "parse_single_sequence_example"
+ argspec: "args=[\'serialized\', \'context_features\', \'sequence_features\', \'example_name\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
name: "parse_tensor"
argspec: "args=[\'serialized\', \'out_type\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -37,7 +105,23 @@
argspec: "args=[\'filename\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "serialize_many_sparse"
+ argspec: "args=[\'sp_input\', \'name\', \'out_type\'], varargs=None, keywords=None, defaults=[\'None\', \"<dtype: \'string\'>\"], "
+ }
+ member_method {
+ name: "serialize_sparse"
+ argspec: "args=[\'sp_input\', \'name\', \'out_type\'], varargs=None, keywords=None, defaults=[\'None\', \"<dtype: \'string\'>\"], "
+ }
+ member_method {
+ name: "tf_record_iterator"
+ argspec: "args=[\'path\', \'options\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "write_file"
argspec: "args=[\'filename\', \'contents\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
+ member_method {
+ name: "write_graph"
+ argspec: "args=[\'graph_or_graph_def\', \'logdir\', \'name\', \'as_text\'], varargs=None, keywords=None, defaults=[\'True\'], "
+ }
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.linalg.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.linalg.pbtxt
index d979116..6ac95d9 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.linalg.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.linalg.pbtxt
@@ -109,10 +109,18 @@
argspec: "args=[\'num_rows\', \'num_columns\', \'batch_shape\', \'dtype\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \"<dtype: \'float32\'>\", \'None\'], "
}
member_method {
+ name: "global_norm"
+ argspec: "args=[\'t_list\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "inv"
argspec: "args=[\'input\', \'adjoint\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], "
}
member_method {
+ name: "l2_normalize"
+ argspec: "args=[\'x\', \'axis\', \'epsilon\', \'name\', \'dim\'], varargs=None, keywords=None, defaults=[\'None\', \'1e-12\', \'None\', \'None\'], "
+ }
+ member_method {
name: "logdet"
argspec: "args=[\'matrix\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -125,6 +133,10 @@
argspec: "args=[\'matrix\', \'rhs\', \'l2_regularizer\', \'fast\', \'name\'], varargs=None, keywords=None, defaults=[\'0.0\', \'True\', \'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\'], "
+ }
+ member_method {
name: "norm"
argspec: "args=[\'tensor\', \'ord\', \'axis\', \'keepdims\', \'name\', \'keep_dims\'], varargs=None, keywords=None, defaults=[\'euclidean\', \'None\', \'None\', \'None\', \'None\'], "
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.math.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.math.pbtxt
index a308c76..459b9e3 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.math.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.math.pbtxt
@@ -1,6 +1,14 @@
path: "tensorflow.math"
tf_module {
member_method {
+ name: "abs"
+ argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "accumulate_n"
+ argspec: "args=[\'inputs\', \'shape\', \'tensor_dtype\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
+ }
+ member_method {
name: "acos"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -13,6 +21,22 @@
argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "add_n"
+ argspec: "args=[\'inputs\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "angle"
+ argspec: "args=[\'input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "argmax"
+ argspec: "args=[\'input\', \'axis\', \'name\', \'dimension\', \'output_type\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \"<dtype: \'int64\'>\"], "
+ }
+ member_method {
+ name: "argmin"
+ argspec: "args=[\'input\', \'axis\', \'name\', \'dimension\', \'output_type\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \"<dtype: \'int64\'>\"], "
+ }
+ member_method {
name: "asin"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -53,10 +77,18 @@
argspec: "args=[\'a\', \'b\', \'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "bincount"
+ argspec: "args=[\'arr\', \'weights\', \'minlength\', \'maxlength\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \"<dtype: \'int32\'>\"], "
+ }
+ member_method {
name: "ceil"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "conj"
+ argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "cos"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -65,14 +97,34 @@
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "count_nonzero"
+ argspec: "args=[\'input_tensor\', \'axis\', \'keepdims\', \'dtype\', \'name\', \'reduction_indices\', \'keep_dims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \"<dtype: \'int64\'>\", \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "cumprod"
+ argspec: "args=[\'x\', \'axis\', \'exclusive\', \'reverse\', \'name\'], varargs=None, keywords=None, defaults=[\'0\', \'False\', \'False\', \'None\'], "
+ }
+ member_method {
+ name: "cumsum"
+ argspec: "args=[\'x\', \'axis\', \'exclusive\', \'reverse\', \'name\'], varargs=None, keywords=None, defaults=[\'0\', \'False\', \'False\', \'None\'], "
+ }
+ member_method {
name: "digamma"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "divide"
+ argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "equal"
argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "erf"
+ argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "erfc"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -89,6 +141,10 @@
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "floordiv"
+ argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "greater"
argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -105,10 +161,26 @@
argspec: "args=[\'a\', \'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "imag"
+ argspec: "args=[\'input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "in_top_k"
+ argspec: "args=[\'predictions\', \'targets\', \'k\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "invert_permutation"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "l2_normalize"
+ argspec: "args=[\'x\', \'axis\', \'epsilon\', \'name\', \'dim\'], varargs=None, keywords=None, defaults=[\'None\', \'1e-12\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "lbeta"
+ argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "less"
argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -129,6 +201,14 @@
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "log_sigmoid"
+ argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "log_softmax"
+ argspec: "args=[\'logits\', \'axis\', \'name\', \'dim\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
+ }
+ member_method {
name: "logical_and"
argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -141,6 +221,10 @@
argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "logical_xor"
+ argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'LogicalXor\'], "
+ }
+ member_method {
name: "maximum"
argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -149,6 +233,14 @@
argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "multiply"
+ argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "negative"
+ argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "not_equal"
argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -161,18 +253,66 @@
argspec: "args=[\'coeffs\', \'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "pow"
+ argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "real"
+ argspec: "args=[\'input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "reciprocal"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "reduce_all"
+ argspec: "args=[\'input_tensor\', \'axis\', \'keepdims\', \'name\', \'reduction_indices\', \'keep_dims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "reduce_any"
+ argspec: "args=[\'input_tensor\', \'axis\', \'keepdims\', \'name\', \'reduction_indices\', \'keep_dims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "reduce_logsumexp"
+ argspec: "args=[\'input_tensor\', \'axis\', \'keepdims\', \'name\', \'reduction_indices\', \'keep_dims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "reduce_max"
+ argspec: "args=[\'input_tensor\', \'axis\', \'keepdims\', \'name\', \'reduction_indices\', \'keep_dims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "reduce_mean"
+ argspec: "args=[\'input_tensor\', \'axis\', \'keepdims\', \'name\', \'reduction_indices\', \'keep_dims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "reduce_min"
+ argspec: "args=[\'input_tensor\', \'axis\', \'keepdims\', \'name\', \'reduction_indices\', \'keep_dims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "reduce_prod"
+ argspec: "args=[\'input_tensor\', \'axis\', \'keepdims\', \'name\', \'reduction_indices\', \'keep_dims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "reduce_sum"
+ argspec: "args=[\'input_tensor\', \'axis\', \'keepdims\', \'name\', \'reduction_indices\', \'keep_dims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
name: "rint"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "round"
+ argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "rsqrt"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "scalar_mul"
+ argspec: "args=[\'scalar\', \'x\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
name: "segment_max"
argspec: "args=[\'data\', \'segment_ids\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -193,6 +333,14 @@
argspec: "args=[\'data\', \'segment_ids\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "sigmoid"
+ argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "sign"
+ argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "sin"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -201,6 +349,10 @@
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "softmax"
+ argspec: "args=[\'logits\', \'axis\', \'name\', \'dim\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
+ }
+ member_method {
name: "softplus"
argspec: "args=[\'features\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -209,18 +361,46 @@
argspec: "args=[\'features\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "sqrt"
+ argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "square"
+ argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "squared_difference"
argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "subtract"
+ argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "tan"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "tanh"
+ argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "top_k"
+ argspec: "args=[\'input\', \'k\', \'sorted\', \'name\'], varargs=None, keywords=None, defaults=[\'1\', \'True\', \'None\'], "
+ }
+ member_method {
+ name: "truediv"
+ argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "unsorted_segment_max"
argspec: "args=[\'data\', \'segment_ids\', \'num_segments\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "unsorted_segment_mean"
+ argspec: "args=[\'data\', \'segment_ids\', \'num_segments\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "unsorted_segment_min"
argspec: "args=[\'data\', \'segment_ids\', \'num_segments\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -229,10 +409,26 @@
argspec: "args=[\'data\', \'segment_ids\', \'num_segments\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "unsorted_segment_sqrt_n"
+ argspec: "args=[\'data\', \'segment_ids\', \'num_segments\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "unsorted_segment_sum"
argspec: "args=[\'data\', \'segment_ids\', \'num_segments\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "xdivy"
+ argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "xlogy"
+ argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "zero_fraction"
+ argspec: "args=[\'value\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "zeta"
argspec: "args=[\'x\', \'q\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.nn.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.nn.pbtxt
index d9e5b0d..9b28ce5 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.nn.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.nn.pbtxt
@@ -101,6 +101,10 @@
argspec: "args=[\'labels\', \'inputs\', \'sequence_length\', \'preprocess_collapse_repeated\', \'ctc_merge_repeated\', \'ignore_longer_outputs_than_inputs\', \'time_major\'], varargs=None, keywords=None, defaults=[\'False\', \'True\', \'False\', \'True\'], "
}
member_method {
+ name: "depth_to_space"
+ argspec: "args=[\'input\', \'block_size\', \'name\', \'data_format\'], varargs=None, keywords=None, defaults=[\'None\', \'NHWC\'], "
+ }
+ member_method {
name: "depthwise_conv2d"
argspec: "args=[\'input\', \'filter\', \'strides\', \'padding\', \'rate\', \'name\', \'data_format\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
}
@@ -305,6 +309,14 @@
argspec: "args=[\'features\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "space_to_batch"
+ argspec: "args=[\'input\', \'paddings\', \'block_size\', \'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\'], "
+ }
+ member_method {
name: "sparse_softmax_cross_entropy_with_logits"
argspec: "args=[\'_sentinel\', \'labels\', \'logits\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.nn.rnn_cell.-basic-r-n-n-cell.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.nn.rnn_cell.-basic-r-n-n-cell.pbtxt
deleted file mode 100644
index a4483fe..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.nn.rnn_cell.-basic-r-n-n-cell.pbtxt
+++ /dev/null
@@ -1,202 +0,0 @@
-path: "tensorflow.nn.rnn_cell.BasicRNNCell"
-tf_class {
- is_instance: "<class \'tensorflow.python.ops.rnn_cell_impl.BasicRNNCell\'>"
- is_instance: "<class \'tensorflow.python.ops.rnn_cell_impl.LayerRNNCell\'>"
- is_instance: "<class \'tensorflow.python.ops.rnn_cell_impl.RNNCell\'>"
- is_instance: "<class \'tensorflow.python.layers.base.Layer\'>"
- 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: "graph"
- 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: "output_size"
- mtype: "<type \'property\'>"
- }
- member {
- name: "scope_name"
- mtype: "<type \'property\'>"
- }
- member {
- name: "state_size"
- 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\', \'num_units\', \'activation\', \'reuse\', \'name\', \'dtype\'], varargs=None, keywords=kwargs, defaults=[\'None\', \'None\', \'None\', \'None\'], "
- }
- member_method {
- name: "add_loss"
- argspec: "args=[\'self\', \'losses\', \'inputs\'], varargs=None, keywords=None, defaults=[\'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\', \'use_resource\', \'synchronization\', \'aggregation\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'VariableSynchronization.AUTO\', \'VariableAggregation.NONE\', \'None\'], "
- }
- member_method {
- name: "apply"
- argspec: "args=[\'self\', \'inputs\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "build"
- argspec: "args=[\'instance\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "call"
- argspec: "args=[\'self\', \'inputs\', \'state\'], varargs=None, keywords=None, 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_initial_state"
- argspec: "args=[\'self\', \'inputs\', \'batch_size\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'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: "set_weights"
- argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "zero_state"
- argspec: "args=[\'self\', \'batch_size\', \'dtype\'], varargs=None, keywords=None, defaults=None"
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.nn.rnn_cell.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.nn.rnn_cell.pbtxt
index 64697e8..24767e2 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.nn.rnn_cell.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.nn.rnn_cell.pbtxt
@@ -5,10 +5,6 @@
mtype: "<type \'type\'>"
}
member {
- name: "BasicRNNCell"
- mtype: "<type \'type\'>"
- }
- member {
name: "DeviceWrapper"
mtype: "<type \'type\'>"
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.pbtxt
index 96212f5..5b3ea75 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.pbtxt
@@ -1,10 +1,6 @@
path: "tensorflow"
tf_module {
member {
- name: "AUTO_REUSE"
- mtype: "<enum \'_ReuseMode\'>"
- }
- member {
name: "AggregationMethod"
mtype: "<type \'type\'>"
}
@@ -233,18 +229,10 @@
mtype: "<type \'type\'>"
}
member {
- name: "Variable"
- mtype: "<class \'tensorflow.python.ops.variables.VariableMetaclass\'>"
- }
- member {
name: "VariableAggregation"
mtype: "<class \'enum.EnumMeta\'>"
}
member {
- name: "VariableScope"
- mtype: "<type \'type\'>"
- }
- member {
name: "VariableSynchronization"
mtype: "<class \'enum.EnumMeta\'>"
}
@@ -469,6 +457,10 @@
mtype: "<class \'tensorflow.python.framework.dtypes.DType\'>"
}
member {
+ name: "random"
+ mtype: "<type \'module\'>"
+ }
+ member {
name: "random_normal_initializer"
mtype: "<type \'type\'>"
}
@@ -553,10 +545,6 @@
mtype: "<type \'module\'>"
}
member {
- name: "variable_scope"
- mtype: "<type \'type\'>"
- }
- member {
name: "variance_scaling_initializer"
mtype: "<type \'type\'>"
}
@@ -617,10 +605,6 @@
argspec: "args=[\'names\', \'value\'], varargs=None, keywords=None, defaults=None"
}
member_method {
- name: "all_variables"
- argspec: "args=[], varargs=None, keywords=None, defaults=None"
- }
- member_method {
name: "angle"
argspec: "args=[\'input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -733,10 +717,6 @@
argspec: "args=[\'tensor\', \'tf_type\', \'message\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
}
member_method {
- name: "assert_variables_initialized"
- argspec: "args=[\'var_list\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
name: "atan"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -1137,10 +1117,6 @@
argspec: "args=[], varargs=None, keywords=None, defaults=None"
}
member_method {
- name: "get_local_variable"
- argspec: "args=[\'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'collections\', \'caching_device\', \'partitioner\', \'validate_shape\', \'use_resource\', \'custom_getter\', \'constraint\', \'synchronization\', \'aggregation\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'False\', \'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'VariableSynchronization.AUTO\', \'VariableAggregation.NONE\'], "
- }
- member_method {
name: "get_seed"
argspec: "args=[\'op_seed\'], varargs=None, keywords=None, defaults=None"
}
@@ -1153,26 +1129,10 @@
argspec: "args=[\'handle\', \'dtype\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
- name: "get_variable"
- argspec: "args=[\'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'collections\', \'caching_device\', \'partitioner\', \'validate_shape\', \'use_resource\', \'custom_getter\', \'constraint\', \'synchronization\', \'aggregation\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'VariableSynchronization.AUTO\', \'VariableAggregation.NONE\'], "
- }
- member_method {
- name: "get_variable_scope"
- argspec: "args=[], varargs=None, keywords=None, defaults=None"
- }
- member_method {
name: "global_norm"
argspec: "args=[\'t_list\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
- name: "global_variables"
- argspec: "args=[\'scope\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "global_variables_initializer"
- argspec: "args=[], varargs=None, keywords=None, defaults=None"
- }
- member_method {
name: "gradients"
argspec: "args=[\'ys\', \'xs\', \'grad_ys\', \'name\', \'colocate_gradients_with_ops\', \'gate_gradients\', \'aggregation_method\', \'stop_gradients\'], varargs=None, keywords=None, defaults=[\'None\', \'gradients\', \'False\', \'False\', \'None\', \'None\'], "
}
@@ -1249,18 +1209,6 @@
argspec: "args=[\'name\'], varargs=None, keywords=None, defaults=[\'init_all_tables\'], "
}
member_method {
- name: "initialize_all_variables"
- argspec: "args=[], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "initialize_local_variables"
- argspec: "args=[], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "initialize_variables"
- argspec: "args=[\'var_list\', \'name\'], varargs=None, keywords=None, defaults=[\'init\'], "
- }
- member_method {
name: "invert_permutation"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -1289,10 +1237,6 @@
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
- name: "is_variable_initialized"
- argspec: "args=[\'variable\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
name: "lbeta"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -1329,14 +1273,6 @@
argspec: "args=[\'library_filename\'], varargs=None, keywords=None, defaults=None"
}
member_method {
- name: "local_variables"
- argspec: "args=[\'scope\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "local_variables_initializer"
- argspec: "args=[], varargs=None, keywords=None, defaults=None"
- }
- member_method {
name: "log"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -1449,14 +1385,6 @@
argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
- name: "model_variables"
- argspec: "args=[\'scope\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "moving_average_variables"
- argspec: "args=[\'scope\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
name: "multinomial"
argspec: "args=[\'logits\', \'num_samples\', \'seed\', \'name\', \'output_dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
}
@@ -1657,10 +1585,6 @@
argspec: "args=[\'base_type\', \'conversion_func\', \'priority\'], varargs=None, keywords=None, defaults=[\'100\'], "
}
member_method {
- name: "report_uninitialized_variables"
- argspec: "args=[\'var_list\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'report_uninitialized_variables\'], "
- }
- member_method {
name: "required_space_to_batch_paddings"
argspec: "args=[\'input_shape\', \'block_shape\', \'base_paddings\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
}
@@ -1689,6 +1613,10 @@
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "roll"
+ argspec: "args=[\'input\', \'shift\', \'axis\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
name: "round"
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -2069,10 +1997,6 @@
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
- name: "trainable_variables"
- argspec: "args=[\'scope\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
name: "transpose"
argspec: "args=[\'a\', \'perm\', \'name\', \'conjugate\'], varargs=None, keywords=None, defaults=[\'None\', \'transpose\', \'False\'], "
}
@@ -2141,14 +2065,6 @@
argspec: "args=[\'max_shard_bytes\', \'axis\', \'bytes_per_string_element\', \'max_shards\'], varargs=None, keywords=None, defaults=[\'0\', \'16\', \'None\'], "
}
member_method {
- name: "variable_op_scope"
- argspec: "args=[\'values\', \'name_or_scope\', \'default_name\', \'initializer\', \'regularizer\', \'caching_device\', \'partitioner\', \'custom_getter\', \'reuse\', \'dtype\', \'use_resource\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\'], "
- }
- member_method {
- name: "variables_initializer"
- argspec: "args=[\'var_list\', \'name\'], varargs=None, keywords=None, defaults=[\'init\'], "
- }
- member_method {
name: "verify_tensor_all_finite"
argspec: "args=[\'t\', \'msg\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.quantization.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.quantization.pbtxt
index 6d865ef..77c92ae 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.quantization.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.quantization.pbtxt
@@ -29,6 +29,10 @@
argspec: "args=[\'gradients\', \'inputs\', \'min\', \'max\', \'num_bits\', \'narrow_range\', \'name\'], varargs=None, keywords=None, defaults=[\'8\', \'False\', \'None\'], "
}
member_method {
+ name: "quantize"
+ argspec: "args=[\'input\', \'min_range\', \'max_range\', \'T\', \'mode\', \'round_mode\', \'name\'], varargs=None, keywords=None, defaults=[\'MIN_COMBINED\', \'HALF_AWAY_FROM_ZERO\', \'None\'], "
+ }
+ member_method {
name: "quantized_concat"
argspec: "args=[\'concat_dim\', \'values\', \'input_mins\', \'input_maxes\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.random.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.random.pbtxt
new file mode 100644
index 0000000..a568dd4
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.random.pbtxt
@@ -0,0 +1,47 @@
+path: "tensorflow.random"
+tf_module {
+ member_method {
+ name: "gamma"
+ argspec: "args=[\'shape\', \'alpha\', \'beta\', \'dtype\', \'seed\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \"<dtype: \'float32\'>\", \'None\', \'None\'], "
+ }
+ member_method {
+ name: "get_seed"
+ argspec: "args=[\'op_seed\'], varargs=None, keywords=None, defaults=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\'], "
+ }
+ member_method {
+ name: "multinomial"
+ argspec: "args=[\'logits\', \'num_samples\', \'seed\', \'name\', \'output_dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "normal"
+ argspec: "args=[\'shape\', \'mean\', \'stddev\', \'dtype\', \'seed\', \'name\'], varargs=None, keywords=None, defaults=[\'0.0\', \'1.0\', \"<dtype: \'float32\'>\", \'None\', \'None\'], "
+ }
+ member_method {
+ name: "poisson"
+ argspec: "args=[\'lam\', \'shape\', \'dtype\', \'seed\', \'name\'], varargs=None, keywords=None, defaults=[\"<dtype: \'float32\'>\", \'None\', \'None\'], "
+ }
+ member_method {
+ name: "set_random_seed"
+ argspec: "args=[\'seed\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "shuffle"
+ argspec: "args=[\'value\', \'seed\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "truncated_normal"
+ argspec: "args=[\'shape\', \'mean\', \'stddev\', \'dtype\', \'seed\', \'name\'], varargs=None, keywords=None, defaults=[\'0.0\', \'1.0\', \"<dtype: \'float32\'>\", \'None\', \'None\'], "
+ }
+ member_method {
+ name: "uniform"
+ argspec: "args=[\'shape\', \'minval\', \'maxval\', \'dtype\', \'seed\', \'name\'], varargs=None, keywords=None, defaults=[\'0\', \'None\', \"<dtype: \'float32\'>\", \'None\', \'None\'], "
+ }
+ member_method {
+ name: "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.saved_model.-builder.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.saved_model.-builder.pbtxt
new file mode 100644
index 0000000..67457de
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.saved_model.-builder.pbtxt
@@ -0,0 +1,21 @@
+path: "tensorflow.saved_model.Builder"
+tf_class {
+ is_instance: "<class \'tensorflow.python.saved_model.builder_impl.SavedModelBuilder\'>"
+ is_instance: "<type \'object\'>"
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'export_dir\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "add_meta_graph"
+ argspec: "args=[\'self\', \'tags\', \'signature_def_map\', \'assets_collection\', \'legacy_init_op\', \'clear_devices\', \'main_op\', \'strip_default_attrs\', \'saver\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'False\', \'None\', \'False\', \'None\'], "
+ }
+ member_method {
+ name: "add_meta_graph_and_variables"
+ argspec: "args=[\'self\', \'sess\', \'tags\', \'signature_def_map\', \'assets_collection\', \'legacy_init_op\', \'clear_devices\', \'main_op\', \'strip_default_attrs\', \'saver\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'False\', \'None\', \'False\', \'None\'], "
+ }
+ member_method {
+ name: "save"
+ argspec: "args=[\'self\', \'as_text\'], varargs=None, keywords=None, defaults=[\'False\'], "
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.saved_model.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.saved_model.pbtxt
index e1a0385..3f4965f 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.saved_model.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.saved_model.pbtxt
@@ -1,6 +1,10 @@
path: "tensorflow.saved_model"
tf_module {
member {
+ name: "Builder"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "builder"
mtype: "<type \'module\'>"
}
@@ -33,6 +37,46 @@
mtype: "<type \'module\'>"
}
member_method {
+ name: "build_signature_def"
+ argspec: "args=[\'inputs\', \'outputs\', \'method_name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "build_tensor_info"
+ argspec: "args=[\'tensor\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "classification_signature_def"
+ argspec: "args=[\'examples\', \'classes\', \'scores\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "get_tensor_from_tensor_info"
+ argspec: "args=[\'tensor_info\', \'graph\', \'import_scope\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "is_valid_signature"
+ argspec: "args=[\'signature_def\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "load"
+ argspec: "args=[\'sess\', \'tags\', \'export_dir\', \'import_scope\'], varargs=None, keywords=saver_kwargs, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "main_op_with_restore"
+ argspec: "args=[\'restore_op_name\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "maybe_saved_model_directory"
+ argspec: "args=[\'export_dir\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "predict_signature_def"
+ argspec: "args=[\'inputs\', \'outputs\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "regression_signature_def"
+ argspec: "args=[\'examples\', \'predictions\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
name: "simple_save"
argspec: "args=[\'session\', \'export_dir\', \'inputs\', \'outputs\', \'legacy_init_op\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.sparse.-sparse-conditional-accumulator.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.sparse.-sparse-conditional-accumulator.pbtxt
new file mode 100644
index 0000000..cd97716
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.sparse.-sparse-conditional-accumulator.pbtxt
@@ -0,0 +1,46 @@
+path: "tensorflow.sparse.SparseConditionalAccumulator"
+tf_class {
+ is_instance: "<class \'tensorflow.python.ops.data_flow_ops.SparseConditionalAccumulator\'>"
+ is_instance: "<class \'tensorflow.python.ops.data_flow_ops.ConditionalAccumulatorBase\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "accumulator_ref"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "dtype"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "name"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'dtype\', \'shape\', \'shared_name\', \'name\', \'reduction_type\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'sparse_conditional_accumulator\', \'MEAN\'], "
+ }
+ member_method {
+ name: "apply_grad"
+ argspec: "args=[\'self\', \'grad_indices\', \'grad_values\', \'grad_shape\', \'local_step\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'0\', \'None\'], "
+ }
+ member_method {
+ name: "apply_indexed_slices_grad"
+ argspec: "args=[\'self\', \'grad\', \'local_step\', \'name\'], varargs=None, keywords=None, defaults=[\'0\', \'None\'], "
+ }
+ member_method {
+ name: "num_accumulated"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "set_global_step"
+ argspec: "args=[\'self\', \'new_global_step\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "take_grad"
+ argspec: "args=[\'self\', \'num_required\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "take_indexed_slices_grad"
+ argspec: "args=[\'self\', \'num_required\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.sparse.-sparse-tensor.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.sparse.-sparse-tensor.pbtxt
new file mode 100644
index 0000000..02e59a6
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.sparse.-sparse-tensor.pbtxt
@@ -0,0 +1,54 @@
+path: "tensorflow.sparse.SparseTensor"
+tf_class {
+ is_instance: "<class \'tensorflow.python.framework.sparse_tensor.SparseTensor\'>"
+ is_instance: "<class \'tensorflow.python.framework.ops._TensorLike\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "dense_shape"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "dtype"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "graph"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "indices"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "op"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "shape"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "values"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'indices\', \'values\', \'dense_shape\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "consumers"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "eval"
+ argspec: "args=[\'self\', \'feed_dict\', \'session\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "from_value"
+ argspec: "args=[\'cls\', \'sparse_tensor_value\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "get_shape"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.sparse.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.sparse.pbtxt
index ba9e651..32bd8d5f 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.sparse.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.sparse.pbtxt
@@ -1,5 +1,21 @@
path: "tensorflow.sparse"
tf_module {
+ member {
+ name: "SparseConditionalAccumulator"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "SparseTensor"
+ mtype: "<type \'type\'>"
+ }
+ member_method {
+ name: "add"
+ argspec: "args=[\'a\', \'b\', \'thresh\'], varargs=None, keywords=None, defaults=[\'0\'], "
+ }
+ member_method {
+ name: "concat"
+ argspec: "args=[\'axis\', \'sp_inputs\', \'name\', \'expand_nonconcat_dim\', \'concat_dim\'], varargs=None, keywords=None, defaults=[\'None\', \'False\', \'None\'], "
+ }
member_method {
name: "cross"
argspec: "args=[\'inputs\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
@@ -16,4 +32,100 @@
name: "eye"
argspec: "args=[\'num_rows\', \'num_columns\', \'dtype\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \"<dtype: \'float32\'>\", \'None\'], "
}
+ member_method {
+ name: "fill_empty_rows"
+ argspec: "args=[\'sp_input\', \'default_value\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "mask"
+ argspec: "args=[\'a\', \'mask_indices\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "matmul"
+ argspec: "args=[\'sp_a\', \'b\', \'adjoint_a\', \'adjoint_b\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'False\', \'None\'], "
+ }
+ member_method {
+ name: "maximum"
+ argspec: "args=[\'sp_a\', \'sp_b\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "merge"
+ argspec: "args=[\'sp_ids\', \'sp_values\', \'vocab_size\', \'name\', \'already_sorted\'], varargs=None, keywords=None, defaults=[\'None\', \'False\'], "
+ }
+ member_method {
+ name: "minimum"
+ argspec: "args=[\'sp_a\', \'sp_b\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "placeholder"
+ argspec: "args=[\'dtype\', \'shape\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "reduce_max"
+ argspec: "args=[\'sp_input\', \'axis\', \'keepdims\', \'reduction_axes\', \'keep_dims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "reduce_max_sparse"
+ argspec: "args=[\'sp_input\', \'axis\', \'keepdims\', \'reduction_axes\', \'keep_dims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "reduce_sum"
+ argspec: "args=[\'sp_input\', \'axis\', \'keepdims\', \'reduction_axes\', \'keep_dims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "reduce_sum_sparse"
+ argspec: "args=[\'sp_input\', \'axis\', \'keepdims\', \'reduction_axes\', \'keep_dims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "reorder"
+ argspec: "args=[\'sp_input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "reset_shape"
+ argspec: "args=[\'sp_input\', \'new_shape\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "reshape"
+ argspec: "args=[\'sp_input\', \'shape\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "retain"
+ argspec: "args=[\'sp_input\', \'to_retain\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "segment_mean"
+ argspec: "args=[\'data\', \'indices\', \'segment_ids\', \'name\', \'num_segments\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "segment_sqrt_n"
+ argspec: "args=[\'data\', \'indices\', \'segment_ids\', \'name\', \'num_segments\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "segment_sum"
+ argspec: "args=[\'data\', \'indices\', \'segment_ids\', \'name\', \'num_segments\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "slice"
+ argspec: "args=[\'sp_input\', \'start\', \'size\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "softmax"
+ argspec: "args=[\'sp_input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "split"
+ argspec: "args=[\'keyword_required\', \'sp_input\', \'num_split\', \'axis\', \'name\', \'split_dim\'], varargs=None, keywords=None, defaults=[\'KeywordRequired()\', \'None\', \'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "to_dense"
+ argspec: "args=[\'sp_input\', \'default_value\', \'validate_indices\', \'name\'], varargs=None, keywords=None, defaults=[\'0\', \'True\', \'None\'], "
+ }
+ member_method {
+ name: "to_indicator"
+ argspec: "args=[\'sp_input\', \'vocab_size\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "transpose"
+ argspec: "args=[\'sp_input\', \'perm\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.strings.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.strings.pbtxt
index c52581d..ebdaf57 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.strings.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.strings.pbtxt
@@ -13,6 +13,10 @@
argspec: "args=[\'input\', \'name\', \'unit\'], varargs=None, keywords=None, defaults=[\'None\', \'BYTE\'], "
}
member_method {
+ name: "reduce_join"
+ argspec: "args=[\'inputs\', \'axis\', \'keep_dims\', \'separator\', \'name\', \'reduction_indices\'], varargs=None, keywords=None, defaults=[\'None\', \'False\', \'\', \'None\', \'None\'], "
+ }
+ member_method {
name: "regex_full_match"
argspec: "args=[\'input\', \'pattern\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -48,4 +52,8 @@
name: "to_number"
argspec: "args=[\'string_tensor\', \'out_type\', \'name\'], varargs=None, keywords=None, defaults=[\"<dtype: \'float32\'>\", \'None\'], "
}
+ member_method {
+ name: "unicode_script"
+ argspec: "args=[\'input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.train.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.train.pbtxt
index cb6da50..7e980fe 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.train.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.train.pbtxt
@@ -253,6 +253,10 @@
argspec: "args=[\'checkpoint_prefix\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "confusion_matrix"
+ argspec: "args=[\'labels\', \'predictions\', \'num_classes\', \'dtype\', \'name\', \'weights\'], varargs=None, keywords=None, defaults=[\'None\', \"<dtype: \'int32\'>\", \'None\', \'None\'], "
+ }
+ 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/golden/v2/tensorflow.variable_scope.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.variable_scope.pbtxt
deleted file mode 100644
index e62dec9..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.variable_scope.pbtxt
+++ /dev/null
@@ -1,9 +0,0 @@
-path: "tensorflow.variable_scope"
-tf_class {
- is_instance: "<class \'tensorflow.python.ops.variable_scope.variable_scope\'>"
- is_instance: "<type \'object\'>"
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'name_or_scope\', \'default_name\', \'values\', \'initializer\', \'regularizer\', \'caching_device\', \'partitioner\', \'custom_getter\', \'reuse\', \'dtype\', \'use_resource\', \'constraint\', \'auxiliary_name_scope\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'True\'], "
- }
-}
diff --git a/tensorflow/tools/api/tests/BUILD b/tensorflow/tools/api/tests/BUILD
index 4efa4a9..3cbea41 100644
--- a/tensorflow/tools/api/tests/BUILD
+++ b/tensorflow/tools/api/tests/BUILD
@@ -19,6 +19,7 @@
"api_compatibility_test.py",
"//tensorflow:tf_python_api_gen_v2",
],
+ args = ["--only_test_core_api=true"],
data = [
"//tensorflow/tools/api/golden:api_golden_v1",
"//tensorflow/tools/api/golden:api_golden_v2",
diff --git a/tensorflow/tools/api/tests/api_compatibility_test.py b/tensorflow/tools/api/tests/api_compatibility_test.py
index d06c7f2..6487a62 100644
--- a/tensorflow/tools/api/tests/api_compatibility_test.py
+++ b/tensorflow/tools/api/tests/api_compatibility_test.py
@@ -56,6 +56,14 @@
have to be authorized by TensorFlow leads.
"""
+# DEFINE_boolean, only_test_core_api, default False:
+_ONLY_TEST_CORE_API_HELP = """
+ Some TF APIs are being moved outside of the tensorflow/ directory. There is
+ no garuntee which versions of these APIs will be present when running this
+ test. Therefore, do not error out on API changes in non-core TF code
+ if this flag is set.
+"""
+
# DEFINE_boolean, verbose_diffs, default True:
_VERBOSE_DIFFS_HELP = """
If set to true, print line by line diffs on all libraries. If set to
@@ -67,6 +75,8 @@
_TEST_README_FILE = 'tensorflow/tools/api/tests/README.txt'
_UPDATE_WARNING_FILE = 'tensorflow/tools/api/tests/API_UPDATE_WARNING.txt'
+_NON_CORE_PACKAGES = ['estimator']
+
def _KeyToFilePath(key, api_version):
"""From a given key, construct a filepath.
@@ -111,6 +121,19 @@
'They are not yet supported by the API tools.' % path)
+def _FilterNonCoreGoldenFiles(golden_file_list):
+ """Filter out non-core API pbtxt files."""
+ filtered_file_list = []
+ filtered_package_prefixes = [
+ 'tensorflow.%s.' % p for p in _NON_CORE_PACKAGES]
+ for f in golden_file_list:
+ if any([f.rsplit('/')[-1].startswith(pre)
+ for pre in filtered_package_prefixes]):
+ continue
+ filtered_file_list.append(f)
+ return filtered_file_list
+
+
class ApiCompatibilityTest(test.TestCase):
def __init__(self, *args, **kwargs):
@@ -233,6 +256,9 @@
return
visitor = public_api.PublicAPIVisitor(_VerifyNoSubclassOfMessageVisitor)
visitor.do_not_descend_map['tf'].append('contrib')
+ if FLAGS.only_test_core_api:
+ visitor.do_not_descend_map['tf'].extend(
+ _NON_CORE_PACKAGES)
traverse.traverse(tf_v2.compat.v1, visitor)
def testNoSubclassOfMessageV2(self):
@@ -240,6 +266,9 @@
return
visitor = public_api.PublicAPIVisitor(_VerifyNoSubclassOfMessageVisitor)
visitor.do_not_descend_map['tf'].append('contrib')
+ if FLAGS.only_test_core_api:
+ visitor.do_not_descend_map['tf'].extend(
+ _NON_CORE_PACKAGES)
traverse.traverse(tf_v2, visitor)
def _checkBackwardsCompatibility(
@@ -252,6 +281,9 @@
public_api_visitor.do_not_descend_map['tf'].append('contrib')
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)
if additional_private_map:
public_api_visitor.private_map.update(additional_private_map)
@@ -260,6 +292,8 @@
# Read all golden files.
golden_file_list = file_io.get_matching_files(golden_file_pattern)
+ if FLAGS.only_test_core_api:
+ golden_file_list = _FilterNonCoreGoldenFiles(golden_file_list)
def _ReadFileToProto(filename):
"""Read a filename, create a protobuf from its contents."""
@@ -325,6 +359,11 @@
parser = argparse.ArgumentParser()
parser.add_argument(
'--update_goldens', type=bool, default=False, help=_UPDATE_GOLDENS_HELP)
+ # TODO(mikecase): Create Estimator's own API compatibility test or
+ # a more general API compatibility test for use for TF components.
+ parser.add_argument(
+ '--only_test_core_api', type=bool, default=False,
+ help=_ONLY_TEST_CORE_API_HELP)
parser.add_argument(
'--verbose_diffs', type=bool, default=True, help=_VERBOSE_DIFFS_HELP)
FLAGS, unparsed = parser.parse_known_args()
diff --git a/tensorflow/tools/benchmark/README.md b/tensorflow/tools/benchmark/README.md
index e64af2b..dee1a20 100644
--- a/tensorflow/tools/benchmark/README.md
+++ b/tensorflow/tools/benchmark/README.md
@@ -32,7 +32,7 @@
(4) Run the benchmark. For example:
```
-adb shell "/data/local/tmp/benchmark_model \
+adb shell /data/local/tmp/benchmark_model \
--graph=/data/local/tmp/tensorflow_inception_graph.pb \
--input_layer="input:0" \
--input_layer_shape="1,224,224,3" \
diff --git a/tensorflow/tools/ci_build/Dockerfile.cmake b/tensorflow/tools/ci_build/Dockerfile.cmake
index b7450c8..ef0024f 100644
--- a/tensorflow/tools/ci_build/Dockerfile.cmake
+++ b/tensorflow/tools/ci_build/Dockerfile.cmake
@@ -28,8 +28,8 @@
RUN pip install --upgrade gast
RUN pip install --upgrade numpy
RUN pip install --upgrade termcolor
-RUN pip install keras_applications==1.0.5
-RUN pip install keras_preprocessing==1.0.3
+RUN pip install --upgrade keras_applications
+RUN pip install --upgrade keras_preprocessing
# Install golang
RUN apt-get install -t xenial-backports -y golang-1.9
diff --git a/tensorflow/tools/ci_build/Dockerfile.rbe.cuda9.0-cudnn7-ubuntu14.04 b/tensorflow/tools/ci_build/Dockerfile.rbe.cuda9.0-cudnn7-ubuntu14.04
index a30858d..dd8d705 100644
--- a/tensorflow/tools/ci_build/Dockerfile.rbe.cuda9.0-cudnn7-ubuntu14.04
+++ b/tensorflow/tools/ci_build/Dockerfile.rbe.cuda9.0-cudnn7-ubuntu14.04
@@ -26,7 +26,7 @@
ENV NVIDIA_DRIVER_CAPABILITIES compute,utility
ENV NVIDIA_REQUIRE_CUDA "cuda>=9.0"
ENV NCCL_VERSION 2.2.13
-ENV CUDNN_VERSION 7.2.1.38
+ENV CUDNN_VERSION 7.1.4.18
# TODO(b/110903506): /usr/loca/cuda/lib64/stubs should not be needed in
# LD_LIBRARY_PATH. The stubs/libcuda.so is not meant to used at runtime. The
diff --git a/tensorflow/tools/ci_build/Dockerfile.rocm b/tensorflow/tools/ci_build/Dockerfile.rocm
new file mode 100644
index 0000000..aadaa8b
--- /dev/null
+++ b/tensorflow/tools/ci_build/Dockerfile.rocm
@@ -0,0 +1,97 @@
+# This Dockerfile provides a starting point for a ROCm installation of
+# MIOpen and tensorflow.
+FROM ubuntu:xenial
+MAINTAINER Jeff Poznanovic <jeffrey.poznanovic@amd.com>
+
+ARG DEB_ROCM_REPO=http://repo.radeon.com/rocm/apt/debian/
+ARG ROCM_PATH=/opt/rocm
+
+ENV DEBIAN_FRONTEND noninteractive
+ENV TF_NEED_ROCM 1
+ENV HOME /root/
+RUN apt update && apt install -y wget software-properties-common
+
+# Add rocm repository
+RUN apt-get clean all
+RUN wget -qO - $DEB_ROCM_REPO/rocm.gpg.key | apt-key add -
+RUN sh -c "echo deb [arch=amd64] $DEB_ROCM_REPO xenial main > /etc/apt/sources.list.d/rocm.list"
+
+# Install misc pkgs
+RUN apt-get update --allow-insecure-repositories && DEBIAN_FRONTEND=noninteractive apt-get install -y \
+ build-essential \
+ clang-3.8 \
+ clang-format-3.8 \
+ clang-tidy-3.8 \
+ cmake \
+ cmake-qt-gui \
+ ssh \
+ curl \
+ apt-utils \
+ pkg-config \
+ g++-multilib \
+ git \
+ libunwind-dev \
+ libfftw3-dev \
+ libelf-dev \
+ libncurses5-dev \
+ libpthread-stubs0-dev \
+ vim \
+ gfortran \
+ libboost-program-options-dev \
+ libssl-dev \
+ libboost-dev \
+ libboost-system-dev \
+ libboost-filesystem-dev \
+ rpm \
+ libnuma-dev \
+ virtualenv \
+ python-pip \
+ python3-pip \
+ wget && \
+ apt-get clean && \
+ rm -rf /var/lib/apt/lists/*
+
+# Install rocm pkgs
+RUN apt-get update --allow-insecure-repositories && \
+ DEBIAN_FRONTEND=noninteractive apt-get install -y --allow-unauthenticated \
+ rocm-dev rocm-libs rocm-utils \
+ rocfft miopen-hip miopengemm rocblas hipblas rocrand \
+ rocm-profiler cxlactivitylogger && \
+ apt-get clean && \
+ rm -rf /var/lib/apt/lists/*
+
+RUN cd ~ && git clone https://github.com/GPUOpen-ProfessionalCompute-Tools/HIP.git
+RUN cd ~/HIP && mkdir -p build && cd build && cmake .. && make package -j && dpkg -i *.deb
+
+ENV HCC_HOME=$ROCM_PATH/hcc
+ENV HIP_PATH=$ROCM_PATH/hip
+ENV OPENCL_ROOT=$ROCM_PATH/opencl
+ENV PATH="$HCC_HOME/bin:$HIP_PATH/bin:${PATH}"
+ENV PATH="$ROCM_PATH/bin:${PATH}"
+ENV PATH="$OPENCL_ROOT/bin:${PATH}"
+
+# Add target file to help determine which device(s) to build for
+RUN echo -e "gfx803\ngfx900" >> /opt/rocm/bin/target.lst
+
+# Setup environment variables, and add those environment variables at the end of ~/.bashrc
+ARG HCC_HOME=/opt/rocm/hcc
+ARG HIP_PATH=/opt/rocm/hip
+ARG PATH=$HCC_HOME/bin:$HIP_PATH/bin:$PATH
+
+# Copy and run the install scripts.
+COPY install/*.sh /install/
+ARG DEBIAN_FRONTEND=noninteractive
+RUN /install/install_bootstrap_deb_packages.sh
+RUN add-apt-repository -y ppa:openjdk-r/ppa && \
+ add-apt-repository -y ppa:george-edison55/cmake-3.x
+RUN /install/install_deb_packages.sh
+RUN /install/install_pip_packages.sh
+RUN /install/install_bazel.sh
+RUN /install/install_golang.sh
+
+# Set up the master bazelrc configuration file.
+COPY install/.bazelrc /etc/bazel.bazelrc
+
+# Configure the build for our CUDA configuration.
+ENV TF_NEED_ROCM 1
+
diff --git a/tensorflow/tools/ci_build/builds/docker_test.sh b/tensorflow/tools/ci_build/builds/docker_test.sh
index e337ea4..38891b6 100755
--- a/tensorflow/tools/ci_build/builds/docker_test.sh
+++ b/tensorflow/tools/ci_build/builds/docker_test.sh
@@ -19,7 +19,7 @@
#
# Usage: docker_test.sh <IMAGE_TYPE> <TAG> <WHL_PATH>
# Arguments:
-# IMAGE_TYPE : Type of the image: (CPU|GPU)
+# IMAGE_TYPE : Type of the image: (CPU|GPU|ROCM)
# TAG : Docker image tag
# WHL_PATH : Path to the whl file to be installed inside the docker image
#
@@ -60,6 +60,8 @@
DOCKERFILE="tensorflow/tools/docker/Dockerfile"
elif [[ "${IMAGE_TYPE}" == "gpu" ]]; then
DOCKERFILE="tensorflow/tools/docker/Dockerfile.gpu"
+elif [[ "${IMAGE_TYPE}" == "rocm" ]]; then
+ DOCKERFILE="tensorflow/tools/docker/Dockerfile.rocm"
else
die "Unrecognized image type: $1"
fi
@@ -106,13 +108,16 @@
devices=$(\ls /dev/nvidia* | xargs -I{} echo '--device {}:{}')
libs=$(\ls /usr/lib/x86_64-linux-gnu/libcuda.* | xargs -I{} echo '-v {}:{}')
GPU_EXTRA_PARAMS="${devices} ${libs}"
+elif [ "${IMAGE_TYPE}" == "rocm" ]; then
+ ROCM_EXTRA_PARAMS="--device=/dev/kfd --device=/dev/dri --group-add video"
else
GPU_EXTRA_PARAMS=""
+ ROCM_EXTRA_PARAMS=""
fi
# Run docker image with source directory mapped
docker run -v ${BASE_DIR}:/tensorflow-src -w /tensorflow-src \
-${GPU_EXTRA_PARAMS} \
+${GPU_EXTRA_PARAMS} ${ROCM_EXTRA_PARAMS} \
"${DOCKER_IMG_TAG}" \
/bin/bash -c "tensorflow/tools/ci_build/builds/run_pip_tests.sh && "\
"tensorflow/tools/ci_build/builds/test_tutorials.sh && "\
diff --git a/tensorflow/tools/ci_build/builds/pip.sh b/tensorflow/tools/ci_build/builds/pip.sh
index fef121a..6543779 100755
--- a/tensorflow/tools/ci_build/builds/pip.sh
+++ b/tensorflow/tools/ci_build/builds/pip.sh
@@ -132,6 +132,7 @@
PIP_BUILD_TARGET="//tensorflow/tools/pip_package:build_pip_package"
GPU_FLAG=""
if [[ ${CONTAINER_TYPE} == "cpu" ]] || \
+ [[ ${CONTAINER_TYPE} == "rocm" ]] || \
[[ ${CONTAINER_TYPE} == "debian.jessie.cpu" ]]; then
bazel build ${BAZEL_FLAGS} ${PIP_BUILD_TARGET} || \
die "Build failed."
@@ -255,7 +256,8 @@
die "ERROR: Cannot find repaired wheel."
fi
# Copy and rename for gpu manylinux as we do not want auditwheel to package in libcudart.so
- elif [[ ${CONTAINER_TYPE} == "gpu" ]]; then
+ elif [[ ${CONTAINER_TYPE} == "gpu" ]] || \
+ [[ ${CONTAINER_TYPE} == "rocm" ]]; then
WHL_PATH=${AUDITED_WHL_NAME}
cp ${WHL_DIR}/${WHL_BASE_NAME} ${WHL_PATH}
echo "Copied manylinx1 wheel file at ${WHL_PATH}"
diff --git a/tensorflow/tools/ci_build/builds/run_pip_tests.sh b/tensorflow/tools/ci_build/builds/run_pip_tests.sh
index 17198a6..7d5cf3f 100755
--- a/tensorflow/tools/ci_build/builds/run_pip_tests.sh
+++ b/tensorflow/tools/ci_build/builds/run_pip_tests.sh
@@ -111,7 +111,6 @@
# virtualenv.
export TF_NEED_GCP=0
export TF_NEED_HDFS=0
-export TF_ENABLE_XLA=0
# Obtain the path to Python binary
if [[ ${IS_VIRTUALENV} == "1" ]]; then
diff --git a/tensorflow/tools/ci_build/builds/with_the_same_user b/tensorflow/tools/ci_build/builds/with_the_same_user
index b216e35..1cc5aed 100755
--- a/tensorflow/tools/ci_build/builds/with_the_same_user
+++ b/tensorflow/tools/ci_build/builds/with_the_same_user
@@ -48,6 +48,12 @@
usermod -a -G sudo "${CI_BUILD_USER}"
echo "${CI_BUILD_USER} ALL=(ALL) NOPASSWD:ALL" > /etc/sudoers.d/90-nopasswd-sudo
+if [[ "${TF_NEED_ROCM}" -eq 1 ]]; then
+ # ROCm requires the video group in order to use the GPU for compute. If it
+ # exists on the host, add it to the container.
+ getent group video || addgroup video && adduser "${CI_BUILD_USER}" video
+fi
+
if [ -e /root/.bazelrc ]; then
cp /root/.bazelrc "${CI_BUILD_HOME}/.bazelrc"
chown "${CI_BUILD_UID}:${CI_BUILD_GID}" "${CI_BUILD_HOME}/.bazelrc"
diff --git a/tensorflow/tools/ci_build/ci_build.sh b/tensorflow/tools/ci_build/ci_build.sh
index 77265e0..eab0616 100755
--- a/tensorflow/tools/ci_build/ci_build.sh
+++ b/tensorflow/tools/ci_build/ci_build.sh
@@ -18,7 +18,7 @@
# <COMMAND>
#
# CONTAINER_TYPE: Type of the docker container used the run the build:
-# e.g., (cpu | gpu | android | tensorboard)
+# e.g., (cpu | gpu | rocm | android | tensorboard)
#
# DOCKERFILE_PATH: (Optional) Path to the Dockerfile used for docker build.
# If this optional value is not supplied (via the
@@ -103,6 +103,14 @@
GPU_EXTRA_PARAMS=""
fi
+# Add extra params for rocm devices and libraries for ROCm container.
+if [[ "${CONTAINER_TYPE}" == "rocm" ]]; then
+ ROCM_EXTRA_PARAMS="--device=/dev/kfd --device=/dev/dri --group-add video"
+else
+ ROCM_EXTRA_PARAMS=""
+fi
+
+
# Determine the docker image name
DOCKER_IMG_NAME="${BUILD_TAG}.${CONTAINER_TYPE}"
@@ -159,6 +167,7 @@
-v ${WORKSPACE}:/workspace \
-w /workspace \
${GPU_EXTRA_PARAMS} \
+ ${ROCM_EXTRA_PARAMS} \
${CI_DOCKER_EXTRA_PARAMS[@]} \
"${DOCKER_IMG_NAME}" \
${CI_COMMAND_PREFIX[@]} \
diff --git a/tensorflow/tools/ci_build/install/install_pip_packages.sh b/tensorflow/tools/ci_build/install/install_pip_packages.sh
index 4ced96f..b90f3f3 100755
--- a/tensorflow/tools/ci_build/install/install_pip_packages.sh
+++ b/tensorflow/tools/ci_build/install/install_pip_packages.sh
@@ -115,10 +115,10 @@
pip3 install --upgrade setuptools==39.1.0
# Keras
-pip2 install keras_applications==1.0.5 --no-deps
-pip3 install keras_applications==1.0.5 --no-deps
-pip2 install keras_preprocessing==1.0.3 --no-deps
-pip3 install keras_preprocessing==1.0.3 --no-deps
+pip2 install keras_applications==1.0.6 --no-deps
+pip3 install keras_applications==1.0.6 --no-deps
+pip2 install keras_preprocessing==1.0.5 --no-deps
+pip3 install keras_preprocessing==1.0.5 --no-deps
pip2 install --upgrade h5py==2.8.0
pip3 install --upgrade h5py==2.8.0
diff --git a/tensorflow/tools/ci_build/install/install_python3.5_pip_packages.sh b/tensorflow/tools/ci_build/install/install_python3.5_pip_packages.sh
index 37e6b51..61d4fe3 100755
--- a/tensorflow/tools/ci_build/install/install_python3.5_pip_packages.sh
+++ b/tensorflow/tools/ci_build/install/install_python3.5_pip_packages.sh
@@ -85,8 +85,8 @@
pip3.5 install --upgrade setuptools==39.1.0
# Keras
-pip3.5 install keras_applications==1.0.5
-pip3.5 install keras_preprocessing==1.0.3
+pip3.5 install keras_applications==1.0.6
+pip3.5 install keras_preprocessing==1.0.5
pip3.5 install --upgrade h5py==2.8.0
# Install last working version of setuptools.
diff --git a/tensorflow/tools/ci_build/install/install_python3.6_pip_packages.sh b/tensorflow/tools/ci_build/install/install_python3.6_pip_packages.sh
index 7520ff7..8949af8 100755
--- a/tensorflow/tools/ci_build/install/install_python3.6_pip_packages.sh
+++ b/tensorflow/tools/ci_build/install/install_python3.6_pip_packages.sh
@@ -102,7 +102,7 @@
pip3 install --upgrade h5py==2.8.0
# Keras
-pip3 install keras_applications==1.0.5
-pip3 install keras_preprocessing==1.0.3
+pip3 install keras_applications==1.0.6
+pip3 install keras_preprocessing==1.0.5
# LINT.ThenChange(//tensorflow/tools/ci_build/install/install_python3.5_pip_packages.sh)
diff --git a/tensorflow/tools/ci_build/linux/cpu/run_cc_core.sh b/tensorflow/tools/ci_build/linux/cpu/run_cc_core.sh
index 8eeddcd..3b5c92d 100755
--- a/tensorflow/tools/ci_build/linux/cpu/run_cc_core.sh
+++ b/tensorflow/tools/ci_build/linux/cpu/run_cc_core.sh
@@ -26,6 +26,7 @@
# Run configure.
export TF_NEED_CUDA=0
+export TF_NEED_ROCM=0
export CC_OPT_FLAGS='-mavx'
# Only running cc tests, python version does not matter.
export PYTHON_BIN_PATH=`which python`
diff --git a/tensorflow/tools/ci_build/linux/cpu/run_py2_core.sh b/tensorflow/tools/ci_build/linux/cpu/run_py2_core.sh
index 8eca198..52eff63 100755
--- a/tensorflow/tools/ci_build/linux/cpu/run_py2_core.sh
+++ b/tensorflow/tools/ci_build/linux/cpu/run_py2_core.sh
@@ -26,6 +26,7 @@
# Run configure.
export TF_NEED_CUDA=0
+export TF_NEED_ROCM=0
export CC_OPT_FLAGS='-mavx'
export PYTHON_BIN_PATH=`which python2`
yes "" | $PYTHON_BIN_PATH configure.py
diff --git a/tensorflow/tools/ci_build/linux/cpu/run_py3_contrib.sh b/tensorflow/tools/ci_build/linux/cpu/run_py3_contrib.sh
index f6fa925..d120275 100755
--- a/tensorflow/tools/ci_build/linux/cpu/run_py3_contrib.sh
+++ b/tensorflow/tools/ci_build/linux/cpu/run_py3_contrib.sh
@@ -26,6 +26,7 @@
# Run configure.
export TF_NEED_CUDA=0
+export TF_NEED_ROCM=0
export CC_OPT_FLAGS='-mavx'
export PYTHON_BIN_PATH=`which python3`
yes "" | $PYTHON_BIN_PATH configure.py
diff --git a/tensorflow/tools/ci_build/linux/cpu/run_py3_core.sh b/tensorflow/tools/ci_build/linux/cpu/run_py3_core.sh
index 51eb2cd..7c531a4 100755
--- a/tensorflow/tools/ci_build/linux/cpu/run_py3_core.sh
+++ b/tensorflow/tools/ci_build/linux/cpu/run_py3_core.sh
@@ -26,6 +26,7 @@
# Run configure.
export TF_NEED_CUDA=0
+export TF_NEED_ROCM=0
export CC_OPT_FLAGS='-mavx'
export PYTHON_BIN_PATH=`which python3`
yes "" | $PYTHON_BIN_PATH configure.py
diff --git a/tensorflow/tools/ci_build/linux/libtensorflow.sh b/tensorflow/tools/ci_build/linux/libtensorflow.sh
index beef8e0..3b6e15f 100755
--- a/tensorflow/tools/ci_build/linux/libtensorflow.sh
+++ b/tensorflow/tools/ci_build/linux/libtensorflow.sh
@@ -27,5 +27,8 @@
if [ "${TF_NEED_CUDA}" == "1" ]; then
SUFFIX="-gpu-linux-"
fi
+if [ "${TF_NEED_ROCM}" == "1" ]; then
+ SUFFIX="-rocm-linux-"
+fi
build_libtensorflow_tarball "${SUFFIX}$(uname -m)"
diff --git a/tensorflow/tools/ci_build/linux/libtensorflow_cpu.sh b/tensorflow/tools/ci_build/linux/libtensorflow_cpu.sh
index 4bf34dd..b76262b 100755
--- a/tensorflow/tools/ci_build/linux/libtensorflow_cpu.sh
+++ b/tensorflow/tools/ci_build/linux/libtensorflow_cpu.sh
@@ -19,4 +19,5 @@
set -ex
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
export TF_NEED_CUDA=0
+export TF_NEED_ROCM=0
"${SCRIPT_DIR}/libtensorflow_docker.sh"
diff --git a/tensorflow/tools/ci_build/linux/libtensorflow_docker.sh b/tensorflow/tools/ci_build/linux/libtensorflow_docker.sh
index 60c974c..467b8dc 100755
--- a/tensorflow/tools/ci_build/linux/libtensorflow_docker.sh
+++ b/tensorflow/tools/ci_build/linux/libtensorflow_docker.sh
@@ -38,6 +38,11 @@
DOCKER_BINARY="nvidia-docker"
DOCKER_FILE="Dockerfile.gpu"
fi
+if [ "${TF_NEED_ROCM}" == "1" ]; then
+ DOCKER_IMAGE="tf-tensorflow-rocm"
+ DOCKER_BINARY="docker"
+ DOCKER_FILE="Dockerfile.rocm"
+fi
docker build \
-t "${DOCKER_IMAGE}" \
@@ -53,6 +58,7 @@
-e "TF_NEED_HDFS=0" \
-e "TF_NEED_CUDA=${TF_NEED_CUDA}" \
-e "TF_NEED_TENSORRT=${TF_NEED_CUDA}" \
+ -e "TF_NEED_ROCM=${TF_NEED_ROCM}" \
-e "TF_NEED_OPENCL_SYCL=0" \
"${DOCKER_IMAGE}" \
"/workspace/tensorflow/tools/ci_build/linux/libtensorflow.sh"
diff --git a/tensorflow/tools/ci_build/linux/libtensorflow_rocm.sh b/tensorflow/tools/ci_build/linux/libtensorflow_rocm.sh
new file mode 100755
index 0000000..c1ebbe3
--- /dev/null
+++ b/tensorflow/tools/ci_build/linux/libtensorflow_rocm.sh
@@ -0,0 +1,22 @@
+#!/usr/bin/env bash
+# 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.
+# ==============================================================================
+#
+# Script to build a binary releases of libtensorflow with GPU support.
+
+set -ex
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+export TF_NEED_ROCM=1
+"${SCRIPT_DIR}/libtensorflow_docker.sh"
diff --git a/tensorflow/tools/ci_build/linux/rocm/run_cc_core.sh b/tensorflow/tools/ci_build/linux/rocm/run_cc_core.sh
new file mode 100755
index 0000000..200089f
--- /dev/null
+++ b/tensorflow/tools/ci_build/linux/rocm/run_cc_core.sh
@@ -0,0 +1,39 @@
+#!/usr/bin/env bash
+# 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.
+#
+# ==============================================================================
+
+set -e
+set -x
+
+N_JOBS=$(grep -c ^processor /proc/cpuinfo)
+
+echo ""
+echo "Bazel will use ${N_JOBS} concurrent job(s)."
+echo ""
+
+# Run configure.
+export PYTHON_BIN_PATH=`which python3`
+export CC_OPT_FLAGS='-mavx'
+
+export TF_NEED_ROCM=1
+
+yes "" | $PYTHON_BIN_PATH configure.py
+
+# Run bazel test command. Double test timeouts to avoid flakes.
+bazel test --config=rocm --test_tag_filters=-no_oss,-oss_serial,-no_gpu,-benchmark-test -k \
+ --test_lang_filters=cc --jobs=${N_JOBS} --test_timeout 300,450,1200,3600 \
+ --build_tests_only --test_output=errors --local_test_jobs=1 --config=opt \
+ //tensorflow/... -//tensorflow/compiler/... -//tensorflow/contrib/...
diff --git a/tensorflow/tools/ci_build/linux/rocm/run_py3_core.sh b/tensorflow/tools/ci_build/linux/rocm/run_py3_core.sh
new file mode 100755
index 0000000..1d0b838
--- /dev/null
+++ b/tensorflow/tools/ci_build/linux/rocm/run_py3_core.sh
@@ -0,0 +1,39 @@
+#!/usr/bin/env bash
+# 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.
+#
+# ==============================================================================
+
+set -e
+set -x
+
+N_JOBS=$(grep -c ^processor /proc/cpuinfo)
+
+echo ""
+echo "Bazel will use ${N_JOBS} concurrent job(s)."
+echo ""
+
+# Run configure.
+export PYTHON_BIN_PATH=`which python3`
+export CC_OPT_FLAGS='-mavx'
+
+export TF_NEED_ROCM=1
+
+yes "" | $PYTHON_BIN_PATH configure.py
+
+# Run bazel test command. Double test timeouts to avoid flakes.
+bazel test --config=rocm --test_tag_filters=-no_oss,-oss_serial,-no_gpu,-benchmark-test -k \
+ --test_lang_filters=py --jobs=${N_JOBS} --test_timeout 300,450,1200,3600 \
+ --build_tests_only --test_output=errors --local_test_jobs=1 --config=opt \
+ //tensorflow/... -//tensorflow/compiler/... -//tensorflow/contrib/...
diff --git a/tensorflow/tools/ci_build/osx/cpu/run_py2_cc_core.sh b/tensorflow/tools/ci_build/osx/cpu/run_py2_cc_core.sh
index c7cc16e..adee0d3 100755
--- a/tensorflow/tools/ci_build/osx/cpu/run_py2_cc_core.sh
+++ b/tensorflow/tools/ci_build/osx/cpu/run_py2_cc_core.sh
@@ -27,6 +27,7 @@
# Run configure.
export TF_NEED_CUDA=0
+export TF_NEED_ROCM=0
export CC_OPT_FLAGS='-mavx'
export PYTHON_BIN_PATH=$(which python2)
yes "" | $PYTHON_BIN_PATH configure.py
diff --git a/tensorflow/tools/ci_build/osx/libtensorflow_cpu.sh b/tensorflow/tools/ci_build/osx/libtensorflow_cpu.sh
index 9ae5fc6..06798ad 100755
--- a/tensorflow/tools/ci_build/osx/libtensorflow_cpu.sh
+++ b/tensorflow/tools/ci_build/osx/libtensorflow_cpu.sh
@@ -26,6 +26,7 @@
export PYTHON_BIN_PATH="/usr/bin/python"
export TF_NEED_HDFS=0
export TF_NEED_CUDA=0
+export TF_NEED_ROCM=0
export TF_NEED_OPENCL_SYCL=0
export TF_NEED_MKL=0
export COMPUTECPP_PATH="/usr/local"
diff --git a/tensorflow/tools/ci_build/osx/libtensorflow_gpu.sh b/tensorflow/tools/ci_build/osx/libtensorflow_gpu.sh
index d95fcde..95f1992 100755
--- a/tensorflow/tools/ci_build/osx/libtensorflow_gpu.sh
+++ b/tensorflow/tools/ci_build/osx/libtensorflow_gpu.sh
@@ -27,6 +27,7 @@
export LD_LIBRARY_PATH="/usr/local/cuda/lib:/usr/local/cuda/extras/CUPTI/lib:${LD_LIBRARY_PATH}"
export PYTHON_BIN_PATH="/usr/bin/python"
export TF_NEED_HDFS=0
+export TF_NEED_ROCM=0
export TF_NEED_OPENCL_SYCL=0
export TF_NEED_MKL=0
export COMPUTECPP_PATH="/usr/local"
diff --git a/tensorflow/tools/ci_build/osx/libtensorflow_rocm.sh b/tensorflow/tools/ci_build/osx/libtensorflow_rocm.sh
new file mode 100755
index 0000000..aeabc0e
--- /dev/null
+++ b/tensorflow/tools/ci_build/osx/libtensorflow_rocm.sh
@@ -0,0 +1,36 @@
+#!/usr/bin/env bash
+# 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.
+# ==============================================================================
+#
+# Script to produce binary release of libtensorflow (C API, Java jars etc.).
+
+set -ex
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+# See comments at the top of this file for details.
+source "${SCRIPT_DIR}/../builds/libtensorflow.sh"
+
+# Configure script
+export TF_NEED_ROCM=1
+export PYTHON_BIN_PATH="/usr/bin/python"
+export TF_NEED_GCP=0
+export TF_NEED_HDFS=0
+export TF_NEED_CUDA=0
+export TF_NEED_OPENCL_SYCL=0
+export TF_NEED_MKL=0
+export COMPUTECPP_PATH="/usr/local"
+
+export PATH="/usr/local/cuda/bin:/usr/local/bin:/usr/bin:/bin:/usr/sbin:/sbin"
+build_libtensorflow_tarball "-gpu-darwin-$(uname -m)"
diff --git a/tensorflow/tools/ci_build/xla/linux/rocm/run_py3.sh b/tensorflow/tools/ci_build/xla/linux/rocm/run_py3.sh
new file mode 100755
index 0000000..a0de128
--- /dev/null
+++ b/tensorflow/tools/ci_build/xla/linux/rocm/run_py3.sh
@@ -0,0 +1,41 @@
+#!/usr/bin/env bash
+# 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.
+#
+# ==============================================================================
+
+set -e
+set -x
+
+N_JOBS=$(grep -c ^processor /proc/cpuinfo)
+
+echo ""
+echo "Bazel will use ${N_JOBS} concurrent job(s)."
+echo ""
+
+# Run configure.
+export PYTHON_BIN_PATH=`which python3`
+
+export TF_NEED_ROCM=1
+
+yes "" | $PYTHON_BIN_PATH configure.py
+echo "build --distinct_host_configuration=false" >> .tf_configure.bazelrc
+
+bazel clean
+# Run bazel test command. Double test timeouts to avoid flakes.
+bazel test --config=rocm --test_tag_filters=-no_gpu,-benchmark-test,-no_oss -k \
+ --jobs=${N_JOBS} --test_timeout 300,450,1200,3600 \
+ --build_tests_only --test_output=errors --local_test_jobs=1 \
+ --config=xla -- \
+ //tensorflow/compiler/...
diff --git a/tensorflow/tools/dist_test/server/BUILD b/tensorflow/tools/dist_test/server/BUILD
index 003a19a..3aa53a5 100644
--- a/tensorflow/tools/dist_test/server/BUILD
+++ b/tensorflow/tools/dist_test/server/BUILD
@@ -8,6 +8,7 @@
exports_files(["LICENSE"])
load("//tensorflow:tensorflow.bzl", "py_test")
+load("//tensorflow:tensorflow.bzl", "py_binary")
py_binary(
name = "grpc_tensorflow_server",
diff --git a/tensorflow/tools/docker/Dockerfile b/tensorflow/tools/docker/Dockerfile
index b5a6c05..205128a 100644
--- a/tensorflow/tools/docker/Dockerfile
+++ b/tensorflow/tools/docker/Dockerfile
@@ -29,8 +29,8 @@
h5py \
ipykernel \
jupyter \
- keras_applications==1.0.5 \
- keras_preprocessing==1.0.3 \
+ keras_applications \
+ keras_preprocessing \
matplotlib \
numpy \
pandas \
diff --git a/tensorflow/tools/docker/Dockerfile.devel b/tensorflow/tools/docker/Dockerfile.devel
index c741e8a..6f8e91f 100644
--- a/tensorflow/tools/docker/Dockerfile.devel
+++ b/tensorflow/tools/docker/Dockerfile.devel
@@ -33,8 +33,8 @@
h5py \
ipykernel \
jupyter \
- keras_applications==1.0.5 \
- keras_preprocessing==1.0.3 \
+ keras_applications \
+ keras_preprocessing \
matplotlib \
mock \
numpy \
diff --git a/tensorflow/tools/docker/Dockerfile.devel-gpu b/tensorflow/tools/docker/Dockerfile.devel-gpu
index f544725..69a117f 100644
--- a/tensorflow/tools/docker/Dockerfile.devel-gpu
+++ b/tensorflow/tools/docker/Dockerfile.devel-gpu
@@ -55,8 +55,8 @@
h5py \
ipykernel \
jupyter \
- keras_applications==1.0.5 \
- keras_preprocessing==1.0.3 \
+ keras_applications \
+ keras_preprocessing \
matplotlib \
mock \
numpy \
diff --git a/tensorflow/tools/docker/Dockerfile.devel-mkl b/tensorflow/tools/docker/Dockerfile.devel-mkl
index db7c701..e433e9e 100755
--- a/tensorflow/tools/docker/Dockerfile.devel-mkl
+++ b/tensorflow/tools/docker/Dockerfile.devel-mkl
@@ -52,8 +52,8 @@
h5py \
ipykernel \
jupyter \
- keras_applications==1.0.5 \
- keras_preprocessing==1.0.3 \
+ keras_applications \
+ keras_preprocessing \
matplotlib \
mock \
numpy \
diff --git a/tensorflow/tools/docker/Dockerfile.devel-mkl-horovod b/tensorflow/tools/docker/Dockerfile.devel-mkl-horovod
index 987b582..48f2400 100755
--- a/tensorflow/tools/docker/Dockerfile.devel-mkl-horovod
+++ b/tensorflow/tools/docker/Dockerfile.devel-mkl-horovod
@@ -45,8 +45,8 @@
h5py \
ipykernel \
jupyter \
- keras_applications==1.0.5 \
- keras_preprocessing==1.0.3 \
+ keras_applications \
+ keras_preprocessing \
matplotlib \
mock \
numpy \
diff --git a/tensorflow/tools/docker/Dockerfile.gpu b/tensorflow/tools/docker/Dockerfile.gpu
index 781bf9e..7dc92a8 100644
--- a/tensorflow/tools/docker/Dockerfile.gpu
+++ b/tensorflow/tools/docker/Dockerfile.gpu
@@ -42,8 +42,8 @@
h5py \
ipykernel \
jupyter \
- keras_applications==1.0.5 \
- keras_preprocessing==1.0.3 \
+ keras_applications \
+ keras_preprocessing \
matplotlib \
numpy \
pandas \
diff --git a/tensorflow/tools/docker/Dockerfile.mkl b/tensorflow/tools/docker/Dockerfile.mkl
index 641c9e3..ac41cff 100755
--- a/tensorflow/tools/docker/Dockerfile.mkl
+++ b/tensorflow/tools/docker/Dockerfile.mkl
@@ -38,8 +38,8 @@
h5py \
ipykernel \
jupyter \
- keras_applications==1.0.5 \
- keras_preprocessing==1.0.3 \
+ keras_applications \
+ keras_preprocessing \
matplotlib \
numpy \
pandas \
diff --git a/tensorflow/tools/docker/Dockerfile.mkl-horovod b/tensorflow/tools/docker/Dockerfile.mkl-horovod
index 2b11679..4daf4fe 100755
--- a/tensorflow/tools/docker/Dockerfile.mkl-horovod
+++ b/tensorflow/tools/docker/Dockerfile.mkl-horovod
@@ -38,8 +38,8 @@
h5py \
ipykernel \
jupyter \
- keras_applications==1.0.5 \
- keras_preprocessing==1.0.3 \
+ keras_applications \
+ keras_preprocessing \
matplotlib \
numpy \
pandas \
diff --git a/tensorflow/tools/docs/BUILD b/tensorflow/tools/docs/BUILD
index b218e90..1a53f24 100644
--- a/tensorflow/tools/docs/BUILD
+++ b/tensorflow/tools/docs/BUILD
@@ -37,6 +37,7 @@
name = "doc_controls",
srcs = ["doc_controls.py"],
srcs_version = "PY2AND3",
+ visibility = ["//visibility:public"],
)
py_test(
@@ -126,7 +127,6 @@
name = "build_docs_test",
size = "small",
srcs = ["build_docs_test.py"],
- data = ["//tensorflow/docs_src"],
srcs_version = "PY2AND3",
tags = [
# No reason to run sanitizers or fastbuild for this test.
diff --git a/tensorflow/tools/docs/build_docs_test.py b/tensorflow/tools/docs/build_docs_test.py
index 0cbf8b4..4d3bedd 100644
--- a/tensorflow/tools/docs/build_docs_test.py
+++ b/tensorflow/tools/docs/build_docs_test.py
@@ -30,9 +30,11 @@
class Flags(object):
resource_root = resource_loader.get_root_dir_with_all_resources()
- src_dir = os.path.join(resource_root, 'tensorflow/docs_src')
+ src_dir = os.path.join(googletest.GetTempDir(), 'input')
+ os.mkdir(src_dir)
base_dir = os.path.join(resource_root, 'tensorflow/')
- output_dir = googletest.GetTempDir()
+ output_dir = os.path.join(googletest.GetTempDir(), 'output')
+ os.mkdir(output_dir)
class BuildDocsTest(googletest.TestCase):
diff --git a/tensorflow/tools/lib_package/BUILD b/tensorflow/tools/lib_package/BUILD
index b450bc4..85514b8 100644
--- a/tensorflow/tools/lib_package/BUILD
+++ b/tensorflow/tools/lib_package/BUILD
@@ -125,6 +125,7 @@
"@gemmlowp//:LICENSE",
"@gif_archive//:COPYING",
"@highwayhash//:LICENSE",
+ "@icu//:icu4c/LICENSE",
"@jpeg//:LICENSE.md",
"@llvm//:LICENSE.TXT",
"@lmdb//:LICENSE",
@@ -136,24 +137,6 @@
"@snappy//:COPYING",
"@zlib_archive//:zlib.h",
] + select({
- "//tensorflow:with_aws_support": [
- "@aws//:LICENSE",
- ],
- "//conditions:default": [],
- }) + select({
- "//tensorflow:with_gcp_support": [
- "@com_github_googlecloudplatform_google_cloud_cpp//:LICENSE",
- ],
- "//conditions:default": [],
- }) + select({
- "//tensorflow:with_jemalloc_linux_x86_64": [
- "@jemalloc//:COPYING",
- ],
- "//tensorflow:with_jemalloc_linux_ppc64le": [
- "@jemalloc//:COPYING",
- ],
- "//conditions:default": [],
- }) + select({
"//tensorflow/core/kernels:xsmm": [
"@libxsmm_archive//:LICENSE.md",
],
@@ -170,7 +153,14 @@
"@grpc//third_party/nanopb:LICENSE.txt",
"@grpc//third_party/address_sorting:LICENSE",
],
- ),
+ ) + select({
+ "//tensorflow:linux_s390x": [],
+ "//tensorflow:windows": [],
+ "//conditions:default": [
+ "@aws//:LICENSE",
+ "@com_github_googlecloudplatform_google_cloud_cpp//:LICENSE",
+ ],
+ }),
outs = ["include/tensorflow/c/LICENSE"],
cmd = "$(location :concat_licenses.sh) $(SRCS) >$@",
tools = [":concat_licenses.sh"],
@@ -192,6 +182,7 @@
"@gemmlowp//:LICENSE",
"@gif_archive//:COPYING",
"@highwayhash//:LICENSE",
+ "@icu//:icu4j/main/shared/licenses/LICENSE",
"@jpeg//:LICENSE.md",
"@llvm//:LICENSE.TXT",
"@lmdb//:LICENSE",
@@ -203,24 +194,6 @@
"@snappy//:COPYING",
"@zlib_archive//:zlib.h",
] + select({
- "//tensorflow:with_aws_support": [
- "@aws//:LICENSE",
- ],
- "//conditions:default": [],
- }) + select({
- "//tensorflow:with_gcp_support": [
- "@com_github_googlecloudplatform_google_cloud_cpp//:LICENSE",
- ],
- "//conditions:default": [],
- }) + select({
- "//tensorflow:with_jemalloc_linux_x86_64": [
- "@jemalloc//:COPYING",
- ],
- "//tensorflow:with_jemalloc_linux_ppc64le": [
- "@jemalloc//:COPYING",
- ],
- "//conditions:default": [],
- }) + select({
"//tensorflow/core/kernels:xsmm": [
"@libxsmm_archive//:LICENSE.md",
],
@@ -230,7 +203,14 @@
]) + if_mkl([
"//third_party/mkl:LICENSE",
"//third_party/mkl_dnn:LICENSE",
- ]),
+ ]) + select({
+ "//tensorflow:linux_s390x": [],
+ "//tensorflow:windows": [],
+ "//conditions:default": [
+ "@aws//:LICENSE",
+ "@com_github_googlecloudplatform_google_cloud_cpp//:LICENSE",
+ ],
+ }),
outs = ["include/tensorflow/jni/LICENSE"],
cmd = "$(location :concat_licenses.sh) $(SRCS) >$@",
tools = [":concat_licenses.sh"],
diff --git a/tensorflow/tools/pip_package/BUILD b/tensorflow/tools/pip_package/BUILD
index 12354a6..164b3d8 100644
--- a/tensorflow/tools/pip_package/BUILD
+++ b/tensorflow/tools/pip_package/BUILD
@@ -64,10 +64,6 @@
"//tensorflow/contrib/cluster_resolver:cluster_resolver_pip",
"//tensorflow/contrib/compiler:xla",
"//tensorflow/contrib/constrained_optimization:constrained_optimization_pip",
- "//tensorflow/contrib/data/python/kernel_tests/serialization:dataset_serialization_test_base",
- "//tensorflow/contrib/data/python/kernel_tests:stats_dataset_test_base",
- "//tensorflow/contrib/data/python/kernel_tests:test_utils",
- "//tensorflow/contrib/data/python/ops:contrib_op_loader",
"//tensorflow/contrib/eager/python/examples:examples_pip",
"//tensorflow/contrib/eager/python:evaluator",
"//tensorflow/contrib/gan:gan",
@@ -108,6 +104,9 @@
"//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/testdata:self_adjoint_eig_op_test_files",
@@ -152,6 +151,7 @@
"@gemmlowp//:LICENSE",
"@gif_archive//:COPYING",
"@highwayhash//:LICENSE",
+ "@icu//:icu4c/LICENSE",
"@jpeg//:LICENSE.md",
"@lmdb//:LICENSE",
"@local_config_sycl//sycl:LICENSE.text",
@@ -167,30 +167,6 @@
"@zlib_archive//:zlib.h",
"@org_python_pypi_backports_weakref//:LICENSE",
] + select({
- "//tensorflow:with_aws_support": [
- "@aws//:LICENSE",
- ],
- "//conditions:default": [],
- }) + select({
- "//tensorflow:with_gcp_support": [
- "@com_github_googleapis_googleapis//:LICENSE",
- "@com_github_googlecloudplatform_google_cloud_cpp//:LICENSE",
- ],
- "//conditions:default": [],
- }) + select({
- "//tensorflow:with_jemalloc_linux_x86_64": [
- "@jemalloc//:COPYING",
- ],
- "//tensorflow:with_jemalloc_linux_ppc64le": [
- "@jemalloc//:COPYING",
- ],
- "//conditions:default": [],
- }) + select({
- "//tensorflow:with_kafka_support": [
- "@kafka//:LICENSE",
- ],
- "//conditions:default": [],
- }) + select({
"//tensorflow/core/kernels:xsmm": [
"@libxsmm_archive//:LICENSE.md",
],
@@ -213,7 +189,16 @@
"@ngraph_tf//:LICENSE",
"@nlohmann_json_lib//:LICENSE.MIT",
"@tbb//:LICENSE",
- ]) + tf_additional_license_deps(),
+ ]) + tf_additional_license_deps() + select({
+ "//tensorflow:linux_s390x": [],
+ "//tensorflow:windows": [],
+ "//conditions:default": [
+ "@aws//:LICENSE",
+ "@com_github_googleapis_googleapis//:LICENSE",
+ "@com_github_googlecloudplatform_google_cloud_cpp//:LICENSE",
+ "@kafka//:LICENSE",
+ ],
+ }),
)
sh_binary(
diff --git a/tensorflow/tools/pip_package/pip_smoke_test.py b/tensorflow/tools/pip_package/pip_smoke_test.py
index bfc007b..c6ef82c 100644
--- a/tensorflow/tools/pip_package/pip_smoke_test.py
+++ b/tensorflow/tools/pip_package/pip_smoke_test.py
@@ -90,6 +90,7 @@
"//tensorflow/contrib/lite/python:interpreter.py",
"//tensorflow/contrib/lite/python:interpreter_test.py",
"//tensorflow/contrib/ffmpeg:test_data",
+ "//tensorflow/contrib/fused_conv:fused_conv2d_bias_activation_op_test_base",
"//tensorflow/contrib/hadoop:test_data",
"//tensorflow/contrib/factorization/examples:mnist",
"//tensorflow/contrib/factorization/examples:mnist.py",
diff --git a/tensorflow/tools/pip_package/setup.py b/tensorflow/tools/pip_package/setup.py
index b95e1f5..d864a7a 100644
--- a/tensorflow/tools/pip_package/setup.py
+++ b/tensorflow/tools/pip_package/setup.py
@@ -51,12 +51,11 @@
'absl-py >= 0.1.6',
'astor >= 0.6.0',
'gast >= 0.2.0',
- 'keras_applications >= 1.0.5',
- 'keras_preprocessing >= 1.0.3',
+ 'keras_applications >= 1.0.6',
+ 'keras_preprocessing >= 1.0.5',
'numpy >= 1.13.3',
'six >= 1.10.0',
- 'protobuf >= 3.6.0',
- 'setuptools <= 39.1.0',
+ 'protobuf >= 3.6.1',
'tensorboard >= 1.11.0, < 1.12.0',
'termcolor >= 1.1.0',
]
diff --git a/tensorflow/tools/quantization/BUILD b/tensorflow/tools/quantization/BUILD
deleted file mode 100644
index 17443a8..0000000
--- a/tensorflow/tools/quantization/BUILD
+++ /dev/null
@@ -1,78 +0,0 @@
-# Description:
-# Utilities for quantizing TensorFlow graphs to lower bit depths.
-
-package(default_visibility = ["//visibility:public"])
-
-licenses(["notice"]) # Apache 2.0
-
-exports_files(["LICENSE"])
-
-load("//tensorflow:tensorflow.bzl", "py_test")
-
-py_library(
- name = "quantize_graph_lib",
- srcs = ["quantize_graph.py"],
- srcs_version = "PY2AND3",
- deps = [
- "//tensorflow/core:protos_all_py",
- "//tensorflow/python:array_ops",
- "//tensorflow/python:constant_op",
- "//tensorflow/python:dtypes",
- "//tensorflow/python:framework",
- "//tensorflow/python:framework_ops",
- "//tensorflow/python:graph_util",
- "//tensorflow/python:platform",
- "//tensorflow/python:session",
- "//tensorflow/python:tensor_shape",
- "//tensorflow/python:tensor_util",
- "//third_party/py/numpy",
- ],
-)
-
-py_binary(
- name = "quantize_graph",
- srcs = ["quantize_graph.py"],
- srcs_version = "PY2AND3",
- deps = [
- "//tensorflow/core:protos_all_py",
- "//tensorflow/python", # TODO(b/34059704): remove when fixed
- "//tensorflow/python:array_ops",
- "//tensorflow/python:client",
- "//tensorflow/python:framework",
- "//tensorflow/python:framework_for_generated_wrappers",
- "//tensorflow/python:graph_util",
- "//tensorflow/python:platform",
- "//tensorflow/python:tensor_util",
- "//third_party/py/numpy",
- ],
-)
-
-py_test(
- name = "quantize_graph_test",
- size = "small",
- srcs = ["quantize_graph_test.py"],
- srcs_version = "PY2AND3",
- tags = ["nomsan"], # http://b/32242946
- deps = [
- ":quantize_graph",
- "//tensorflow/core:protos_all_py",
- "//tensorflow/python:client",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:framework",
- "//tensorflow/python:framework_for_generated_wrappers",
- "//tensorflow/python:graph_util",
- "//tensorflow/python:platform",
- "//third_party/py/numpy",
- ],
-)
-
-py_binary(
- name = "graph_to_dot",
- srcs = ["graph_to_dot.py"],
- main = "graph_to_dot.py",
- srcs_version = "PY2AND3",
- deps = [
- "//tensorflow/core:protos_all_py",
- "//tensorflow/python:platform",
- ],
-)
diff --git a/tensorflow/tools/quantization/graph_to_dot.py b/tensorflow/tools/quantization/graph_to_dot.py
deleted file mode 100644
index 81d6aa6..0000000
--- a/tensorflow/tools/quantization/graph_to_dot.py
+++ /dev/null
@@ -1,68 +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.
-# ==============================================================================
-"""Converts a GraphDef file into a DOT format suitable for visualization.
-
-This script takes a GraphDef representing a network, and produces a DOT file
-that can then be visualized by GraphViz tools like dot and xdot.
-
-"""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import re
-
-from google.protobuf import text_format
-
-from tensorflow.core.framework import graph_pb2
-from tensorflow.python.platform import app
-from tensorflow.python.platform import flags
-from tensorflow.python.platform import gfile
-
-FLAGS = flags.FLAGS
-
-flags.DEFINE_string("graph", "", """TensorFlow 'GraphDef' file to load.""")
-flags.DEFINE_bool("input_binary", True,
- """Whether the input files are in binary format.""")
-flags.DEFINE_string("dot_output", "", """Where to write the DOT output.""")
-
-
-def main(unused_args):
- if not gfile.Exists(FLAGS.graph):
- print("Input graph file '" + FLAGS.graph + "' does not exist!")
- return -1
-
- graph = graph_pb2.GraphDef()
- with open(FLAGS.graph, "r") as f:
- if FLAGS.input_binary:
- graph.ParseFromString(f.read())
- else:
- text_format.Merge(f.read(), graph)
-
- with open(FLAGS.dot_output, "wb") as f:
- print("digraph graphname {", file=f)
- for node in graph.node:
- output_name = node.name
- print(" \"" + output_name + "\" [label=\"" + node.op + "\"];", file=f)
- for input_full_name in node.input:
- parts = input_full_name.split(":")
- input_name = re.sub(r"^\^", "", parts[0])
- print(" \"" + input_name + "\" -> \"" + output_name + "\";", file=f)
- print("}", file=f)
- print("Created DOT file '" + FLAGS.dot_output + "'.")
-
-
-if __name__ == "__main__":
- app.run()
diff --git a/tensorflow/tools/quantization/quantize_graph.py b/tensorflow/tools/quantization/quantize_graph.py
deleted file mode 100644
index 3acb532..0000000
--- a/tensorflow/tools/quantization/quantize_graph.py
+++ /dev/null
@@ -1,1302 +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.
-# ==============================================================================
-r"""Transforms a float-trained graph into an equivalent quantized version.
-
-An example of command-line usage is:
-bazel build tensorflow/tools/quantization:quantize_graph \
-&& bazel-bin/tensorflow/tools/quantization/quantize_graph \
---input=tensorflow_inception_graph.pb
---output_node_names="softmax2" --print_nodes --output=/tmp/quantized_graph.pb \
---mode=eightbit --logtostderr
-
-"""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import collections
-import re
-import numpy as np
-
-from tensorflow.core.framework import attr_value_pb2
-from tensorflow.core.framework import graph_pb2
-from tensorflow.core.framework import node_def_pb2
-from tensorflow.python.client import session
-from tensorflow.python.framework import constant_op
-from tensorflow.python.framework import dtypes
-from tensorflow.python.framework import graph_util
-from tensorflow.python.framework import importer
-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.platform import app
-from tensorflow.python.platform import flags as flags_lib
-from tensorflow.python.platform import gfile
-
-flags = flags_lib
-FLAGS = flags.FLAGS
-
-flags.DEFINE_boolean("print_nodes", False, """Lists all nodes in the model.""")
-flags.DEFINE_string("input", "", """TensorFlow 'GraphDef' file to load.""")
-flags.DEFINE_string("output_node_names", "",
- """Output node names, comma separated.""")
-flags.DEFINE_string("output", "", """File to save the output graph to.""")
-flags.DEFINE_integer("bitdepth", 8,
- """How many bits to quantize the graph to.""")
-flags.DEFINE_string("mode", "round",
- """What transformation to apply (round, quantize,"""
- """ eightbit, weights, or weights_rounded).""")
-flags.DEFINE_string("test_input_dims", "1,224,224,3",
- """The size of the input tensor to use when testing a"""
- """ graph loaded from a file.""")
-flags.DEFINE_boolean("strip_redundant_quantization", True,
- """Removes redundant dequantize/quantize pairs.""")
-flags.DEFINE_boolean("quantized_input", False,
- "If true, assume Placeholders are quantized with values "
- "covering [--quantized_input_min,--quantized_input_max]. "
- "Only supported when --mode=eightbit")
-flags.DEFINE_float("quantized_input_min", 0,
- "The minimum of the actual input range when "
- "--quantized_input")
-flags.DEFINE_float("quantized_input_max", 1,
- "The maximum of the actual input range when "
- "--quantized_input")
-flags.DEFINE_float(
- "quantized_fallback_min", None,
- "The fallback 'min' value to use for layers which lack min-max "
- "information. Note: this should be considered a coarse tool just good "
- "enough for experimentation purposes, since graphs quantized in this way "
- "would be very inaccurate.")
-flags.DEFINE_float(
- "quantized_fallback_max", None,
- "The fallback 'max' value to use for layers which lack min-max "
- "information. Note: this should be considered a coarse tool just good "
- "enough for experimentation purposes, since graphs quantized in this way "
- "would be very inaccurate.")
-
-
-def print_input_nodes(current_node, nodes_map, indent, already_visited):
- print(" " * indent + current_node.op + ":" + current_node.name)
- already_visited[current_node.name] = True
- for input_node_name in current_node.input:
- if input_node_name in already_visited:
- continue
- input_node = nodes_map[input_node_name]
- print_input_nodes(input_node, nodes_map, indent + 1, already_visited)
-
-
-def create_node(op, name, inputs):
- new_node = node_def_pb2.NodeDef()
- new_node.op = op
- new_node.name = name
- for input_name in inputs:
- new_node.input.extend([input_name])
- return new_node
-
-
-def create_constant_node(name, value, dtype, shape=None):
- node = create_node("Const", name, [])
- set_attr_dtype(node, "dtype", dtype)
- set_attr_tensor(node, "value", value, dtype, shape)
- return node
-
-
-def copy_attr(node, key, attr_value):
- try:
- node.attr[key].CopyFrom(attr_value)
- except KeyError:
- pass
-
-
-def set_attr_dtype(node, key, value):
- try:
- node.attr[key].CopyFrom(
- attr_value_pb2.AttrValue(type=value.as_datatype_enum))
- except KeyError:
- pass
-
-
-def set_attr_shape(node, key, value):
- try:
- node.attr[key].CopyFrom(
- attr_value_pb2.AttrValue(shape=tensor_shape.as_shape(value).as_proto()))
- except KeyError:
- pass
-
-
-def set_attr_tensor(node, key, value, dtype, shape=None):
- try:
- node.attr[key].CopyFrom(
- attr_value_pb2.AttrValue(tensor=tensor_util.make_tensor_proto(
- value, dtype=dtype, shape=shape)))
- except KeyError:
- pass
-
-
-def set_attr_string(node, key, value):
- try:
- node.attr[key].CopyFrom(attr_value_pb2.AttrValue(s=value))
- except KeyError:
- pass
-
-
-def set_attr_int_list(node, key, value):
- list_value = attr_value_pb2.AttrValue.ListValue(i=value)
- try:
- node.attr[key].CopyFrom(attr_value_pb2.AttrValue(list=list_value))
- except KeyError:
- pass
-
-
-def set_attr_bool(node, key, value):
- try:
- node.attr[key].CopyFrom(attr_value_pb2.AttrValue(b=value))
- except KeyError:
- pass
-
-
-def set_attr_int(node, key, value):
- try:
- node.attr[key].CopyFrom(attr_value_pb2.AttrValue(i=value))
- except KeyError:
- pass
-
-
-def set_attr_float(node, key, value):
- try:
- node.attr[key].CopyFrom(attr_value_pb2.AttrValue(f=value))
- except KeyError:
- pass
-
-
-def node_name_from_input(node_name):
- """Strips off ports and other decorations to get the underlying node name."""
- if node_name.startswith("^"):
- node_name = node_name[1:]
- m = re.search(r"(.*):\d+$", node_name)
- if m:
- node_name = m.group(1)
- return node_name
-
-
-def ensure_tensor_name_has_port(node_name):
- """Makes sure that a tensor name has :0 if no explicit port exists."""
- m = re.search(r"(.*):\d+$", node_name)
- if m:
- name_with_port = node_name
- else:
- name_with_port = node_name + ":0"
- return name_with_port
-
-
-def unique_node_name_from_input(node_name):
- """Replaces invalid characters in input names to get a unique node name."""
- return node_name.replace(":", "__port__").replace("^", "__hat__")
-
-
-def quantize_array(arr, num_buckets):
- """Quantizes a numpy array.
-
- This function maps each scalar in arr to the center of one of num_buckets
- buckets. For instance,
- quantize_array([0, 0.3, 0.6, 1], 2) => [0.25, 0.25, 0.75, 0.75]
-
- Args:
- arr: The numpy array to quantize.
- num_buckets: The number of buckets to map "var" to.
- Returns:
- The quantized numpy array.
- Raises:
- ValueError: when num_buckets < 1.
- """
- if num_buckets < 1:
- raise ValueError("num_buckets must be >= 1")
- arr_max = arr.max()
- arr_min = arr.min()
- if arr_max == arr_min:
- return arr
- bucket_width = (arr_max - arr_min) / num_buckets
- # Map scalars to bucket indices. Take special care of max(arr).
- bucket_indices = np.floor((arr - arr_min) / bucket_width)
- bucket_indices[bucket_indices == num_buckets] = num_buckets - 1
- # Map each scalar to the center of a bucket.
- arr = arr_min + bucket_width * (bucket_indices + 0.5)
- return arr
-
-
-def quantize_weight_rounded(input_node):
- """Returns a replacement node for input_node containing bucketed floats."""
- input_tensor = input_node.attr["value"].tensor
- tensor_value = tensor_util.MakeNdarray(input_tensor)
- shape = input_tensor.tensor_shape
- # Currently, the parameter FLAGS.bitdepth is used to compute the
- # number of buckets as 1 << FLAGS.bitdepth, meaning the number of
- # buckets can only be a power of 2.
- # This could be fixed by introducing a new parameter, num_buckets,
- # which would allow for more flexibility in chosing the right model
- # size/accuracy tradeoff. But I didn't want to add more parameters
- # to this script than absolutely necessary.
- num_buckets = 1 << FLAGS.bitdepth
- tensor_value_rounded = quantize_array(tensor_value, num_buckets)
- tensor_shape_list = tensor_util.TensorShapeProtoToList(shape)
- return [
- create_constant_node(
- input_node.name,
- tensor_value_rounded,
- dtypes.float32,
- shape=tensor_shape_list)
- ]
-
-
-def quantize_weight_eightbit(input_node, quantization_mode):
- """Returns replacement nodes for input_node using the Dequantize op."""
- base_name = input_node.name + "_"
- quint8_const_name = base_name + "quint8_const"
- min_name = base_name + "min"
- max_name = base_name + "max"
- float_tensor = tensor_util.MakeNdarray(input_node.attr["value"].tensor)
- min_value = np.min(float_tensor.flatten())
- max_value = np.max(float_tensor.flatten())
- # Make sure that the range includes zero.
- if min_value > 0.0:
- min_value = 0.0
- # min_value == max_value is a tricky case. It can occur for general
- # tensors, and of course for scalars. The quantized ops cannot deal
- # with this case, so we set max_value to something else.
- # It's a tricky question what is the numerically best solution to
- # deal with this degeneracy.
- # TODO(petewarden): Better use a tolerance than a hard comparison?
- if min_value == max_value:
- if abs(min_value) < 0.000001:
- max_value = min_value + 1.0
- elif min_value > 0:
- max_value = 2 * min_value
- else:
- max_value = min_value / 2.0
-
- sess = session.Session()
- with sess.as_default():
- quantize_op = array_ops.quantize_v2(
- float_tensor,
- min_value,
- max_value,
- dtypes.quint8,
- mode=quantization_mode)
- quint8_tensor = quantize_op[0].eval()
- shape = tensor_util.TensorShapeProtoToList(input_node.attr["value"]
- .tensor.tensor_shape)
- quint8_const_node = create_constant_node(
- quint8_const_name, quint8_tensor, dtypes.quint8, shape=shape)
- min_node = create_constant_node(min_name, min_value, dtypes.float32)
- max_node = create_constant_node(max_name, max_value, dtypes.float32)
- dequantize_node = create_node("Dequantize", input_node.name,
- [quint8_const_name, min_name, max_name])
- set_attr_dtype(dequantize_node, "T", dtypes.quint8)
- set_attr_string(dequantize_node, "mode", quantization_mode)
- return [quint8_const_node, min_node, max_node, dequantize_node]
-
-
-EightbitizeRecursionState = collections.namedtuple(
- "EightbitizeRecursionState",
- ["already_visited", "output_node_stack", "merged_with_fake_quant"])
-
-
-class GraphRewriter(object):
- """Takes a float graph, and rewrites it in quantized form."""
-
- def __init__(self,
- input_graph,
- mode,
- quantized_input_range,
- fallback_quantization_range=None):
- """Sets up the class to rewrite a float graph.
-
- Args:
- input_graph: A float graph to transform.
- mode: A string controlling how quantization is performed -
- round, quantize, eightbit, or weights.
- quantized_input_range: if set, assume the input is
- quantized and represents the range
- [quantized_input_range[0], quantized_input_range[1]]
- fallback_quantization_range: if set, then for nodes where the quantization
- range can't be inferred from the graph, use the range
- [fallback_quantization_range[0], fallback_quantization_range[1]) instead
- of using a RequantizationRange node in the graph.
-
- Raises:
- ValueError: Two nodes with the same name were found in the graph.
- """
- self.input_graph = input_graph
- self.nodes_map = self.create_nodes_map(input_graph)
- self.output_graph = None
- self.mode = mode
- self.final_node_renames = {}
- if quantized_input_range:
- self.input_range = (quantized_input_range[0], quantized_input_range[1])
- if self.input_range[0] >= self.input_range[1]:
- raise ValueError("Invalid quantized_input_range: [%s,%s]" %
- self.input_range)
- if self.mode != "eightbit":
- raise ValueError(
- "quantized_input_range can only be specified in eightbit mode")
- else:
- self.input_range = None
-
- if fallback_quantization_range:
- self.fallback_quantization_range = [
- fallback_quantization_range[0], fallback_quantization_range[1]
- ]
- if (self.fallback_quantization_range[0] >=
- self.fallback_quantization_range[1]):
- raise ValueError("Invalid fallback_quantization_range: [%s,%s]" %
- self.fallback_quantization_range)
- if self.mode != "eightbit":
- raise ValueError("fallback_quantization_range can only be "
- "specified in eightbit mode")
- else:
- self.fallback_quantization_range = None
-
- # Data that is valid only during the recursive call to rewrite the graph.
- self.state = None
-
- def create_nodes_map(self, graph):
- """Builds a mapping of node names to their defs from the graph."""
- nodes_map = {}
- for node in graph.node:
- if node.name not in nodes_map.keys():
- nodes_map[node.name] = node
- else:
- raise ValueError("Duplicate node names detected.")
- return nodes_map
-
- def rewrite(self, output_node_names):
- """Triggers rewriting of the float graph.
-
- Args:
- output_node_names: A list of names of the nodes that produce the final
- results.
-
- Returns:
- A quantized version of the float graph.
- """
- self.output_graph = graph_pb2.GraphDef()
- output_nodes = [
- self.nodes_map[output_node_name]
- for output_node_name in output_node_names
- ]
- if self.mode == "round":
- self.already_visited = {}
- for output_node in output_nodes:
- self.round_nodes_recursively(output_node)
- elif self.mode == "quantize":
- self.already_visited = {}
- self.already_quantized = {}
- for output_node in output_nodes:
- self.quantize_nodes_recursively(output_node)
- elif self.mode == "eightbit":
- self.set_input_graph(graph_util.remove_training_nodes(
- self.input_graph, protected_nodes=output_node_names))
- output_nodes = [
- self.nodes_map[output_node_name]
- for output_node_name in output_node_names
- ]
-
- self.state = EightbitizeRecursionState(
- already_visited={}, output_node_stack=[], merged_with_fake_quant={})
- for output_node in output_nodes:
- self.eightbitize_nodes_recursively(output_node)
- self.state = None
- if self.input_range:
- self.add_output_graph_node(
- create_constant_node("quantized_input_min_value", self.input_range[
- 0], dtypes.float32, []))
- self.add_output_graph_node(
- create_constant_node("quantized_input_max_value", self.input_range[
- 1], dtypes.float32, []))
- if self.fallback_quantization_range:
- self.add_output_graph_node(
- create_constant_node("fallback_quantization_min_value",
- self.fallback_quantization_range[0],
- dtypes.float32, []))
- self.add_output_graph_node(
- create_constant_node("fallback_quantization_max_value",
- self.fallback_quantization_range[1],
- dtypes.float32, []))
- if FLAGS.strip_redundant_quantization:
- self.output_graph = self.remove_redundant_quantization(
- self.output_graph)
- self.remove_dead_nodes(output_node_names)
- self.apply_final_node_renames()
- elif self.mode == "weights":
- self.output_graph = self.quantize_weights(self.input_graph,
- b"MIN_COMBINED")
- self.remove_dead_nodes(output_node_names)
- elif self.mode == "weights_rounded":
- self.output_graph = self.quantize_weights(self.input_graph, self.mode)
- self.remove_dead_nodes(output_node_names)
- else:
- print("Bad mode - " + self.mode + ".")
- return self.output_graph
-
- def round_nodes_recursively(self, current_node):
- """The entry point for simple rounding quantization."""
- if (current_node.name in self.already_visited
- ) and self.already_visited[current_node.name]:
- return
- self.already_visited[current_node.name] = True
- for input_node_name in current_node.input:
- input_node_name = node_name_from_input(input_node_name)
- input_node = self.nodes_map[input_node_name]
- self.round_nodes_recursively(input_node)
- nodes_to_quantize = ["Conv2D", "BiasAdd", "MatMul"]
- if any(current_node.op in s for s in nodes_to_quantize):
- new_node = node_def_pb2.NodeDef()
- new_node.CopyFrom(current_node)
- new_node.name = current_node.name + "_original"
- self.add_output_graph_node(new_node)
- levels = 1 << FLAGS.bitdepth
- constant_name = current_node.name + "_round_depth"
- constant_tensor = constant_op.constant(
- levels, dtype=dtypes.int32, name=constant_name)
- constant_node = constant_tensor.op.node_def
- self.add_output_graph_node(constant_node)
- quantize_node = node_def_pb2.NodeDef()
- quantize_node.op = "RoundToSteps"
- quantize_node.name = current_node.name
- quantize_node.input.extend([current_node.name + "_original"])
- quantize_node.input.extend([constant_node.name])
- self.add_output_graph_node(quantize_node)
- else:
- new_node = node_def_pb2.NodeDef()
- new_node.CopyFrom(current_node)
- self.add_output_graph_node(new_node)
-
- def quantize_nodes_recursively(self, current_node):
- """The entry point for quantizing nodes to eight bit and back."""
- if self.already_visited[current_node.name]:
- return
- self.already_visited[current_node.name] = True
- for input_node_name in current_node.input:
- input_node_name = node_name_from_input(input_node_name)
- input_node = self.nodes_map[input_node_name]
- self.quantize_nodes_recursively(input_node)
- nodes_to_quantize = ["Conv2D", "BiasAdd", "MatMul"]
- if any(current_node.op in s for s in nodes_to_quantize):
- for input_name in current_node.input:
- input_name = node_name_from_input(input_name)
- input_node = self.nodes_map[input_name]
- self.quantize_node(input_node)
- self.quantize_node(current_node)
- else:
- new_node = node_def_pb2.NodeDef()
- new_node.CopyFrom(current_node)
- self.add_output_graph_node(new_node)
-
- def quantize_node(self, input_node):
- """Handles quantizing a single node."""
- input_name = input_node.name
- if input_name in self.already_quantized:
- return
- self.already_quantized[input_name] = True
- original_input_name = input_name + "_original"
- reshape_name = input_name + "_reshape"
- reshape_dims_name = input_name + "_reshape_dims"
- max_name = input_name + "_max"
- min_name = input_name + "_min"
- dims_name = input_name + "_dims"
- quantize_name = input_name + "_quantize"
- dequantize_name = input_name
- original_input_node = node_def_pb2.NodeDef()
- original_input_node.CopyFrom(input_node)
- original_input_node.name = original_input_name
- self.add_output_graph_node(original_input_node)
- reshape_dims_node = create_constant_node(reshape_dims_name, -1,
- dtypes.int32, [1])
- self.add_output_graph_node(reshape_dims_node)
- reshape_node = create_node("Reshape", reshape_name,
- [original_input_name, reshape_dims_name])
- set_attr_dtype(reshape_node, "T", dtypes.float32)
- self.add_output_graph_node(reshape_node)
- dims_node = create_constant_node(dims_name, 0, dtypes.int32, [1])
- self.add_output_graph_node(dims_node)
- max_node = create_node("Max", max_name, [reshape_name, dims_name])
- set_attr_dtype(max_node, "T", dtypes.float32)
- set_attr_bool(max_node, "keep_dims", False)
- self.add_output_graph_node(max_node)
- min_node = create_node("Min", min_name, [reshape_name, dims_name])
- set_attr_dtype(min_node, "T", dtypes.float32)
- set_attr_bool(min_node, "keep_dims", False)
- self.add_output_graph_node(min_node)
- quantize_node = create_node("Quantize", quantize_name,
- [original_input_name, min_name, max_name])
- set_attr_dtype(quantize_node, "T", dtypes.quint8)
- set_attr_string(quantize_node, "mode", b"MIN_FIRST")
- self.add_output_graph_node(quantize_node)
- dequantize_node = create_node("Dequantize", dequantize_name,
- [quantize_name, min_name, max_name])
- set_attr_dtype(dequantize_node, "T", dtypes.quint8)
- set_attr_string(dequantize_node, "mode", b"MIN_FIRST")
- self.add_output_graph_node(dequantize_node)
-
- def should_merge_with_fake_quant_node(self):
- """Should the current node merge with self.state.output_node_stack[-1]?"""
- if not self.state.output_node_stack:
- return False
- top = self.state.output_node_stack[-1]
- return top[1] == 0 and top[0].op in ["FakeQuantWithMinMaxVars"]
-
- def should_quantize_const(self, node):
- if not self.state.output_node_stack:
- return False
- top = self.state.output_node_stack[-1]
- if not top[2]:
- return False
- dtype = dtypes.as_dtype(node.attr["dtype"].type)
- assert dtype == dtypes.float32, (
- "Failed to quantized constant %s of type %s" % (node.name, dtype))
- return True
-
- def eightbitize_nodes_recursively(self, current_node):
- """The entry point for transforming a graph into full eight bit."""
- if current_node.name in self.state.already_visited:
- if (self.should_merge_with_fake_quant_node() or
- current_node.name in self.state.merged_with_fake_quant):
- raise ValueError("Unsupported graph structure: output of node %s "
- "is processed by a FakeQuant* node and should have "
- "no other outputs.", current_node.name)
- return
- self.state.already_visited[current_node.name] = True
-
- for i, input_node_name in enumerate(current_node.input):
- quantize_input = False
- if current_node.op in ("MatMul", "Conv2D", "BiasAdd", "MaxPool",
- "AvgPool", "Relu", "Relu6",
- "BatchNormWithGlobalNormalization"):
- quantize_input = True
- elif current_node.op == "Concat" and i > 0:
- quantize_input = (
- dtypes.as_dtype(current_node.attr["T"].type) == dtypes.float32)
- elif current_node.op == "Reshape" and i == 0:
- quantize_input = (
- dtypes.as_dtype(current_node.attr["T"].type) == dtypes.float32)
-
- self.state.output_node_stack.append((current_node, i, quantize_input))
-
- input_node_name = node_name_from_input(input_node_name)
- input_node = self.nodes_map[input_node_name]
- self.eightbitize_nodes_recursively(input_node)
-
- self.state.output_node_stack.pop()
-
- if current_node.op == "MatMul":
- self.eightbitize_mat_mul_node(current_node)
- elif current_node.op == "Conv2D":
- self.eightbitize_conv_node(current_node)
- elif current_node.op == "BiasAdd":
- self.eightbitize_bias_add_node(current_node)
- elif current_node.op == "MaxPool" or current_node.op == "AvgPool":
- self.eightbitize_single_input_tensor_node(current_node,
- self.add_pool_function)
- elif current_node.op == "Relu" or current_node.op == "Relu6":
- self.eightbitize_single_input_tensor_node(current_node,
- self.add_relu_function)
- elif (current_node.op == "Concat" and
- dtypes.as_dtype(current_node.attr["T"].type) == dtypes.float32):
- self.eightbitize_concat_node(current_node)
- elif current_node.op == "BatchNormWithGlobalNormalization":
- self.eightbitize_batch_norm_node(current_node)
- elif (current_node.op == "Reshape" and
- dtypes.as_dtype(current_node.attr["T"].type) == dtypes.float32):
- self.eightbitize_reshape_node(current_node)
- elif (self.input_range and
- current_node.op in ("Placeholder", "PlaceholderV2")):
- self.eightbitize_placeholder_node(current_node)
- elif current_node.op == "FakeQuantWithMinMaxVars":
- # It will have been merged into the underlying node.
- pass
- elif current_node.op == "Const":
- if self.should_quantize_const(current_node):
- for n in quantize_weight_eightbit(current_node, b"MIN_FIRST"):
- self.add_output_graph_node(n)
- else:
- new_node = node_def_pb2.NodeDef()
- new_node.CopyFrom(current_node)
- self.add_output_graph_node(new_node)
-
- ###################################################################
- # Note: if more cases are added here, you may need to update the op
- # name lists in the loop over children at the start of the function.
- ###################################################################
- else:
- new_node = node_def_pb2.NodeDef()
- new_node.CopyFrom(current_node)
- self.add_output_graph_node(new_node)
-
- if (self.should_merge_with_fake_quant_node() and
- current_node.name not in self.state.merged_with_fake_quant):
- raise ValueError(
- "FakeQuant* node %s failed to merge with node %s of type %s" %
- (self.state.output_node_stack[-1][0], current_node.name,
- current_node.op))
-
- def add_eightbit_prologue_nodes(self, original_node):
- """Adds input conversion nodes to handle quantizing the underlying node."""
- namespace_prefix = original_node.name + "_eightbit"
- reshape_dims_name, reduction_dims_name = self.add_common_quantization_nodes(
- namespace_prefix)
- input_names = []
- min_max_names = []
- for original_input_name in original_node.input:
- quantize_input_name, min_input_name, max_input_name = (
- self.eightbitize_input_to_node(namespace_prefix, original_input_name,
- reshape_dims_name,
- reduction_dims_name))
- input_names.append(quantize_input_name)
- min_max_names.append(min_input_name)
- min_max_names.append(max_input_name)
- all_input_names = []
- all_input_names.extend(input_names)
- all_input_names.extend(min_max_names)
- return all_input_names
-
- def add_common_quantization_nodes(self, namespace_prefix):
- """Builds constant nodes needed for quantization of inputs."""
- reshape_dims_name = namespace_prefix + "_reshape_dims"
- reduction_dims_name = namespace_prefix + "_reduction_dims"
-
- reshape_dims_node = create_constant_node(reshape_dims_name, -1,
- dtypes.int32, [1])
- self.add_output_graph_node(reshape_dims_node)
- reduction_dims_node = create_constant_node(reduction_dims_name, 0,
- dtypes.int32, [1])
- self.add_output_graph_node(reduction_dims_node)
- return reshape_dims_name, reduction_dims_name
-
- def eightbitize_input_to_node(self, namespace_prefix, original_input_name,
- reshape_dims_name, reduction_dims_name):
- """Takes one float input to an op, and converts it to quantized form."""
- unique_input_name = unique_node_name_from_input(original_input_name)
- reshape_input_name = namespace_prefix + "_reshape_" + unique_input_name
- min_input_name = namespace_prefix + "_min_" + unique_input_name
- max_input_name = namespace_prefix + "_max_" + unique_input_name
- quantize_input_name = namespace_prefix + "_quantize_" + unique_input_name
- reshape_input_node = create_node("Reshape", reshape_input_name,
- [original_input_name, reshape_dims_name])
- set_attr_dtype(reshape_input_node, "T", dtypes.float32)
- self.add_output_graph_node(reshape_input_node)
- min_input_node = create_node("Min", min_input_name,
- [reshape_input_name, reduction_dims_name])
- set_attr_dtype(min_input_node, "T", dtypes.float32)
- set_attr_bool(min_input_node, "keep_dims", False)
- self.add_output_graph_node(min_input_node)
- max_input_node = create_node("Max", max_input_name,
- [reshape_input_name, reduction_dims_name])
- set_attr_dtype(max_input_node, "T", dtypes.float32)
- set_attr_bool(max_input_node, "keep_dims", False)
- self.add_output_graph_node(max_input_node)
- quantize_input_node = create_node(
- "QuantizeV2", quantize_input_name,
- [original_input_name, min_input_name, max_input_name])
- set_attr_dtype(quantize_input_node, "T", dtypes.quint8)
- set_attr_string(quantize_input_node, "mode", b"MIN_FIRST")
- self.add_output_graph_node(quantize_input_node)
- min_output_name = quantize_input_name + ":1"
- max_output_name = quantize_input_name + ":2"
- return quantize_input_name, min_output_name, max_output_name
-
- def add_quantize_down_nodes(self, original_node, quantized_output_name):
- quantized_outputs = [
- quantized_output_name, quantized_output_name + ":1",
- quantized_output_name + ":2"
- ]
- min_max_inputs = None
- if self.should_merge_with_fake_quant_node():
- # Use the inputs to the FakeQuantWithMinMaxVars node as the inputs to
- # Requantize.
- fake_quant_node = self.state.output_node_stack[-1][0]
- min_max_inputs = [fake_quant_node.input[1], fake_quant_node.input[2]]
- assert original_node.name not in self.state.merged_with_fake_quant
- self.state.merged_with_fake_quant[original_node.name] = True
- elif self.fallback_quantization_range:
- min_max_inputs = [
- "fallback_quantization_min_value:0",
- "fallback_quantization_max_value:0"
- ]
- else:
- # Add a RequantizationRange node for finding the min and max values.
- requant_range_node = create_node(
- "RequantizationRange", original_node.name + "_eightbit_requant_range",
- quantized_outputs)
- set_attr_dtype(requant_range_node, "Tinput", dtypes.qint32)
- self.add_output_graph_node(requant_range_node)
- min_max_inputs = [
- requant_range_node.name + ":0", requant_range_node.name + ":1"
- ]
- requantize_node = create_node("Requantize",
- original_node.name + "_eightbit_requantize",
- quantized_outputs + min_max_inputs)
- set_attr_dtype(requantize_node, "Tinput", dtypes.qint32)
- set_attr_dtype(requantize_node, "out_type", dtypes.quint8)
- self.add_output_graph_node(requantize_node)
- return requantize_node.name
-
- def add_dequantize_result_node(self,
- quantized_output_name,
- original_node_name,
- min_tensor_index=1):
- min_max_inputs = [
- "%s:%s" % (quantized_output_name, min_tensor_index),
- "%s:%s" % (quantized_output_name, (min_tensor_index + 1))
- ]
- dequantize_name = original_node_name
- if self.should_merge_with_fake_quant_node():
- fake_quant_node = self.state.output_node_stack[-1][0]
- if original_node_name not in self.state.merged_with_fake_quant:
- min_max_inputs = [fake_quant_node.input[1], fake_quant_node.input[2]]
- self.state.merged_with_fake_quant[original_node_name] = True
- dequantize_name = fake_quant_node.name
-
- dequantize_node = create_node(
- "Dequantize", dequantize_name,
- [quantized_output_name, min_max_inputs[0], min_max_inputs[1]])
- set_attr_dtype(dequantize_node, "T", dtypes.quint8)
- set_attr_string(dequantize_node, "mode", b"MIN_FIRST")
- self.add_output_graph_node(dequantize_node)
-
- def eightbitize_mat_mul_node(self, original_node):
- """Replaces a MatMul node with the eight bit equivalent sub-graph."""
- quantized_mat_mul_name = original_node.name + "_eightbit_quantized_mat_mul"
- all_input_names = self.add_eightbit_prologue_nodes(original_node)
- quantized_mat_mul_node = create_node("QuantizedMatMul",
- quantized_mat_mul_name,
- all_input_names)
- set_attr_dtype(quantized_mat_mul_node, "T1", dtypes.quint8)
- set_attr_dtype(quantized_mat_mul_node, "T2", dtypes.quint8)
- set_attr_dtype(quantized_mat_mul_node, "Toutput", dtypes.qint32)
- copy_attr(quantized_mat_mul_node, "transpose_a",
- original_node.attr["transpose_a"])
- copy_attr(quantized_mat_mul_node, "transpose_b",
- original_node.attr["transpose_b"])
- self.add_output_graph_node(quantized_mat_mul_node)
- quantize_down_name = self.add_quantize_down_nodes(original_node,
- quantized_mat_mul_name)
- self.add_dequantize_result_node(quantize_down_name, original_node.name)
-
- def eightbitize_conv_node(self, original_node):
- """Replaces a Conv2D node with the eight bit equivalent sub-graph."""
- all_input_names = self.add_eightbit_prologue_nodes(original_node)
- quantized_conv_name = original_node.name + "_eightbit_quantized_conv"
- quantized_conv_node = create_node("QuantizedConv2D", quantized_conv_name,
- all_input_names)
- copy_attr(quantized_conv_node, "strides", original_node.attr["strides"])
- copy_attr(quantized_conv_node, "padding", original_node.attr["padding"])
- set_attr_dtype(quantized_conv_node, "Tinput", dtypes.quint8)
- set_attr_dtype(quantized_conv_node, "Tfilter", dtypes.quint8)
- set_attr_dtype(quantized_conv_node, "out_type", dtypes.qint32)
- self.add_output_graph_node(quantized_conv_node)
- quantize_down_name = self.add_quantize_down_nodes(original_node,
- quantized_conv_name)
- self.add_dequantize_result_node(quantize_down_name, original_node.name)
-
- def eightbitize_bias_add_node(self, original_node):
- """Replaces a BiasAdd node with the eight bit equivalent sub-graph."""
- quantized_bias_add_name = (
- original_node.name + "_eightbit_quantized_bias_add")
- all_input_names = self.add_eightbit_prologue_nodes(original_node)
- quantized_bias_add_node = create_node("QuantizedBiasAdd",
- quantized_bias_add_name,
- all_input_names)
- set_attr_dtype(quantized_bias_add_node, "T1", dtypes.quint8)
- set_attr_dtype(quantized_bias_add_node, "T2", dtypes.quint8)
- set_attr_dtype(quantized_bias_add_node, "out_type", dtypes.qint32)
- self.add_output_graph_node(quantized_bias_add_node)
- quantize_down_name = self.add_quantize_down_nodes(original_node,
- quantized_bias_add_name)
- self.add_dequantize_result_node(quantize_down_name, original_node.name)
-
- def eightbitize_single_input_tensor_node(self, original_node,
- add_op_function):
- """Replaces a single-tensor node with the eight bit equivalent sub-graph.
-
- Converts a node like this:
-
- Shape(f) Input(f)
- | |
- +--------v v
- Operation
- |
- v
- (f)
-
- Into a quantized equivalent:
-
- Input(f) ReshapeDims
- +------v v-------------+
- | Reshape
- | |
- | | ReductionDims
- | +-----+ |
- | | +---c---------+
- | v v v v-------+
- | Min Max
- | +----+ |
- v v v--------+
- Quantize
- |
- v
- QuantizedOperation
- | | |
- v v v
- Dequantize
- |
- v
- (f)
-
-
- Args:
- original_node: Float node to be converted.
- add_op_function: Function to create the actual node.
-
- Returns:
- Subgraph representing the quantized version of the original node.
-
- """
- quantized_op_name = original_node.name + "_eightbit_quantized"
- quantized_op_type = "Quantized" + original_node.op
- all_input_names = self.add_eightbit_prologue_nodes(original_node)
- quantized_op_node = create_node(quantized_op_type, quantized_op_name,
- all_input_names)
- add_op_function(original_node, quantized_op_node)
- self.add_output_graph_node(quantized_op_node)
- self.add_dequantize_result_node(quantized_op_name, original_node.name)
-
- def add_pool_function(self, original_node, quantized_op_node):
- set_attr_dtype(quantized_op_node, "T", dtypes.quint8)
- copy_attr(quantized_op_node, "ksize", original_node.attr["ksize"])
- copy_attr(quantized_op_node, "strides", original_node.attr["strides"])
- copy_attr(quantized_op_node, "padding", original_node.attr["padding"])
-
- def add_relu_function(self, unused_arg_node, quantized_op_node):
- set_attr_dtype(quantized_op_node, "Tinput", dtypes.quint8)
-
- def eightbitize_concat_node(self, original_node):
- """Replaces a Concat node with the eight bit equivalent sub-graph.
-
- Converts a node like this:
-
- Shape(f) Input0(f) Input1(f)
- | | |
- +--------v v v----------+
- Concat
- |
- v
- (f)
-
- Into a quantized equivalent:
-
- Shape(f) Input0(f) ReshapeDims Input1(f)
- | +------v v--------------+------------------v v------+
- | | Reshape Reshape |
- | | | | |
- | | | ReductionDims | |
- | | +------+ | +--------+ |
- | | | +---c---------+-----------c-----+ | |
- | | +v v v v-------+---------v v v v+ |
- | | Min Max Min Max |
- | | +----+ | | +-----+ |
- | v v v--------+ +----------v v v
- | Quantize Quantize
- | +------------------+ +----------------------+
- +-------------------------------+ | |
- v v v
- QuantizedConcat
- | | |
- v v v
- Dequantize
- |
- v
- (f)
- Args:
- original_node: Float node to be converted.
-
- Returns:
- Subgraph representing the quantized version of the original node.
-
- """
- namespace_prefix = original_node.name + "_eightbit"
- quantized_concat_name = namespace_prefix + "_quantized_concat"
- reshape_dims_name, reduction_dims_name = self.add_common_quantization_nodes(
- namespace_prefix)
- shape_input_name = original_node.input[0]
- original_inputs = original_node.input[1:]
- input_names = []
- min_names = []
- max_names = []
- for original_input_name in original_inputs:
- quantize_input_name, min_input_name, max_input_name = (
- self.eightbitize_input_to_node(namespace_prefix, original_input_name,
- reshape_dims_name,
- reduction_dims_name))
- input_names.append(quantize_input_name)
- min_names.append(min_input_name)
- max_names.append(max_input_name)
- all_input_names = [shape_input_name]
- all_input_names.extend(input_names)
- all_input_names.extend(min_names)
- all_input_names.extend(max_names)
- quantized_concat_node = create_node("QuantizedConcat",
- quantized_concat_name, all_input_names)
- set_attr_int(quantized_concat_node, "N", len(original_inputs))
- set_attr_dtype(quantized_concat_node, "T", dtypes.quint8)
- self.add_output_graph_node(quantized_concat_node)
- self.add_dequantize_result_node(quantized_concat_name, original_node.name)
-
- def eightbitize_placeholder_node(self, current_node):
- """Replaces a placeholder node with a quint8 placeholder node+dequantize."""
- name = current_node.name
-
- # Convert the placeholder into a quantized type.
- output_node = node_def_pb2.NodeDef()
- output_node.CopyFrom(current_node)
- set_attr_dtype(output_node, "dtype", dtypes.quint8)
- output_node.name += "_original_input"
- self.add_output_graph_node(output_node)
-
- # Add a dequantize to convert back to float.
- dequantize_node = create_node("Dequantize", name, [
- output_node.name, "quantized_input_min_value",
- "quantized_input_max_value"
- ])
- set_attr_dtype(dequantize_node, "T", dtypes.quint8)
- set_attr_string(dequantize_node, "mode", b"MIN_FIRST")
- self.add_output_graph_node(dequantize_node)
-
- # For the descent over the graph to work, the dequantize node must be named
- # current_node.name. However, for the feeding of the graph to work, the
- # placeholder must have the name current_node.name; so record a final set
- # of renames to apply after all processing has been done.
- self.final_node_renames[output_node.name] = name
- self.final_node_renames[dequantize_node.name] = name + "_dequantize"
-
- def eightbitize_reshape_node(self, original_node):
- """Replaces a Reshape node with the eight bit equivalent sub-graph.
-
- Args:
- original_node: Float node to be converted.
-
- Returns:
- Subgraph representing the quantized version of the original node.
-
- """
- namespace_prefix = original_node.name + "_eightbit"
- quantized_reshape_name = namespace_prefix + "_quantized_reshape"
- reshape_dims_name, reduction_dims_name = self.add_common_quantization_nodes(
- namespace_prefix)
- shape_input_name = original_node.input[1]
- quantize_input_name, min_input_name, max_input_name = (
- self.eightbitize_input_to_node(namespace_prefix, original_node.input[0],
- reshape_dims_name, reduction_dims_name))
- quantized_reshape_node = create_node(
- "QuantizedReshape", quantized_reshape_name,
- [quantize_input_name, shape_input_name, min_input_name, max_input_name])
- set_attr_dtype(quantized_reshape_node, "T", dtypes.quint8)
- self.add_output_graph_node(quantized_reshape_node)
- self.add_dequantize_result_node(quantized_reshape_name, original_node.name)
-
- def eightbitize_batch_norm_node(self, original_node):
- """Replaces a MatMul node with the eight bit equivalent sub-graph."""
- namespace_prefix = original_node.name + "_eightbit"
- original_input_name = original_node.input[0]
- original_mean_name = original_node.input[1]
- original_variance_name = original_node.input[2]
- original_beta_name = original_node.input[3]
- original_gamma_name = original_node.input[4]
- quantized_batch_norm_name = namespace_prefix + "_quantized_batch_norm"
-
- reshape_dims_name, reduction_dims_name = self.add_common_quantization_nodes(
- namespace_prefix)
- quantize_input_name, min_input_name, max_input_name = (
- self.eightbitize_input_to_node(namespace_prefix, original_input_name,
- reshape_dims_name, reduction_dims_name))
- quantize_mean_name, min_mean_name, max_mean_name = (
- self.eightbitize_input_to_node(namespace_prefix, original_mean_name,
- reshape_dims_name, reduction_dims_name))
- quantize_variance_name, min_variance_name, max_variance_name = (
- self.eightbitize_input_to_node(namespace_prefix, original_variance_name,
- reshape_dims_name, reduction_dims_name))
- quantize_beta_name, min_beta_name, max_beta_name = (
- self.eightbitize_input_to_node(namespace_prefix, original_beta_name,
- reshape_dims_name, reduction_dims_name))
- quantize_gamma_name, min_gamma_name, max_gamma_name = (
- self.eightbitize_input_to_node(namespace_prefix, original_gamma_name,
- reshape_dims_name, reduction_dims_name))
- quantized_batch_norm_node = create_node(
- "QuantizedBatchNormWithGlobalNormalization", quantized_batch_norm_name,
- [
- quantize_input_name, min_input_name, max_input_name,
- quantize_mean_name, min_mean_name, max_mean_name,
- quantize_variance_name, min_variance_name, max_variance_name,
- quantize_beta_name, min_beta_name, max_beta_name,
- quantize_gamma_name, min_gamma_name, max_gamma_name
- ])
- set_attr_dtype(quantized_batch_norm_node, "Tinput", dtypes.quint8)
- set_attr_dtype(quantized_batch_norm_node, "out_type", dtypes.qint32)
- copy_attr(quantized_batch_norm_node, "scale_after_normalization",
- original_node.attr["scale_after_normalization"])
- copy_attr(quantized_batch_norm_node, "variance_epsilon",
- original_node.attr["variance_epsilon"])
- self.add_output_graph_node(quantized_batch_norm_node)
- quantize_down_name = self.add_quantize_down_nodes(original_node,
- quantized_batch_norm_name)
- self.add_dequantize_result_node(quantize_down_name, original_node.name)
-
- def add_output_graph_node(self, output_node):
- """Inserts one node into the new graph."""
- self.output_graph.node.extend([output_node])
-
- def remove_redundant_quantization(self, old_graph):
- """Removes unneeded pairs of quantize/dequantize ops from the graph.
-
- This is a bit of a tricky function, because it's attempting to spot the
- pattern of dequantizing from eight-bit up to float, and then immediately
- quantizing back down to eight bits again, that's introduced by previous
- passes that do 'key-hole' conversions of individual nodes but have to
- convert back to float to match the previous output interface, since they
- don't know that the next op can handle quantized tensors.
- It works by:
- - Looking for Quantize nodes.
- - Checking to see if their first input is a Dequantize node.
- - Seeing if their min/max inputs come from Min/Max nodes.
- - Making sure those Min/Max nodes are being fed from the same Dequantize.
- - Or that the Min is indirectly being fed from the same Dequantize as Max.
- - Making sure the Dequantize is going through a Reshape (which we add
- during the previous pass when we create the quantize sub-graph).
- - Looking for the dims Const op for the Min/Max dims.
- If all of these conditions are met, then it's a sub-graph pattern that
- we know how to optimize out (and is likely the common one we've introduced).
- We then rewire the graph to skip it entirely, and then rely on the dead node
- removal pass to get rid of any nodes that are no longer needed.
-
- Args:
- old_graph: The model we'll be stripping redundant nodes from.
-
- Returns:
- A graph with the unnecessary nodes removed.
-
- Raises:
- ValueError: Two nodes with the same name were found in the graph.
- """
- old_nodes_map = self.create_nodes_map(old_graph)
- self.output_graph = graph_pb2.GraphDef()
- inputs_to_rename = {}
- # We go through all the nodes, looking for any that match the patterns we
- # know how to optimize away.
- for node in old_graph.node:
- # We always start with a Quantize node, and examine its inputs to see if
- # they are in a form that can be removed.
- if node.op not in ["Quantize", "QuantizeV2"]:
- continue
- dequantize_node_name = node_name_from_input(node.input[0])
- if dequantize_node_name not in old_nodes_map:
- raise ValueError("Input node name '" + dequantize_node_name +
- "' not found in node '" + node.name + "'")
- dequantize_node = old_nodes_map[dequantize_node_name]
- # Do we have a Dequantize feeding in, with the same type as the Quantize?
- if dequantize_node.op != "Dequantize":
- continue
- if node.attr["T"] != dequantize_node.attr["T"]:
- continue
- # Now look at the other inputs, and ensure they're Min/Max nodes.
- min_node_name = node_name_from_input(node.input[1])
- max_node_name = node_name_from_input(node.input[2])
- min_node = old_nodes_map[min_node_name]
- max_node = old_nodes_map[max_node_name]
- is_min_right_type = (min_node.op in ["Min", "Dequantize"])
- is_max_right_type = (max_node.op in ["Max", "Dequantize"])
- if not is_min_right_type or not is_max_right_type:
- print("Didn't find expected types on inputs : %s, %s." % (min_node.op,
- max_node.op))
- continue
- min_node_input_name = node_name_from_input(min_node.input[0])
- max_node_input_name = node_name_from_input(max_node.input[0])
- # There are two different patterns for Min nodes we can recognize, one
- # where the input comes directly from the same one as the Max, and
- # another where we run it through another Min first, so check for both.
- is_same_input = False
- if min_node_input_name == max_node_input_name:
- is_same_input = True
- else:
- first_min_node_input = old_nodes_map[min_node_input_name]
- if first_min_node_input.op == "Concat":
- second_min_node_name = node_name_from_input(
- first_min_node_input.input[1])
- second_min_node = old_nodes_map[second_min_node_name]
- if second_min_node.op == "Min":
- second_min_node_input_name = node_name_from_input(
- second_min_node.input[0])
- is_same_input = (second_min_node_input_name == max_node_input_name)
- if not is_same_input:
- print("Different min/max inputs: " + min_node_input_name)
- continue
- # We recognize this pattern, so mark the graph edges to be rewired to
- # route around it entirely, since we know it's a no-op.
- dequantize_source_name = node_name_from_input(dequantize_node.input[0])
- node_tensor_name = ensure_tensor_name_has_port(node.name)
- min_tensor_name = node.name + ":1"
- max_tensor_name = node.name + ":2"
- inputs_to_rename[node_tensor_name] = dequantize_source_name
- inputs_to_rename[min_tensor_name] = dequantize_node.input[1]
- inputs_to_rename[max_tensor_name] = dequantize_node.input[2]
- # Finally we apply all the rewiring we've marked to the graph.
- for node in old_graph.node:
- for index, input_full_name in enumerate(node.input):
- input_name = ensure_tensor_name_has_port(input_full_name)
- if input_name in inputs_to_rename:
- node.input[index] = inputs_to_rename[input_name]
- self.add_output_graph_node(node)
- return self.output_graph
-
- def apply_final_node_renames(self):
- """Applies node renames in self.final_node_renames to self.output_graph."""
- old_graph = self.output_graph
- self.output_graph = graph_pb2.GraphDef()
- for node in old_graph.node:
- node.name = self.final_node_renames.get(node.name, node.name)
- for index, input_name in enumerate(node.input):
- node_name = node_name_from_input(input_name)
- input_full_name = ensure_tensor_name_has_port(input_name)
- if node_name in self.final_node_renames:
- node.input[index] = "%s%s" % (self.final_node_renames[node_name],
- input_full_name[len(node_name):])
- self.add_output_graph_node(node)
- return self.output_graph
-
- def remove_dead_nodes(self, output_names):
- """Removes nodes that are no longer needed for inference from the graph."""
- old_output_graph = self.output_graph
- self.output_graph = graph_util.extract_sub_graph(old_output_graph,
- output_names)
-
- def quantize_weights(self, input_graph, quantization_mode):
- """Quantize float Const ops.
-
- There are two modes of operations, both replace float Const ops with
- quantized values.
- 1. If quantization_mode is "weights_rounded", this function replaces float
- Const ops with quantized float Const ops - same as the original op, but
- float values being mapped to the center of one of 1<<FLAGS.bitdepth buckets.
- This does not change the raw model size, but compression algorithms such as
- zip (as used for compressing apks) or bzip2 will achieve a very good
- compression ratio.
- 2. For other quantization modes ("MIN_COMBINED" or "MIN_FIRST"), float
- Const ops are quantized and replaced by a tuple of four ops to perform
- the dequantization at runtime:
- * eight-bit Const (bucket indices, same shape as original float Const op
- * two float Const ops (min and max value of original float Const op)
- * Dequantize op to convert the eight-bit consts to float tensors.
- The quantization mode is important because we see accuracy problems when
- quantizing weights for different situations depending on the algorithm
- used. We haven't figured out exactly what the underlying cause is yet,
- unfortunately.
-
- Args:
- input_graph: A GraphDef of the model containing float Const ops.
- quantization_mode: How to quantize and dequantize the values.
-
- Returns:
- A GraphDef of the converted graph.
-
- Raises:
- ValueError: If quantization_mode is unsupported.
- """
- output_graph = graph_pb2.GraphDef()
- for input_node in input_graph.node:
- should_quantize = False
- if input_node.op == "Const":
- dtype = dtypes.as_dtype(input_node.attr["dtype"].type)
- if dtype == dtypes.float32:
- should_quantize = True
- if should_quantize:
- if quantization_mode == "weights_rounded":
- output_graph.node.extend(quantize_weight_rounded(input_node))
- elif quantization_mode in (b"MIN_COMBINED", b"MIN_FIRST"):
- output_graph.node.extend(
- quantize_weight_eightbit(input_node, quantization_mode))
- else:
- raise ValueError("Unsupported quantization mode %s." %
- quantization_mode)
- else:
- output_node = node_def_pb2.NodeDef()
- output_node.CopyFrom(input_node)
- output_graph.node.extend([output_node])
- return output_graph
-
- def set_input_graph(self, new_input_graph):
- self.input_graph = new_input_graph
- self.nodes_map = self.create_nodes_map(self.input_graph)
-
-
-def main(unused_args):
- if not gfile.Exists(FLAGS.input):
- print("Input graph file '" + FLAGS.input + "' does not exist!")
- return -1
-
- known_modes = [
- "round", "quantize", "eightbit", "weights", "test", "weights_rounded"
- ]
- if not any(FLAGS.mode in s for s in known_modes):
- print("mode is '" + FLAGS.mode + "', not in " + ", ".join(known_modes) +
- ".")
- return -1
-
- tf_graph = graph_pb2.GraphDef()
- with gfile.Open(FLAGS.input, "rb") as f:
- data = f.read()
- tf_graph.ParseFromString(data)
-
- graph = ops.Graph()
- with graph.as_default():
- importer.import_graph_def(tf_graph, input_map={}, name="")
-
- quantized_input_range = None
- if FLAGS.quantized_input:
- quantized_input_range = [
- FLAGS.quantized_input_min, FLAGS.quantized_input_max
- ]
-
- fallback_quantization_range = None
- if (FLAGS.quantized_fallback_min is not None or
- FLAGS.quantized_fallback_max is not None):
- assert FLAGS.quantized_fallback_min is not None
- assert FLAGS.quantized_fallback_max is not None
- fallback_quantization_range = [
- FLAGS.quantized_fallback_min, FLAGS.quantized_fallback_max
- ]
-
- rewriter = GraphRewriter(tf_graph, FLAGS.mode, quantized_input_range,
- fallback_quantization_range)
-
- output_graph = rewriter.rewrite(FLAGS.output_node_names.split(","))
-
- f = gfile.FastGFile(FLAGS.output, "wb")
- f.write(output_graph.SerializeToString())
-
- return 0
-
-
-if __name__ == "__main__":
- app.run()
diff --git a/tensorflow/tools/quantization/quantize_graph_test.py b/tensorflow/tools/quantization/quantize_graph_test.py
deleted file mode 100644
index 92bb512..0000000
--- a/tensorflow/tools/quantization/quantize_graph_test.py
+++ /dev/null
@@ -1,966 +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.
-# ==============================================================================
-"""Tests the graph quantization script.
-
-"""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import sys
-import numpy as np
-
-from tensorflow.core.framework import graph_pb2
-from tensorflow.python.client import session
-from tensorflow.python.framework import dtypes
-from tensorflow.python.framework import graph_util
-from tensorflow.python.framework import importer
-from tensorflow.python.framework import ops as ops_lib
-from tensorflow.python.platform import flags as flags_lib
-from tensorflow.python.platform import test
-from tensorflow.python.platform import tf_logging
-from tensorflow.tools.quantization import quantize_graph
-
-flags = flags_lib
-FLAGS = flags.FLAGS
-
-
-def run_graph_def(graph_def, input_map, outputs):
- graph = ops_lib.Graph()
- with graph.as_default():
- importer.import_graph_def(graph_def, input_map={}, name="")
- with session.Session(graph=graph) as sess:
- results = sess.run(outputs, feed_dict=input_map)
- return results
-
-
-def test_mat_mul(m, n, k, a, b):
- """Tests a MatMul replacement."""
- a_constant_name = "a_constant"
- b_constant_name = "b_constant"
- mat_mul_name = "mat_mul"
-
- float_graph_def = graph_pb2.GraphDef()
- a_constant = quantize_graph.create_constant_node(
- a_constant_name, value=a, dtype=dtypes.float32, shape=[m, k])
- float_graph_def.node.extend([a_constant])
- b_constant = quantize_graph.create_constant_node(
- b_constant_name, value=b, dtype=dtypes.float32, shape=[k, n])
- float_graph_def.node.extend([b_constant])
- mat_mul_node = quantize_graph.create_node("MatMul", mat_mul_name,
- [a_constant_name, b_constant_name])
- quantize_graph.set_attr_dtype(mat_mul_node, "T", dtypes.float32)
- quantize_graph.set_attr_bool(mat_mul_node, "transpose_a", False)
- quantize_graph.set_attr_bool(mat_mul_node, "transpose_b", False)
- float_graph_def.node.extend([mat_mul_node])
-
- test_graph(float_graph_def, {}, [mat_mul_name])
-
-
-def test_conv(depth, image_width, image_height, image_batch_count, filter_size,
- filter_count, stride, padding, input_values, filter_values):
- """Tests a Conv replacement."""
- input_constant_name = "input_constant"
- filter_constant_name = "filter_constant"
- conv_name = "conv"
-
- float_graph_def = graph_pb2.GraphDef()
- input_constant = quantize_graph.create_constant_node(
- input_constant_name,
- value=input_values,
- dtype=dtypes.float32,
- shape=[image_batch_count, image_height, image_width, depth])
- float_graph_def.node.extend([input_constant])
- filter_constant = quantize_graph.create_constant_node(
- filter_constant_name,
- value=filter_values,
- dtype=dtypes.float32,
- shape=[filter_size, filter_size, depth, filter_count])
- float_graph_def.node.extend([filter_constant])
- conv_node = quantize_graph.create_node(
- "Conv2D", conv_name, [input_constant_name, filter_constant_name])
- quantize_graph.set_attr_dtype(conv_node, "T", dtypes.float32)
- quantize_graph.set_attr_int_list(conv_node, "strides", [1, stride, stride, 1])
- quantize_graph.set_attr_string(conv_node, "padding", padding)
- float_graph_def.node.extend([conv_node])
-
- test_graph(float_graph_def, {}, [conv_name])
-
-
-def are_tensors_near(a, b, tolerance):
- """Tests whether two tensors are nearly identical.
-
- This is a specialized comparison function designed to help debug problems with
- quantization. It prints out information about the differences between tensors
- on failure, paying special attention to possible biases by looking at the mean
- and absolute average errors.
-
- Args:
- a: First comparison tensor.
- b: Second comparison tensor.
- tolerance: Float value indicating how large an error between values is ok.
-
- Returns:
- Boolean indicating whether the two inputs were close enough.
- """
- flat_a = a.flatten()
- flat_b = b.flatten()
- if len(flat_a) != len(flat_b):
- tf_logging.info("Tensors are different sizes: " + str(len(flat_a)) + " vs "
- + str(len(flat_b)))
- return False
- value_count = len(flat_a)
- how_many_different = 0
- total_difference = 0
- total_abs_difference = 0
- for index in range(value_count):
- a_value = flat_a[index]
- b_value = flat_b[index]
- difference = a_value - b_value
- total_difference += difference
- total_abs_difference += abs(difference)
- if abs(difference) > tolerance:
- how_many_different += 1
- mean_difference = total_difference / value_count
- mean_abs_difference = total_abs_difference / value_count
- proportion_different = (how_many_different * 1.0) / value_count
- if how_many_different == 0:
- return True
- else:
- tf_logging.info("Tensors have {0} different values ({1}%), with mean"
- " difference {2} and mean absolute difference {3}".format(
- how_many_different, proportion_different * 100,
- mean_difference, mean_abs_difference))
- return False
-
-
-def get_top_value(input_values):
- max_value = None
- max_index = None
- for index, value in enumerate(input_values.flatten()):
- if max_value is None or value > max:
- max_value = value
- max_index = index
- return max_index, max_value
-
-
-def test_graph(float_graph_def, input_map, output_names, log_graph=False):
- """Runs the float graph through the rewriter and tests the results."""
- float_results = run_graph_def(
- float_graph_def, input_map,
- [output_name + ":0" for output_name in output_names])
- # TODO(petewarden): round test is currently failing because there is no
- # RoundToSteps op available.
- # round_rewriter = quantize_graph.GraphRewriter(float_graph_def, "round")
- # round_graph_def = round_rewriter.rewrite(output_name)
- # round_results = run_graph_def(round_graph_def, input_map,
- # [output_name + ":0"])
- # assert are_tensors_near(expected, round_results[0], 1.0)
- #
- # TODO(petewarden): Add test for "quantize" mode.
-
- eightbit_rewriter = quantize_graph.GraphRewriter(
- float_graph_def, "eightbit", quantized_input_range=None)
- eightbit_graph_def = eightbit_rewriter.rewrite(output_names)
- eightbit_results = run_graph_def(
- eightbit_graph_def, input_map,
- [output_name + ":0" for output_name in output_names])
- for expected, result in zip(float_results, eightbit_results):
- assert are_tensors_near(expected, result, 1.0)
-
- if log_graph:
- tf_logging.info("8bit:\n%s", str(eightbit_graph_def))
-
- # Test the weights_rounded mode. This uses the default bit_depth.
- weights_rounded_rewriter = quantize_graph.GraphRewriter(
- float_graph_def, "weights_rounded", quantized_input_range=None)
- weights_rounded_graph_def = weights_rounded_rewriter.rewrite(output_names)
- weights_rounded_results = run_graph_def(
- weights_rounded_graph_def, input_map,
- [output_name + ":0" for output_name in output_names])
- for expected, result in zip(float_results, weights_rounded_results):
- assert are_tensors_near(expected, result, 1.0)
-
-
-class QuantizeGraphTest(test.TestCase):
-
- def test_negative_const_problem(self):
- shape_constant_name = "shape_constant"
- shape_constant = quantize_graph.create_constant_node(
- shape_constant_name, value=-0.8, dtype=dtypes.float32, shape=[1])
- quantization_result = quantize_graph.quantize_weight_eightbit(
- shape_constant, b"MIN_COMBINED")
- self.assertEqual(4, len(quantization_result))
-
- def test_odd_padding_problem(self):
- """Tests one error case we ran into in a real graph."""
- test_conv(1, 4, 4, 1, 3, 1, 2, b"SAME",
- [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],
- [1, 2, 3, 4, 5, 6, 7, 8, 9])
-
- def test_mat_mul_tiny(self):
- # These tests are added to test the generate case where
- # min(matrix) == max(matrix), which used to cause problems.
- test_mat_mul(1, 1, 1, [2], [3])
- test_mat_mul(1, 2, 1, [1], [2, 3])
- test_mat_mul(1, 1, 2, [1, 1], [1, 1])
- test_mat_mul(1, 1, 2, [0, 0], [1, 1])
- # The general case.
- test_mat_mul(1, 1, 2, [1, 2], [1, 2])
-
- def test_mat_mul_small(self):
- test_mat_mul(2, 4, 3, [1, 2, 3, 4, 5, 6],
- [7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18])
-
- def test_conv(self):
- test_conv(1, 4, 3, 1, 3, 1, 1, b"SAME",
- [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
- [1, 4, 7, 2, 5, 8, 3, 6, 9])
-
- def test_reshape(self):
- """Tests that MatMul->Reshape->MatMul avoids extra quantize/dequantize."""
-
- def make_matmul(name, a, b):
- n = quantize_graph.create_node("MatMul", name, [a.name, b.name])
- quantize_graph.set_attr_dtype(n, "T", dtypes.float32)
- quantize_graph.set_attr_bool(n, "transpose_a", False)
- quantize_graph.set_attr_bool(n, "transpose_b", False)
- return n
-
- # matmul_1 = input*weight_1
- input_node = quantize_graph.create_constant_node(
- "input", value=[0, 1, 2, 3], dtype=dtypes.float32, shape=[4, 1])
- weight_1_node = quantize_graph.create_constant_node(
- "weight_1",
- value=[.5, .6, .7, .8, .9],
- dtype=dtypes.float32,
- shape=[1, 5])
- matmul_1_node = make_matmul("matmul_1", input_node, weight_1_node)
-
- # Reshape 4x5 to 10x2.
- new_shape_node = quantize_graph.create_constant_node(
- "new_shape_node", value=[10, 2], dtype=dtypes.int32, shape=[2])
- reshape_node = quantize_graph.create_node(
- "Reshape", "reshape", [matmul_1_node.name, new_shape_node.name])
- quantize_graph.set_attr_dtype(reshape_node, "T", dtypes.float32)
-
- # matmul_2_node = reshape*weight_2
- weight_2_node = quantize_graph.create_constant_node(
- "weight_2", value=[1.5, 2.5], dtype=dtypes.float32, shape=[2, 1])
- matmul_2_node = make_matmul("matmul_2", reshape_node, weight_2_node)
-
- g = graph_pb2.GraphDef()
- g.node.extend([
- input_node, weight_1_node, matmul_1_node, new_shape_node, reshape_node,
- weight_2_node, matmul_2_node
- ])
-
- # Test the graph
- test_graph(g, {}, ["matmul_2"])
-
- # Verify there is only one Quantize and one Requantize op.
- eightbit_rewriter = quantize_graph.GraphRewriter(
- g, "eightbit", quantized_input_range=None)
- eightbit_graph_def = eightbit_rewriter.rewrite(["matmul_2"])
-
- ops = [node.op for node in eightbit_graph_def.node]
- # No quantize since all inputs are const and can be quantized up-front.
- self.assertEqual(0, ops.count("QuantizeV2") + ops.count("Quantize"))
- self.assertEqual(1, ops.count("QuantizedReshape"))
-
- # One dequantize at the end.
- self.assertEqual(1, ops.count("Dequantize"))
-
- def test_quantize_array(self):
- # Test invalid parameters (empty array, or 0 buckets.
- self.assertRaises(ValueError, quantize_graph.quantize_array, np.array([]),
- 2)
- self.assertRaises(ValueError, quantize_graph.quantize_array,
- np.array([1, 2]), 0)
- # Test input array of length 1.
- arr = np.array([1])
- qarr = quantize_graph.quantize_array(arr, 1)
- self.assertEqual(arr, qarr)
- qarr = quantize_graph.quantize_array(arr, 2)
- self.assertEqual(arr, qarr)
- # Test input array with all elements equal.
- arr = np.array([1, 1, 1])
- qarr = quantize_graph.quantize_array(arr, 10)
- self.assertTrue((np.array([1, 1, 1]) == qarr).all())
- # Test "normal" input arrays.
- arr = np.array([0, 0.3, 0.6, 1])
- qarr = quantize_graph.quantize_array(arr, 1)
- self.assertTrue((np.array([0.5, 0.5, 0.5, 0.5]) == qarr).all())
- qarr = quantize_graph.quantize_array(arr, 2)
- self.assertTrue((np.array([0.25, 0.25, 0.75, 0.75]) == qarr).all())
- qarr = quantize_graph.quantize_array(arr.reshape((2, 2)), 2)
- self.assertTrue((np.array([[0.25, 0.25], [0.75, 0.75]]) == qarr).all())
-
- def test_non_float_concat(self):
- concat_dim = quantize_graph.create_constant_node(
- "concat_dim", value=0, dtype=dtypes.int32, shape=[])
- a = quantize_graph.create_constant_node(
- "a",
- value=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
- dtype=dtypes.int32,
- shape=[2, 2, 3])
- b = quantize_graph.create_constant_node(
- "b",
- value=[13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24],
- dtype=dtypes.int32,
- shape=[2, 2, 3])
- concat = quantize_graph.create_node("Concat", "concat",
- [concat_dim.name, a.name, b.name])
- quantize_graph.set_attr_int(concat, "N", 2)
- quantize_graph.set_attr_dtype(concat, "T", dtypes.int32)
-
- g = graph_pb2.GraphDef()
- g.node.extend([concat_dim, a, b, concat])
- test_graph(g, {}, [concat.name])
-
- def test_non_float_reshape(self):
- a = quantize_graph.create_constant_node(
- "a",
- value=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
- dtype=dtypes.int32,
- shape=[2, 2, 3])
- shape = quantize_graph.create_constant_node(
- "shape", value=[12], dtype=dtypes.int32, shape=[1])
- reshape = quantize_graph.create_node("Reshape", "reshape",
- [a.name, shape.name])
- quantize_graph.set_attr_dtype(reshape, "T", dtypes.int32)
-
- g = graph_pb2.GraphDef()
- g.node.extend([a, shape, reshape])
- test_graph(g, {}, [reshape.name])
-
- def test_concat(self):
- shape_constant_name = "shape_constant"
- a_constant_name = "a_constant"
- b_constant_name = "b_constant"
- concat_name = "concat"
-
- float_graph_def = graph_pb2.GraphDef()
- shape_constant = quantize_graph.create_constant_node(
- shape_constant_name, value=0, dtype=dtypes.int32, shape=[])
- float_graph_def.node.extend([shape_constant])
- a_constant = quantize_graph.create_constant_node(
- a_constant_name,
- value=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
- dtype=dtypes.float32,
- shape=[2, 2, 3])
- float_graph_def.node.extend([a_constant])
- b_constant = quantize_graph.create_constant_node(
- b_constant_name,
- value=[13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24],
- dtype=dtypes.float32,
- shape=[2, 2, 3])
- float_graph_def.node.extend([b_constant])
- concat_node = quantize_graph.create_node(
- "Concat", concat_name,
- [shape_constant_name, a_constant_name, b_constant_name])
- quantize_graph.set_attr_int(concat_node, "N", 2)
- quantize_graph.set_attr_dtype(concat_node, "T", dtypes.float32)
- float_graph_def.node.extend([concat_node])
-
- test_graph(float_graph_def, {}, [concat_name])
-
- # Verify the concat is quantized.
- eightbit_rewriter = quantize_graph.GraphRewriter(
- float_graph_def, "eightbit", quantized_input_range=None)
- eightbit_graph_def = eightbit_rewriter.rewrite([concat_name])
-
- ops = [node.op for node in eightbit_graph_def.node]
- self.assertEqual(1, ops.count("QuantizedConcat"))
-
- def test_multiple_outputs(self):
- input_constant_name = "input_constant"
- split_constant_name = "split_constant"
- split_name = "split"
- concat_constant_name = "concat_constant"
- concat_name = "concat"
-
- float_graph_def = graph_pb2.GraphDef()
- input_constant = quantize_graph.create_constant_node(
- input_constant_name,
- value=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
- dtype=dtypes.float32,
- shape=[2, 6])
- float_graph_def.node.extend([input_constant])
- split_constant = quantize_graph.create_constant_node(
- split_constant_name, value=1, dtype=dtypes.int32, shape=[])
- float_graph_def.node.extend([split_constant])
- split_node = quantize_graph.create_node(
- "Split", split_name, [split_constant_name, input_constant_name])
- quantize_graph.set_attr_int(split_node, "num_split", 2)
- quantize_graph.set_attr_dtype(split_node, "T", dtypes.float32)
- float_graph_def.node.extend([split_node])
- concat_constant = quantize_graph.create_constant_node(
- concat_constant_name, value=1, dtype=dtypes.int32, shape=[])
- float_graph_def.node.extend([concat_constant])
- concat_node = quantize_graph.create_node(
- "Concat", concat_name,
- [concat_constant_name, split_name + ":0", split_name + ":1"])
- quantize_graph.set_attr_int(concat_node, "N", 2)
- quantize_graph.set_attr_dtype(concat_node, "T", dtypes.float32)
- float_graph_def.node.extend([concat_node])
-
- test_graph(float_graph_def, {}, [concat_name])
-
- def test_node_name_from_input(self):
- self.assertEqual("SomeName",
- quantize_graph.node_name_from_input("^SomeName:2"))
-
- def test_unique_node_name_from_input(self):
- self.assertEqual("__hat__SomeName__port__2",
- quantize_graph.unique_node_name_from_input("^SomeName:2"))
-
- def test_identity(self):
- input_constant_name = "input_constant"
- identity_name = "identity"
- float_graph_def = graph_pb2.GraphDef()
- input_constant = quantize_graph.create_constant_node(
- input_constant_name,
- value=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
- dtype=dtypes.float32,
- shape=[2, 6])
- float_graph_def.node.extend([input_constant])
- identity_node = quantize_graph.create_node("Identity", identity_name,
- [input_constant_name])
- quantize_graph.set_attr_dtype(identity_node, "T", dtypes.float32)
- float_graph_def.node.extend([identity_node])
-
- mul_name = "mul"
- mul_node = quantize_graph.create_node("Mul", mul_name,
- [identity_name, identity_name])
- quantize_graph.set_attr_dtype(mul_node, "T", dtypes.float32)
- float_graph_def.node.extend([mul_node])
-
- test_graph(float_graph_def, {}, [mul_name])
-
- def test_keep_control_edges(self):
- no_op_name = "no_op"
- a_constant_name = "a_constant"
- b_constant_name = "b_constant"
- a_check_name = "a_check"
- b_check_name = "b_check"
- a_identity_name = "a_identity"
- b_identity_name = "b_identity"
- add_name = "add"
- graph_def = graph_pb2.GraphDef()
- no_op = quantize_graph.create_node("NoOp", no_op_name, [])
- graph_def.node.extend([no_op])
- a_constant = quantize_graph.create_constant_node(
- a_constant_name, value=1, dtype=dtypes.float32, shape=[])
- graph_def.node.extend([a_constant])
- a_check_node = quantize_graph.create_node("CheckNumerics", a_check_name,
- [a_constant_name])
- graph_def.node.extend([a_check_node])
- a_identity_node = quantize_graph.create_node(
- "Identity", a_identity_name,
- [a_constant_name, "^" + a_check_name, "^" + no_op_name])
- graph_def.node.extend([a_identity_node])
- b_constant = quantize_graph.create_constant_node(
- b_constant_name, value=1, dtype=dtypes.float32, shape=[])
- graph_def.node.extend([b_constant])
- b_check_node = quantize_graph.create_node("CheckNumerics", b_check_name,
- [b_constant_name])
- graph_def.node.extend([b_check_node])
- b_identity_node = quantize_graph.create_node(
- "Identity", b_identity_name, [b_constant_name, "^" + b_check_name])
- graph_def.node.extend([b_identity_node])
- add_node = quantize_graph.create_node("Add", add_name,
- [a_identity_name, b_identity_name])
- quantize_graph.set_attr_dtype(add_node, "T", dtypes.float32)
- graph_def.node.extend([add_node])
-
- expected_output = graph_pb2.GraphDef()
- no_op = quantize_graph.create_node("NoOp", no_op_name, [])
- expected_output.node.extend([no_op])
- a_constant = quantize_graph.create_constant_node(
- a_constant_name, value=1, dtype=dtypes.float32, shape=[])
- expected_output.node.extend([a_constant])
- a_identity_node = quantize_graph.create_node(
- "Identity", a_identity_name, [a_constant_name, "^" + no_op_name])
- expected_output.node.extend([a_identity_node])
- b_constant = quantize_graph.create_constant_node(
- b_constant_name, value=1, dtype=dtypes.float32, shape=[])
- expected_output.node.extend([b_constant])
- add_node = quantize_graph.create_node("Add", add_name,
- [a_identity_name, b_constant_name])
- quantize_graph.set_attr_dtype(add_node, "T", dtypes.float32)
- expected_output.node.extend([add_node])
- expected_output.versions.CopyFrom(graph_def.versions)
- expected_output.library.CopyFrom(graph_def.library)
-
- output = graph_util.remove_training_nodes(graph_def)
- stripped_output = graph_util.extract_sub_graph(output, [add_name])
- self.assertProtoEquals(expected_output, stripped_output)
-
- def test_batch_norm(self):
- input_constant_name = "input_constant"
- mean_constant_name = "mean_constant"
- variance_constant_name = "variance_constant"
- beta_constant_name = "beta_constant"
- gamma_constant_name = "gamma_constant"
- batch_norm_name = "batch_norm"
- float_graph_def = graph_pb2.GraphDef()
- input_constant = quantize_graph.create_constant_node(
- input_constant_name,
- value=[1, 4, 2, 5, 3, 6, -1, -4, -2, -5, -3, -6],
- dtype=dtypes.float32,
- shape=[1, 1, 6, 2])
- float_graph_def.node.extend([input_constant])
- mean_constant = quantize_graph.create_constant_node(
- mean_constant_name, value=[10, 20], dtype=dtypes.float32, shape=[2])
- float_graph_def.node.extend([mean_constant])
- variance_constant = quantize_graph.create_constant_node(
- variance_constant_name,
- value=[0.25, 0.5],
- dtype=dtypes.float32,
- shape=[2])
- float_graph_def.node.extend([variance_constant])
- beta_constant = quantize_graph.create_constant_node(
- beta_constant_name, value=[0.1, 0.6], dtype=dtypes.float32, shape=[2])
- float_graph_def.node.extend([beta_constant])
- gamma_constant = quantize_graph.create_constant_node(
- gamma_constant_name, value=[0, 0], dtype=dtypes.float32, shape=[2])
- float_graph_def.node.extend([gamma_constant])
- batch_norm_node = quantize_graph.create_node(
- "BatchNormWithGlobalNormalization", batch_norm_name, [
- input_constant_name, mean_constant_name, variance_constant_name,
- beta_constant_name, gamma_constant_name
- ])
- quantize_graph.set_attr_dtype(batch_norm_node, "T", dtypes.float32)
- quantize_graph.set_attr_bool(batch_norm_node, "scale_after_normalization",
- False)
- quantize_graph.set_attr_float(batch_norm_node, "variance_epsilon", 0.001)
- float_graph_def.node.extend([batch_norm_node])
- test_graph(float_graph_def, {}, [batch_norm_name])
-
- def test_max_pool(self):
- input_constant_name = "input_constant"
- max_pool_name = "max_pool"
- float_graph_def = graph_pb2.GraphDef()
- input_constant = quantize_graph.create_constant_node(
- input_constant_name,
- value=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
- dtype=dtypes.float32,
- shape=[1, 2, 6, 1])
- float_graph_def.node.extend([input_constant])
- max_pool_node = quantize_graph.create_node("MaxPool", max_pool_name,
- [input_constant_name])
- quantize_graph.set_attr_int_list(max_pool_node, "ksize", [1, 2, 2, 1])
- quantize_graph.set_attr_int_list(max_pool_node, "strides", [1, 1, 1, 1])
- quantize_graph.set_attr_string(max_pool_node, "padding", b"SAME")
- float_graph_def.node.extend([max_pool_node])
- test_graph(float_graph_def, {}, [max_pool_name])
-
- def test_avg_pool(self):
- input_constant_name = "input_constant"
- avg_pool_name = "avg_pool"
- float_graph_def = graph_pb2.GraphDef()
- input_constant = quantize_graph.create_constant_node(
- input_constant_name,
- value=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
- dtype=dtypes.float32,
- shape=[1, 2, 6, 1])
- float_graph_def.node.extend([input_constant])
- avg_pool_node = quantize_graph.create_node("AvgPool", avg_pool_name,
- [input_constant_name])
- quantize_graph.set_attr_dtype(avg_pool_node, "T", dtypes.float32)
- quantize_graph.set_attr_int_list(avg_pool_node, "ksize", [1, 2, 2, 1])
- quantize_graph.set_attr_int_list(avg_pool_node, "strides", [1, 1, 1, 1])
- quantize_graph.set_attr_string(avg_pool_node, "padding", b"SAME")
- float_graph_def.node.extend([avg_pool_node])
- test_graph(float_graph_def, {}, [avg_pool_name])
-
- def test_relu(self):
- input_constant_name = "input_constant"
- relu_name = "relu"
- float_graph_def = graph_pb2.GraphDef()
- input_constant = quantize_graph.create_constant_node(
- input_constant_name,
- value=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
- dtype=dtypes.float32,
- shape=[1, 2, 6, 1])
- float_graph_def.node.extend([input_constant])
- relu_node = quantize_graph.create_node("Relu", relu_name,
- [input_constant_name])
- quantize_graph.set_attr_dtype(relu_node, "T", dtypes.float32)
- float_graph_def.node.extend([relu_node])
- test_graph(float_graph_def, {}, [relu_name])
-
- def test_relu_w_fake_quant_w_min_max_vars(self):
- input_node = quantize_graph.create_constant_node(
- "input",
- value=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
- dtype=dtypes.float32,
- shape=[1, 2, 6, 1])
- relu_node = quantize_graph.create_node("Relu", "relu", [input_node.name])
- quantize_graph.set_attr_dtype(relu_node, "T", dtypes.float32)
-
- min_node = quantize_graph.create_constant_node(
- "min_bias_add", value=0, dtype=dtypes.float32, shape=[])
- max_node = quantize_graph.create_constant_node(
- "max_bias_add", value=12, dtype=dtypes.float32, shape=[])
- fake_quant_node = quantize_graph.create_node(
- "FakeQuantWithMinMaxVars", "fake_quant",
- [relu_node.name, min_node.name, max_node.name])
-
- float_graph_def = graph_pb2.GraphDef()
- float_graph_def.node.extend(
- [input_node, relu_node, min_node, max_node, fake_quant_node])
- test_graph(float_graph_def, {}, [fake_quant_node.name], log_graph=True)
-
- # Verify there is only one Quantize and one Requantize op.
- eightbit_rewriter = quantize_graph.GraphRewriter(
- float_graph_def, "eightbit", quantized_input_range=None)
- eightbit_graph_def = eightbit_rewriter.rewrite([fake_quant_node.name])
-
- ops = [node.op for node in eightbit_graph_def.node]
- # No quantize since all inputs are const and can be quantized up-front.
- self.assertEqual(0, ops.count("QuantizeV2") + ops.count("Quantize"))
-
- # One dequantize at the end.
- self.assertEqual(1, ops.count("Dequantize"))
-
- def test_relu6(self):
- input_constant_name = "input_constant"
- relu6_name = "relu6"
- float_graph_def = graph_pb2.GraphDef()
- input_constant = quantize_graph.create_constant_node(
- input_constant_name,
- value=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
- dtype=dtypes.float32,
- shape=[1, 2, 6, 1])
- float_graph_def.node.extend([input_constant])
- relu6_node = quantize_graph.create_node("Relu6", relu6_name,
- [input_constant_name])
- quantize_graph.set_attr_dtype(relu6_node, "T", dtypes.float32)
- float_graph_def.node.extend([relu6_node])
- test_graph(float_graph_def, {}, [relu6_name])
-
- def test_bias_add(self):
- input_constant_name = "input_constant"
- offset_constant_name = "offset_constant"
- bias_add_name = "bias_add"
- float_graph_def = graph_pb2.GraphDef()
- input_constant = quantize_graph.create_constant_node(
- input_constant_name,
- value=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
- dtype=dtypes.float32,
- shape=[1, 1, 2, 6])
- float_graph_def.node.extend([input_constant])
- offset_constant = quantize_graph.create_constant_node(
- offset_constant_name,
- value=[1, 2, 3, 4, 5, 6],
- dtype=dtypes.float32,
- shape=[6])
- float_graph_def.node.extend([offset_constant])
- bias_add_node = quantize_graph.create_node(
- "BiasAdd", bias_add_name, [input_constant_name, offset_constant_name])
- quantize_graph.set_attr_dtype(bias_add_node, "T", dtypes.float32)
- float_graph_def.node.extend([bias_add_node])
- test_graph(float_graph_def, {}, [bias_add_name])
-
- def test_quantized_input_range_errors(self):
- with self.assertRaises(ValueError):
- # Invalid mode.
- quantize_graph.GraphRewriter(graph_pb2.GraphDef(), "weights_rounded",
- [0, 1])
- with self.assertRaises(ValueError):
- # Invalid range.
- quantize_graph.GraphRewriter(graph_pb2.GraphDef(), "eightbit", [0, -1])
-
- def test_quantized_input_range_bias_add(self):
- input_shape = [1, 1, 2, 6]
- input_n = quantize_graph.create_node("Placeholder", "input", [])
- quantize_graph.set_attr_dtype(input_n, "dtype", dtypes.float32)
- quantize_graph.set_attr_shape(input_n, "shape", input_shape)
- offset_n = quantize_graph.create_constant_node(
- "offset", value=[1, 2, 3, 4, 5, 6], dtype=dtypes.float32, shape=[6])
- bias_add_n = quantize_graph.create_node("BiasAdd", "bias_add",
- [input_n.name, offset_n.name])
- quantize_graph.set_attr_dtype(bias_add_n, "T", dtypes.float32)
-
- float_graph_def = graph_pb2.GraphDef()
- float_graph_def.node.extend([input_n, offset_n, bias_add_n])
-
- input_map = {
- input_n.name + ":0":
- np.reshape([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], input_shape)
- }
- self._RunTestsForQuantizedInputRange(float_graph_def, input_map,
- [bias_add_n.name], [-1, 20.])
- self._RunTestsForQuantizedInputRange(float_graph_def, input_map,
- [bias_add_n.name], [0, 12.])
-
- def test_quantized_input_range_mat_mul(self):
- shapes = [[3, 2], [2, 4]]
- inputs = []
- for i, shape in enumerate(shapes):
- node = quantize_graph.create_node("Placeholder", "input_%s" % i, [])
- quantize_graph.set_attr_dtype(node, "dtype", dtypes.float32)
- quantize_graph.set_attr_shape(node, "shape", shape)
- inputs.append(node)
- mat_mul_node = quantize_graph.create_node("MatMul", "mat_mul",
- [n.name for n in inputs])
- quantize_graph.set_attr_dtype(mat_mul_node, "T", dtypes.float32)
-
- float_graph_def = graph_pb2.GraphDef()
- float_graph_def.node.extend(inputs + [mat_mul_node])
-
- input_map = {
- inputs[0].name + ":0":
- np.reshape([1, 2, 3, 4, 5, 6], shapes[0]),
- inputs[1].name + ":0":
- np.reshape([.8, .7, .6, .5, .4, .3, .2, .1], shapes[1])
- }
- self._RunTestsForQuantizedInputRange(float_graph_def, input_map,
- [mat_mul_node.name], [-1, 20.])
- self._RunTestsForQuantizedInputRange(float_graph_def, input_map,
- [mat_mul_node.name], [0, 6.])
-
- def _RunTestsForQuantizedInputRange(self, float_graph_def, input_map,
- output_names, input_range):
- if sys.version_info[0] == 3:
- # uint8->quint8 conversion for numpy is not working currently.
- return
-
- quantized_input_map = {}
- for k, v in input_map.items():
- arr = [
- int(
- round((n - input_range[0]) * 255 / (input_range[1] - input_range[
- 0]))) for n in v.flat
- ]
- arr = np.array(arr, np.uint8)
- arr = arr.reshape(v.shape)
- arr = arr.astype(dtypes.quint8.as_numpy_dtype)
- quantized_input_map[k] = arr
- output_tensors = [output_name + ":0" for output_name in output_names]
- float_results = run_graph_def(float_graph_def, input_map, output_tensors)
-
- # Quantize treating the input as quantized in range <input_range>.
- rewriter = quantize_graph.GraphRewriter(float_graph_def, "eightbit",
- input_range)
- graph_def = rewriter.rewrite(output_names)
- results = run_graph_def(graph_def, quantized_input_map, output_tensors)
- for expected, result in zip(float_results, results):
- assert are_tensors_near(expected, result, .5)
- ops = [node.op for node in graph_def.node]
- self.assertEqual(0, ops.count("QuantizeV2") + ops.count("Quantize"))
- self.assertEqual(len(output_names), ops.count("Dequantize"))
-
- # Quantize without treating input as quantized.
- rewriter = quantize_graph.GraphRewriter(
- float_graph_def, "eightbit", quantized_input_range=None)
- graph_def = rewriter.rewrite(output_names)
- results = run_graph_def(graph_def, input_map, output_tensors)
- for expected, result in zip(float_results, results):
- assert are_tensors_near(expected, result, .5)
- ops = [node.op for node in graph_def.node]
- self.assertEqual(
- len(input_map), ops.count("QuantizeV2") + ops.count("Quantize"))
- self.assertEqual(len(output_names), ops.count("Dequantize"))
-
- def test_bias_add_w_fake_quant_w_min_max_vars(self):
- input_node = quantize_graph.create_constant_node(
- "input",
- value=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
- dtype=dtypes.float32,
- shape=[1, 1, 2, 5])
- offset_node = quantize_graph.create_constant_node(
- "offset", value=[1, 2, 3, 4, 5], dtype=dtypes.float32, shape=[5])
- bias_add_node = quantize_graph.create_node(
- "BiasAdd", "bias_add", [input_node.name, offset_node.name])
- quantize_graph.set_attr_dtype(bias_add_node, "T", dtypes.float32)
-
- min_node = quantize_graph.create_constant_node(
- "min_bias_add", value=-.5, dtype=dtypes.float32, shape=[])
- max_node = quantize_graph.create_constant_node(
- "max_bias_add", value=15.5, dtype=dtypes.float32, shape=[])
- fake_quant_node = quantize_graph.create_node(
- "FakeQuantWithMinMaxVars", "fake_quant",
- [bias_add_node.name, min_node.name, max_node.name])
-
- float_graph_def = graph_pb2.GraphDef()
- float_graph_def.node.extend([
- input_node, offset_node, bias_add_node, min_node, max_node,
- fake_quant_node
- ])
- test_graph(float_graph_def, {}, [fake_quant_node.name], log_graph=True)
-
- # Verify there is only one Quantize and one Requantize op.
- # Pass in fallback_quantization_range, although it will have no effect
- # because the FakeQuantWithMinMaxVars are used instead.
- eightbit_rewriter = quantize_graph.GraphRewriter(
- float_graph_def,
- "eightbit",
- quantized_input_range=None,
- fallback_quantization_range=[-100, 100])
- eightbit_graph_def = eightbit_rewriter.rewrite([fake_quant_node.name])
-
- ops = [node.op for node in eightbit_graph_def.node]
- node_names = [node.name for node in eightbit_graph_def.node]
- # No quantize since all inputs are const and can be quantized up-front.
- self.assertEqual(0, ops.count("QuantizeV2") + ops.count("Quantize"))
-
- # One dequantize at the end.
- self.assertEqual(1, ops.count("Dequantize"))
-
- # The fallback constants are not in the graph.
- self.assertEqual(0, node_names.count("fallback_quantization_min_value"))
- self.assertEqual(0, node_names.count("fallback_quantization_max_value"))
-
- def test_bias_add_w_fallback_min_max_vars(self):
- input_node = quantize_graph.create_constant_node(
- "input",
- value=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
- dtype=dtypes.float32,
- shape=[1, 1, 2, 5])
- offset_node = quantize_graph.create_constant_node(
- "offset", value=[1, 2, 3, 4, 5], dtype=dtypes.float32, shape=[5])
- bias_add_node = quantize_graph.create_node(
- "BiasAdd", "bias_add", [input_node.name, offset_node.name])
- quantize_graph.set_attr_dtype(bias_add_node, "T", dtypes.float32)
-
- float_graph_def = graph_pb2.GraphDef()
- float_graph_def.node.extend([input_node, offset_node, bias_add_node])
- test_graph(float_graph_def, {}, [bias_add_node.name], log_graph=True)
-
- # Verify there is only one Quantize, one Requantize op, and no
- # RequantizationRange op.
- eightbit_rewriter = quantize_graph.GraphRewriter(
- float_graph_def,
- "eightbit",
- quantized_input_range=None,
- fallback_quantization_range=[-.5, 15.5])
- eightbit_graph_def = eightbit_rewriter.rewrite([bias_add_node.name])
-
- ops = [node.op for node in eightbit_graph_def.node]
- node_names = [node.name for node in eightbit_graph_def.node]
- # No quantize since all inputs are const and can be quantized up-front.
- self.assertEqual(0, ops.count("QuantizeV2") + ops.count("Quantize"))
-
- # One dequantize at the end.
- self.assertEqual(1, ops.count("Dequantize"))
-
- # No RequantizationRange
- self.assertEqual(0, ops.count("RequantizationRange"))
-
- # The fallback constants are in the graph.
- self.assertEqual(1, node_names.count("fallback_quantization_min_value"))
- self.assertEqual(1, node_names.count("fallback_quantization_max_value"))
-
- def test_remove_redundant_quantization(self):
- a_constant_name = "a_constant"
- a_constant_min_name = "a_constant_min"
- a_constant_max_name = "a_constant_max"
- a_dequantize_name = "a_dequantize"
- a_quantize_name = "a_quantize"
- b_constant_name = "b_constant"
- b_constant_min_name = "b_constant_min"
- b_constant_max_name = "b_constant_max"
- b_dequantize_name = "b_dequantize"
- b_quantize_name = "b_quantize"
- mat_mul_name = "mat_mul"
- graph_def = graph_pb2.GraphDef()
- a_constant = quantize_graph.create_constant_node(
- a_constant_name, value=(0,), dtype=dtypes.quint8, shape=[])
- graph_def.node.extend([a_constant])
- a_constant_min = quantize_graph.create_constant_node(
- a_constant_min_name, value=2, dtype=dtypes.float32, shape=[])
- graph_def.node.extend([a_constant_min])
- a_constant_max = quantize_graph.create_constant_node(
- a_constant_max_name, value=2, dtype=dtypes.float32, shape=[])
- graph_def.node.extend([a_constant_max])
- a_dequantize_node = quantize_graph.create_node(
- "Dequantize", a_dequantize_name,
- [a_constant_name, a_constant_min_name, a_constant_max_name])
- quantize_graph.set_attr_dtype(a_dequantize_node, "T", dtypes.uint8)
- graph_def.node.extend([a_dequantize_node])
- a_quantize_node = quantize_graph.create_node(
- "QuantizeV2", a_quantize_name,
- [a_dequantize_name, a_dequantize_name + ":1", a_dequantize_name + ":2"])
- quantize_graph.set_attr_dtype(a_quantize_node, "T", dtypes.uint8)
- graph_def.node.extend([a_quantize_node])
- b_constant = quantize_graph.create_constant_node(
- b_constant_name, value=(0,), dtype=dtypes.quint8, shape=[])
- graph_def.node.extend([b_constant])
- b_constant_min = quantize_graph.create_constant_node(
- b_constant_min_name, value=3, dtype=dtypes.float32, shape=[])
- graph_def.node.extend([b_constant_min])
- b_constant_max = quantize_graph.create_constant_node(
- b_constant_max_name, value=3, dtype=dtypes.float32, shape=[])
- graph_def.node.extend([b_constant_max])
- b_dequantize_node = quantize_graph.create_node(
- "Dequantize", b_dequantize_name,
- [b_constant_name, b_constant_min_name, b_constant_max_name])
- quantize_graph.set_attr_dtype(b_dequantize_node, "T", dtypes.uint8)
- graph_def.node.extend([b_dequantize_node])
- b_quantize_node = quantize_graph.create_node(
- "QuantizeV2", b_quantize_name,
- [b_dequantize_name, b_dequantize_name + ":1", b_dequantize_name + ":2"])
- quantize_graph.set_attr_dtype(b_quantize_node, "T", dtypes.uint8)
- graph_def.node.extend([b_quantize_node])
- mat_mul_node = quantize_graph.create_node("QuantizedMatMul", mat_mul_name, [
- a_quantize_name, b_quantize_name, a_quantize_name + ":1",
- a_quantize_name + ":2", b_quantize_name + ":1", b_quantize_name + ":2"
- ])
- quantize_graph.set_attr_dtype(mat_mul_node, "T1", dtypes.uint8)
- quantize_graph.set_attr_dtype(mat_mul_node, "T2", dtypes.int32)
- graph_def.node.extend([mat_mul_node])
-
- expected_output = graph_pb2.GraphDef()
- a_constant = quantize_graph.create_constant_node(
- a_constant_name, value=(0,), dtype=dtypes.quint8, shape=[])
- expected_output.node.extend([a_constant])
- a_constant_min = quantize_graph.create_constant_node(
- a_constant_min_name, value=2, dtype=dtypes.float32, shape=[])
- expected_output.node.extend([a_constant_min])
- a_constant_max = quantize_graph.create_constant_node(
- a_constant_max_name, value=2, dtype=dtypes.float32, shape=[])
- expected_output.node.extend([a_constant_max])
- b_constant = quantize_graph.create_constant_node(
- b_constant_name, value=(0,), dtype=dtypes.quint8, shape=[])
- expected_output.node.extend([b_constant])
- b_constant_min = quantize_graph.create_constant_node(
- b_constant_min_name, value=3, dtype=dtypes.float32, shape=[])
- expected_output.node.extend([b_constant_min])
- b_constant_max = quantize_graph.create_constant_node(
- b_constant_max_name, value=3, dtype=dtypes.float32, shape=[])
- expected_output.node.extend([b_constant_max])
- mat_mul_node = quantize_graph.create_node("QuantizedMatMul", mat_mul_name, [
- a_constant_name, b_constant_name, a_constant_min_name,
- a_constant_max_name, b_constant_min_name, b_constant_max_name
- ])
- quantize_graph.set_attr_dtype(mat_mul_node, "T1", dtypes.uint8)
- quantize_graph.set_attr_dtype(mat_mul_node, "T2", dtypes.int32)
- expected_output.node.extend([mat_mul_node])
- expected_output.versions.CopyFrom(graph_def.versions)
- expected_output.library.CopyFrom(graph_def.library)
-
- rewriter = quantize_graph.GraphRewriter(
- graph_def, [mat_mul_name], quantized_input_range=None)
- output = rewriter.remove_redundant_quantization(graph_def)
- stripped_output = graph_util.extract_sub_graph(output, [mat_mul_name])
- self.assertProtoEquals(expected_output, stripped_output)
-
-
-if __name__ == "__main__":
- test.main()
diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl
index 915fee6..bcc89ef 100755
--- a/tensorflow/workspace.bzl
+++ b/tensorflow/workspace.bzl
@@ -1,6 +1,7 @@
# TensorFlow external dependencies that can be loaded in WORKSPACE files.
load("//third_party/gpus:cuda_configure.bzl", "cuda_configure")
+load("//third_party/gpus:rocm_configure.bzl", "rocm_configure")
load("//third_party/tensorrt:tensorrt_configure.bzl", "tensorrt_configure")
load("//third_party:nccl/nccl_configure.bzl", "nccl_configure")
load("//third_party/mkl:build_defs.bzl", "mkl_repository")
@@ -20,9 +21,11 @@
"def_file_filter_configure",
)
load("//third_party/flatbuffers:workspace.bzl", flatbuffers = "repo")
+load("//third_party/icu:workspace.bzl", icu = "repo")
def initialize_third_party():
flatbuffers()
+ icu()
# Sanitize a dependency so that it works correctly from code that includes
# TensorFlow as a submodule.
@@ -43,6 +46,7 @@
sycl_configure(name = "local_config_sycl")
syslibs_configure(name = "local_config_syslibs")
python_configure(name = "local_config_python")
+ rocm_configure(name = "local_config_rocm")
initialize_third_party()
@@ -53,39 +57,39 @@
# Point //external/local_config_arm_compiler to //external/arm_compiler
arm_compiler_configure(
name = "local_config_arm_compiler",
- remote_config_repo = "../arm_compiler",
build_file = clean_dep("//third_party/toolchains/cpus/arm:BUILD"),
+ remote_config_repo = "../arm_compiler",
)
mkl_repository(
name = "mkl_linux",
+ build_file = clean_dep("//third_party/mkl:mkl.BUILD"),
+ sha256 = "e2233534a9d15c387e22260997af4312a39e9f86f791768409be273b5453c4e6",
+ strip_prefix = "mklml_lnx_2019.0.20180710",
urls = [
"https://mirror.bazel.build/github.com/intel/mkl-dnn/releases/download/v0.16/mklml_lnx_2019.0.20180710.tgz",
"https://github.com/intel/mkl-dnn/releases/download/v0.16/mklml_lnx_2019.0.20180710.tgz",
],
- sha256 = "e2233534a9d15c387e22260997af4312a39e9f86f791768409be273b5453c4e6",
- strip_prefix = "mklml_lnx_2019.0.20180710",
- build_file = clean_dep("//third_party/mkl:mkl.BUILD"),
)
mkl_repository(
name = "mkl_windows",
+ build_file = clean_dep("//third_party/mkl:mkl.BUILD"),
+ sha256 = "3fdcff17b018a0082491adf3ba143358265336a801646e46e0191ec8d58d24a2",
+ strip_prefix = "mklml_win_2019.0.20180710",
urls = [
"https://mirror.bazel.build/github.com/intel/mkl-dnn/releases/download/v0.16/mklml_win_2019.0.20180710.zip",
"https://github.com/intel/mkl-dnn/releases/download/v0.16/mklml_win_2019.0.20180710.zip",
],
- sha256 = "3fdcff17b018a0082491adf3ba143358265336a801646e46e0191ec8d58d24a2",
- strip_prefix = "mklml_win_2019.0.20180710",
- build_file = clean_dep("//third_party/mkl:mkl.BUILD"),
)
mkl_repository(
name = "mkl_darwin",
+ build_file = clean_dep("//third_party/mkl:mkl.BUILD"),
+ sha256 = "411a30014a938eb83fb9f37b3dbe8e371b106fc1dd621fc23123cadc72737ce6",
+ strip_prefix = "mklml_mac_2019.0.20180710",
urls = [
"https://mirror.bazel.build/github.com/intel/mkl-dnn/releases/download/v0.16/mklml_mac_2019.0.20180710.tgz",
"https://github.com/intel/mkl-dnn/releases/download/v0.16/mklml_mac_2019.0.20180710.tgz",
],
- sha256 = "411a30014a938eb83fb9f37b3dbe8e371b106fc1dd621fc23123cadc72737ce6",
- strip_prefix = "mklml_mac_2019.0.20180710",
- build_file = clean_dep("//third_party/mkl:mkl.BUILD"),
)
if path_prefix:
@@ -94,39 +98,40 @@
tf_http_archive(
name = "mkl_dnn",
+ build_file = clean_dep("//third_party/mkl_dnn:mkldnn.BUILD"),
+ sha256 = "363cc9239eacf8e7917753c6d8c94f767e4cd049160d0654a61ef32d5e1b3049",
+ strip_prefix = "mkl-dnn-4e333787e0d66a1dca1218e99a891d493dbc8ef1",
urls = [
"https://mirror.bazel.build/github.com/intel/mkl-dnn/archive/4e333787e0d66a1dca1218e99a891d493dbc8ef1.tar.gz",
"https://github.com/intel/mkl-dnn/archive/4e333787e0d66a1dca1218e99a891d493dbc8ef1.tar.gz",
],
- sha256 = "363cc9239eacf8e7917753c6d8c94f767e4cd049160d0654a61ef32d5e1b3049",
- strip_prefix = "mkl-dnn-4e333787e0d66a1dca1218e99a891d493dbc8ef1",
- build_file = clean_dep("//third_party/mkl_dnn:mkldnn.BUILD"),
)
tf_http_archive(
name = "com_google_absl",
- urls = [
- "https://mirror.bazel.build/github.com/abseil/abseil-cpp/archive/e01d95528ea2137a4a27a88d1f57c6cb260aafed.tar.gz",
- "https://github.com/abseil/abseil-cpp/archive/e01d95528ea2137a4a27a88d1f57c6cb260aafed.tar.gz",
- ],
- sha256 = "84043ed402d2a2a6ba4cdddb7e85118b1158fd81fe4ac3a14adc343d054c1e2e",
- strip_prefix = "abseil-cpp-e01d95528ea2137a4a27a88d1f57c6cb260aafed",
build_file = clean_dep("//third_party:com_google_absl.BUILD"),
+ sha256 = "7dd09690ae7ca4551de3111d4a86b75b23ec17445f273d3c42bdcdc1c7b02e4e",
+ strip_prefix = "abseil-cpp-48cd2c3f351ff188bc85684b84a91b6e6d17d896",
+ urls = [
+ "https://mirror.bazel.build/github.com/abseil/abseil-cpp/archive/48cd2c3f351ff188bc85684b84a91b6e6d17d896.tar.gz",
+ "https://github.com/abseil/abseil-cpp/archive/48cd2c3f351ff188bc85684b84a91b6e6d17d896.tar.gz",
+ ],
)
tf_http_archive(
name = "eigen_archive",
+ build_file = clean_dep("//third_party:eigen.BUILD"),
+ sha256 = "d956415d784fa4e42b6a2a45c32556d6aec9d0a3d8ef48baee2522ab762556a9",
+ strip_prefix = "eigen-eigen-fd6845384b86",
urls = [
"https://mirror.bazel.build/bitbucket.org/eigen/eigen/get/fd6845384b86.tar.gz",
"https://bitbucket.org/eigen/eigen/get/fd6845384b86.tar.gz",
],
- sha256 = "d956415d784fa4e42b6a2a45c32556d6aec9d0a3d8ef48baee2522ab762556a9",
- strip_prefix = "eigen-eigen-fd6845384b86",
- build_file = clean_dep("//third_party:eigen.BUILD"),
)
tf_http_archive(
name = "arm_compiler",
+ build_file = clean_dep("//:arm_compiler.BUILD"),
sha256 = "970285762565c7890c6c087d262b0a18286e7d0384f13a37786d8521773bc969",
strip_prefix = "tools-0e906ebc527eab1cdbf7adabff5b474da9562e9f/arm-bcm2708/arm-rpi-4.9.3-linux-gnueabihf",
urls = [
@@ -135,216 +140,211 @@
# remove the whitelist entry in third_party/repo.bzl.
# "https://github.com/raspberrypi/tools/archive/0e906ebc527eab1cdbf7adabff5b474da9562e9f.tar.gz",
],
- build_file = clean_dep("//:arm_compiler.BUILD"),
)
tf_http_archive(
name = "libxsmm_archive",
+ build_file = clean_dep("//third_party:libxsmm.BUILD"),
+ sha256 = "cd8532021352b4a0290d209f7f9bfd7c2411e08286a893af3577a43457287bfa",
+ strip_prefix = "libxsmm-1.9",
urls = [
"https://mirror.bazel.build/github.com/hfp/libxsmm/archive/1.9.tar.gz",
"https://github.com/hfp/libxsmm/archive/1.9.tar.gz",
],
- sha256 = "cd8532021352b4a0290d209f7f9bfd7c2411e08286a893af3577a43457287bfa",
- strip_prefix = "libxsmm-1.9",
- build_file = clean_dep("//third_party:libxsmm.BUILD"),
)
tf_http_archive(
name = "ortools_archive",
+ build_file = clean_dep("//third_party:ortools.BUILD"),
+ sha256 = "d025a95f78b5fc5eaa4da5f395f23d11c23cf7dbd5069f1f627f002de87b86b9",
+ strip_prefix = "or-tools-6.7.2/src",
urls = [
"https://mirror.bazel.build/github.com/google/or-tools/archive/v6.7.2.tar.gz",
"https://github.com/google/or-tools/archive/v6.7.2.tar.gz",
],
- sha256 = "d025a95f78b5fc5eaa4da5f395f23d11c23cf7dbd5069f1f627f002de87b86b9",
- strip_prefix = "or-tools-6.7.2/src",
- build_file = clean_dep("//third_party:ortools.BUILD"),
)
tf_http_archive(
name = "com_googlesource_code_re2",
+ sha256 = "803c7811146edeef8f91064de37c6f19136ff01a2a8cdb3230e940b2fd9f07fe",
+ strip_prefix = "re2-2018-07-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",
],
- sha256 = "803c7811146edeef8f91064de37c6f19136ff01a2a8cdb3230e940b2fd9f07fe",
- strip_prefix = "re2-2018-07-01",
- system_build_file = clean_dep("//third_party/systemlibs:re2.BUILD"),
)
tf_http_archive(
name = "com_github_googlecloudplatform_google_cloud_cpp",
- urls = [
- "https://mirror.bazel.build/github.com/GoogleCloudPlatform/google-cloud-cpp/archive/14760a86c4ffab9943b476305c4fe927ad95db1c.tar.gz",
- "https://github.com/GoogleCloudPlatform/google-cloud-cpp/archive/14760a86c4ffab9943b476305c4fe927ad95db1c.tar.gz",
- ],
sha256 = "fdd3b3aecce60987e5525e55bf3a21d68a8695320bd5b980775af6507eec3944",
strip_prefix = "google-cloud-cpp-14760a86c4ffab9943b476305c4fe927ad95db1c",
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/14760a86c4ffab9943b476305c4fe927ad95db1c.tar.gz",
+ "https://github.com/GoogleCloudPlatform/google-cloud-cpp/archive/14760a86c4ffab9943b476305c4fe927ad95db1c.tar.gz",
+ ],
)
tf_http_archive(
name = "com_github_googleapis_googleapis",
+ build_file = clean_dep("//third_party:googleapis.BUILD"),
+ sha256 = "824870d87a176f26bcef663e92051f532fac756d1a06b404055dc078425f4378",
+ strip_prefix = "googleapis-f81082ea1e2f85c43649bee26e0d9871d4b41cdb",
+ system_build_file = clean_dep("//third_party/systemlibs:googleapis.BUILD"),
urls = [
"https://mirror.bazel.build/github.com/googleapis/googleapis/archive/f81082ea1e2f85c43649bee26e0d9871d4b41cdb.zip",
"https://github.com/googleapis/googleapis/archive/f81082ea1e2f85c43649bee26e0d9871d4b41cdb.zip",
],
- sha256 = "824870d87a176f26bcef663e92051f532fac756d1a06b404055dc078425f4378",
- strip_prefix = "googleapis-f81082ea1e2f85c43649bee26e0d9871d4b41cdb",
- build_file = clean_dep("//third_party:googleapis.BUILD"),
- system_build_file = clean_dep("//third_party/systemlibs:googleapis.BUILD"),
)
tf_http_archive(
name = "gemmlowp",
+ sha256 = "b87faa7294dfcc5d678f22a59d2c01ca94ea1e2a3b488c38a95a67889ed0a658",
+ strip_prefix = "gemmlowp-38ebac7b059e84692f53e5938f97a9943c120d98",
urls = [
"https://mirror.bazel.build/github.com/google/gemmlowp/archive/38ebac7b059e84692f53e5938f97a9943c120d98.zip",
"https://github.com/google/gemmlowp/archive/38ebac7b059e84692f53e5938f97a9943c120d98.zip",
],
- sha256 = "b87faa7294dfcc5d678f22a59d2c01ca94ea1e2a3b488c38a95a67889ed0a658",
- strip_prefix = "gemmlowp-38ebac7b059e84692f53e5938f97a9943c120d98",
)
tf_http_archive(
name = "farmhash_archive",
+ build_file = clean_dep("//third_party:farmhash.BUILD"),
+ sha256 = "6560547c63e4af82b0f202cb710ceabb3f21347a4b996db565a411da5b17aba0",
+ strip_prefix = "farmhash-816a4ae622e964763ca0862d9dbd19324a1eaf45",
urls = [
"https://mirror.bazel.build/github.com/google/farmhash/archive/816a4ae622e964763ca0862d9dbd19324a1eaf45.tar.gz",
"https://github.com/google/farmhash/archive/816a4ae622e964763ca0862d9dbd19324a1eaf45.tar.gz",
],
- sha256 = "6560547c63e4af82b0f202cb710ceabb3f21347a4b996db565a411da5b17aba0",
- strip_prefix = "farmhash-816a4ae622e964763ca0862d9dbd19324a1eaf45",
- build_file = clean_dep("//third_party:farmhash.BUILD"),
)
tf_http_archive(
name = "highwayhash",
+ build_file = clean_dep("//third_party:highwayhash.BUILD"),
+ sha256 = "9c3e0e87d581feeb0c18d814d98f170ff23e62967a2bd6855847f0b2fe598a37",
+ strip_prefix = "highwayhash-fd3d9af80465e4383162e4a7c5e2f406e82dd968",
urls = [
"http://mirror.bazel.build/github.com/google/highwayhash/archive/fd3d9af80465e4383162e4a7c5e2f406e82dd968.tar.gz",
"https://github.com/google/highwayhash/archive/fd3d9af80465e4383162e4a7c5e2f406e82dd968.tar.gz",
],
- sha256 = "9c3e0e87d581feeb0c18d814d98f170ff23e62967a2bd6855847f0b2fe598a37",
- strip_prefix = "highwayhash-fd3d9af80465e4383162e4a7c5e2f406e82dd968",
- build_file = clean_dep("//third_party:highwayhash.BUILD"),
)
tf_http_archive(
name = "nasm",
+ build_file = clean_dep("//third_party:nasm.BUILD"),
+ sha256 = "63ec86477ad3f0f6292325fd89e1d93aea2e2fd490070863f17d48f7cd387011",
+ strip_prefix = "nasm-2.13.03",
+ system_build_file = clean_dep("//third_party/systemlibs:nasm.BUILD"),
urls = [
"https://mirror.bazel.build/www.nasm.us/pub/nasm/releasebuilds/2.13.03/nasm-2.13.03.tar.bz2",
"http://pkgs.fedoraproject.org/repo/pkgs/nasm/nasm-2.13.03.tar.bz2/sha512/d7a6b4cee8dfd603d8d4c976e5287b5cc542fa0b466ff989b743276a6e28114e64289bf02a7819eca63142a5278aa6eed57773007e5f589e15768e6456a8919d/nasm-2.13.03.tar.bz2",
"http://www.nasm.us/pub/nasm/releasebuilds/2.13.03/nasm-2.13.03.tar.bz2",
],
- sha256 = "63ec86477ad3f0f6292325fd89e1d93aea2e2fd490070863f17d48f7cd387011",
- strip_prefix = "nasm-2.13.03",
- build_file = clean_dep("//third_party:nasm.BUILD"),
- system_build_file = clean_dep("//third_party/systemlibs:nasm.BUILD"),
)
tf_http_archive(
name = "jpeg",
+ build_file = clean_dep("//third_party/jpeg:jpeg.BUILD"),
+ sha256 = "f892fff427ab3adffc289363eac26d197ce3ccacefe5f5822377348a8166069b",
+ strip_prefix = "libjpeg-turbo-2.0.0",
+ system_build_file = clean_dep("//third_party/systemlibs:jpeg.BUILD"),
urls = [
"https://mirror.bazel.build/github.com/libjpeg-turbo/libjpeg-turbo/archive/2.0.0.tar.gz",
"https://github.com/libjpeg-turbo/libjpeg-turbo/archive/2.0.0.tar.gz",
],
- sha256 = "f892fff427ab3adffc289363eac26d197ce3ccacefe5f5822377348a8166069b",
- strip_prefix = "libjpeg-turbo-2.0.0",
- build_file = clean_dep("//third_party/jpeg:jpeg.BUILD"),
- system_build_file = clean_dep("//third_party/systemlibs:jpeg.BUILD"),
)
tf_http_archive(
name = "png_archive",
+ build_file = clean_dep("//third_party:png.BUILD"),
+ patch_file = clean_dep("//third_party:png_fix_rpi.patch"),
+ sha256 = "e45ce5f68b1d80e2cb9a2b601605b374bdf51e1798ef1c2c2bd62131dfcf9eef",
+ strip_prefix = "libpng-1.6.34",
+ system_build_file = clean_dep("//third_party/systemlibs:png.BUILD"),
urls = [
"https://mirror.bazel.build/github.com/glennrp/libpng/archive/v1.6.34.tar.gz",
"https://github.com/glennrp/libpng/archive/v1.6.34.tar.gz",
],
- sha256 = "e45ce5f68b1d80e2cb9a2b601605b374bdf51e1798ef1c2c2bd62131dfcf9eef",
- strip_prefix = "libpng-1.6.34",
- build_file = clean_dep("//third_party:png.BUILD"),
- patch_file = clean_dep("//third_party:png_fix_rpi.patch"),
- system_build_file = clean_dep("//third_party/systemlibs:png.BUILD"),
)
tf_http_archive(
name = "org_sqlite",
+ build_file = clean_dep("//third_party:sqlite.BUILD"),
+ sha256 = "ad68c1216c3a474cf360c7581a4001e952515b3649342100f2d7ca7c8e313da6",
+ strip_prefix = "sqlite-amalgamation-3240000",
+ system_build_file = clean_dep("//third_party/systemlibs:sqlite.BUILD"),
urls = [
"https://mirror.bazel.build/www.sqlite.org/2018/sqlite-amalgamation-3240000.zip",
"https://www.sqlite.org/2018/sqlite-amalgamation-3240000.zip",
],
- sha256 = "ad68c1216c3a474cf360c7581a4001e952515b3649342100f2d7ca7c8e313da6",
- strip_prefix = "sqlite-amalgamation-3240000",
- build_file = clean_dep("//third_party:sqlite.BUILD"),
- system_build_file = clean_dep("//third_party/systemlibs:sqlite.BUILD"),
)
tf_http_archive(
name = "gif_archive",
+ build_file = clean_dep("//third_party:gif.BUILD"),
+ sha256 = "34a7377ba834397db019e8eb122e551a49c98f49df75ec3fcc92b9a794a4f6d1",
+ strip_prefix = "giflib-5.1.4",
+ system_build_file = clean_dep("//third_party/systemlibs:gif.BUILD"),
urls = [
"https://mirror.bazel.build/ufpr.dl.sourceforge.net/project/giflib/giflib-5.1.4.tar.gz",
"http://pilotfiber.dl.sourceforge.net/project/giflib/giflib-5.1.4.tar.gz",
],
- sha256 = "34a7377ba834397db019e8eb122e551a49c98f49df75ec3fcc92b9a794a4f6d1",
- strip_prefix = "giflib-5.1.4",
- build_file = clean_dep("//third_party:gif.BUILD"),
- system_build_file = clean_dep("//third_party/systemlibs:gif.BUILD"),
)
tf_http_archive(
name = "six_archive",
+ build_file = clean_dep("//third_party:six.BUILD"),
+ sha256 = "105f8d68616f8248e24bf0e9372ef04d3cc10104f1980f54d57b2ce73a5ad56a",
+ strip_prefix = "six-1.10.0",
+ system_build_file = clean_dep("//third_party/systemlibs:six.BUILD"),
urls = [
"https://mirror.bazel.build/pypi.python.org/packages/source/s/six/six-1.10.0.tar.gz",
"https://pypi.python.org/packages/source/s/six/six-1.10.0.tar.gz",
],
- sha256 = "105f8d68616f8248e24bf0e9372ef04d3cc10104f1980f54d57b2ce73a5ad56a",
- strip_prefix = "six-1.10.0",
- build_file = clean_dep("//third_party:six.BUILD"),
- system_build_file = clean_dep("//third_party/systemlibs:six.BUILD"),
)
tf_http_archive(
name = "astor_archive",
+ build_file = clean_dep("//third_party:astor.BUILD"),
+ sha256 = "ff6d2e2962d834acb125cc4dcc80c54a8c17c253f4cc9d9c43b5102a560bb75d",
+ strip_prefix = "astor-0.6.2",
+ system_build_file = clean_dep("//third_party/systemlibs:astor.BUILD"),
urls = [
"https://mirror.bazel.build/pypi.python.org/packages/d8/be/c4276b3199ec3feee2a88bc64810fbea8f26d961e0a4cd9c68387a9f35de/astor-0.6.2.tar.gz",
"https://pypi.python.org/packages/d8/be/c4276b3199ec3feee2a88bc64810fbea8f26d961e0a4cd9c68387a9f35de/astor-0.6.2.tar.gz",
],
- sha256 = "ff6d2e2962d834acb125cc4dcc80c54a8c17c253f4cc9d9c43b5102a560bb75d",
- strip_prefix = "astor-0.6.2",
- build_file = clean_dep("//third_party:astor.BUILD"),
- system_build_file = clean_dep("//third_party/systemlibs:astor.BUILD"),
)
tf_http_archive(
name = "gast_archive",
+ build_file = clean_dep("//third_party:gast.BUILD"),
+ sha256 = "7068908321ecd2774f145193c4b34a11305bd104b4551b09273dfd1d6a374930",
+ strip_prefix = "gast-0.2.0",
+ system_build_file = clean_dep("//third_party/systemlibs:gast.BUILD"),
urls = [
"https://mirror.bazel.build/pypi.python.org/packages/5c/78/ff794fcae2ce8aa6323e789d1f8b3b7765f601e7702726f430e814822b96/gast-0.2.0.tar.gz",
"https://pypi.python.org/packages/5c/78/ff794fcae2ce8aa6323e789d1f8b3b7765f601e7702726f430e814822b96/gast-0.2.0.tar.gz",
],
- sha256 = "7068908321ecd2774f145193c4b34a11305bd104b4551b09273dfd1d6a374930",
- strip_prefix = "gast-0.2.0",
- build_file = clean_dep("//third_party:gast.BUILD"),
- system_build_file = clean_dep("//third_party/systemlibs:gast.BUILD"),
)
tf_http_archive(
name = "termcolor_archive",
+ build_file = clean_dep("//third_party:termcolor.BUILD"),
+ sha256 = "1d6d69ce66211143803fbc56652b41d73b4a400a2891d7bf7a1cdf4c02de613b",
+ strip_prefix = "termcolor-1.1.0",
+ system_build_file = clean_dep("//third_party/systemlibs:termcolor.BUILD"),
urls = [
"https://mirror.bazel.build/pypi.python.org/packages/8a/48/a76be51647d0eb9f10e2a4511bf3ffb8cc1e6b14e9e4fab46173aa79f981/termcolor-1.1.0.tar.gz",
"https://pypi.python.org/packages/8a/48/a76be51647d0eb9f10e2a4511bf3ffb8cc1e6b14e9e4fab46173aa79f981/termcolor-1.1.0.tar.gz",
],
- sha256 = "1d6d69ce66211143803fbc56652b41d73b4a400a2891d7bf7a1cdf4c02de613b",
- strip_prefix = "termcolor-1.1.0",
- build_file = clean_dep("//third_party:termcolor.BUILD"),
- system_build_file = clean_dep("//third_party/systemlibs:termcolor.BUILD"),
)
tf_http_archive(
name = "absl_py",
- urls = [
- "https://mirror.bazel.build/github.com/abseil/abseil-py/archive/pypi-v0.2.2.tar.gz",
- "https://github.com/abseil/abseil-py/archive/pypi-v0.2.2.tar.gz",
- ],
sha256 = "95160f778a62c7a60ddeadc7bf2d83f85a23a27359814aca12cf949e896fa82c",
strip_prefix = "abseil-py-pypi-v0.2.2",
system_build_file = clean_dep("//third_party/systemlibs:absl_py.BUILD"),
@@ -352,17 +352,21 @@
"//third_party/systemlibs:absl_py.absl.flags.BUILD": "absl/flags/BUILD",
"//third_party/systemlibs:absl_py.absl.testing.BUILD": "absl/testing/BUILD",
},
+ urls = [
+ "https://mirror.bazel.build/github.com/abseil/abseil-py/archive/pypi-v0.2.2.tar.gz",
+ "https://github.com/abseil/abseil-py/archive/pypi-v0.2.2.tar.gz",
+ ],
)
tf_http_archive(
name = "org_python_pypi_backports_weakref",
+ build_file = clean_dep("//third_party:backports_weakref.BUILD"),
+ sha256 = "8813bf712a66b3d8b85dc289e1104ed220f1878cf981e2fe756dfaabe9a82892",
+ strip_prefix = "backports.weakref-1.0rc1/src",
urls = [
"https://mirror.bazel.build/pypi.python.org/packages/bc/cc/3cdb0a02e7e96f6c70bd971bc8a90b8463fda83e264fa9c5c1c98ceabd81/backports.weakref-1.0rc1.tar.gz",
"https://pypi.python.org/packages/bc/cc/3cdb0a02e7e96f6c70bd971bc8a90b8463fda83e264fa9c5c1c98ceabd81/backports.weakref-1.0rc1.tar.gz",
],
- sha256 = "8813bf712a66b3d8b85dc289e1104ed220f1878cf981e2fe756dfaabe9a82892",
- strip_prefix = "backports.weakref-1.0rc1/src",
- build_file = clean_dep("//third_party:backports_weakref.BUILD"),
)
filegroup_external(
@@ -385,9 +389,9 @@
tf_http_archive(
name = "protobuf_archive",
- urls = PROTOBUF_URLS,
sha256 = PROTOBUF_SHA256,
strip_prefix = PROTOBUF_STRIP_PREFIX,
+ urls = PROTOBUF_URLS,
)
# We need to import the protobuf library under the names com_google_protobuf
@@ -395,222 +399,222 @@
# Unfortunately there is no way to alias http_archives at the moment.
tf_http_archive(
name = "com_google_protobuf",
- urls = PROTOBUF_URLS,
sha256 = PROTOBUF_SHA256,
strip_prefix = PROTOBUF_STRIP_PREFIX,
+ urls = PROTOBUF_URLS,
)
tf_http_archive(
name = "com_google_protobuf_cc",
- urls = PROTOBUF_URLS,
sha256 = PROTOBUF_SHA256,
strip_prefix = PROTOBUF_STRIP_PREFIX,
+ urls = PROTOBUF_URLS,
)
tf_http_archive(
name = "nsync",
+ sha256 = "692f9b30e219f71a6371b98edd39cef3cbda35ac3abc4cd99ce19db430a5591a",
+ strip_prefix = "nsync-1.20.1",
+ system_build_file = clean_dep("//third_party/systemlibs:nsync.BUILD"),
urls = [
"https://mirror.bazel.build/github.com/google/nsync/archive/1.20.1.tar.gz",
"https://github.com/google/nsync/archive/1.20.1.tar.gz",
],
- sha256 = "692f9b30e219f71a6371b98edd39cef3cbda35ac3abc4cd99ce19db430a5591a",
- strip_prefix = "nsync-1.20.1",
- system_build_file = clean_dep("//third_party/systemlibs:nsync.BUILD"),
)
tf_http_archive(
name = "com_google_googletest",
+ sha256 = "353ab86e35cea1cd386115279cf4b16695bbf21b897bfbf2721cf4cb5f64ade8",
+ strip_prefix = "googletest-997d343dd680e541ef96ce71ee54a91daf2577a0",
urls = [
"https://mirror.bazel.build/github.com/google/googletest/archive/997d343dd680e541ef96ce71ee54a91daf2577a0.zip",
"https://github.com/google/googletest/archive/997d343dd680e541ef96ce71ee54a91daf2577a0.zip",
],
- sha256 = "353ab86e35cea1cd386115279cf4b16695bbf21b897bfbf2721cf4cb5f64ade8",
- strip_prefix = "googletest-997d343dd680e541ef96ce71ee54a91daf2577a0",
)
tf_http_archive(
name = "com_github_gflags_gflags",
+ sha256 = "ae27cdbcd6a2f935baa78e4f21f675649271634c092b1be01469440495609d0e",
+ strip_prefix = "gflags-2.2.1",
urls = [
"https://mirror.bazel.build/github.com/gflags/gflags/archive/v2.2.1.tar.gz",
"https://github.com/gflags/gflags/archive/v2.2.1.tar.gz",
],
- sha256 = "ae27cdbcd6a2f935baa78e4f21f675649271634c092b1be01469440495609d0e",
- strip_prefix = "gflags-2.2.1",
)
tf_http_archive(
name = "pcre",
+ build_file = clean_dep("//third_party:pcre.BUILD"),
sha256 = "69acbc2fbdefb955d42a4c606dfde800c2885711d2979e356c0636efde9ec3b5",
+ strip_prefix = "pcre-8.42",
+ system_build_file = clean_dep("//third_party/systemlibs:pcre.BUILD"),
urls = [
"https://mirror.bazel.build/ftp.exim.org/pub/pcre/pcre-8.42.tar.gz",
"http://ftp.exim.org/pub/pcre/pcre-8.42.tar.gz",
],
- strip_prefix = "pcre-8.42",
- build_file = clean_dep("//third_party:pcre.BUILD"),
- system_build_file = clean_dep("//third_party/systemlibs:pcre.BUILD"),
)
tf_http_archive(
name = "swig",
+ build_file = clean_dep("//third_party:swig.BUILD"),
sha256 = "58a475dbbd4a4d7075e5fe86d4e54c9edde39847cdb96a3053d87cb64a23a453",
+ strip_prefix = "swig-3.0.8",
+ system_build_file = clean_dep("//third_party/systemlibs:swig.BUILD"),
urls = [
"https://mirror.bazel.build/ufpr.dl.sourceforge.net/project/swig/swig/swig-3.0.8/swig-3.0.8.tar.gz",
"http://ufpr.dl.sourceforge.net/project/swig/swig/swig-3.0.8/swig-3.0.8.tar.gz",
"http://pilotfiber.dl.sourceforge.net/project/swig/swig/swig-3.0.8/swig-3.0.8.tar.gz",
],
- strip_prefix = "swig-3.0.8",
- build_file = clean_dep("//third_party:swig.BUILD"),
- system_build_file = clean_dep("//third_party/systemlibs:swig.BUILD"),
)
tf_http_archive(
name = "curl",
+ build_file = clean_dep("//third_party:curl.BUILD"),
sha256 = "e9c37986337743f37fd14fe8737f246e97aec94b39d1b71e8a5973f72a9fc4f5",
+ strip_prefix = "curl-7.60.0",
+ system_build_file = clean_dep("//third_party/systemlibs:curl.BUILD"),
urls = [
"https://mirror.bazel.build/curl.haxx.se/download/curl-7.60.0.tar.gz",
"https://curl.haxx.se/download/curl-7.60.0.tar.gz",
],
- strip_prefix = "curl-7.60.0",
- build_file = clean_dep("//third_party:curl.BUILD"),
- system_build_file = clean_dep("//third_party/systemlibs:curl.BUILD"),
)
tf_http_archive(
name = "grpc",
+ sha256 = "50db9cf2221354485eb7c3bd55a4c27190caef7048a2a1a15fbe60a498f98b44",
+ strip_prefix = "grpc-1.13.0",
+ system_build_file = clean_dep("//third_party/systemlibs:grpc.BUILD"),
urls = [
"https://mirror.bazel.build/github.com/grpc/grpc/archive/v1.13.0.tar.gz",
"https://github.com/grpc/grpc/archive/v1.13.0.tar.gz",
],
- sha256 = "50db9cf2221354485eb7c3bd55a4c27190caef7048a2a1a15fbe60a498f98b44",
- strip_prefix = "grpc-1.13.0",
- system_build_file = clean_dep("//third_party/systemlibs:grpc.BUILD"),
)
tf_http_archive(
name = "linenoise",
+ build_file = clean_dep("//third_party:linenoise.BUILD"),
sha256 = "7f51f45887a3d31b4ce4fa5965210a5e64637ceac12720cfce7954d6a2e812f7",
+ strip_prefix = "linenoise-c894b9e59f02203dbe4e2be657572cf88c4230c3",
urls = [
"https://mirror.bazel.build/github.com/antirez/linenoise/archive/c894b9e59f02203dbe4e2be657572cf88c4230c3.tar.gz",
"https://github.com/antirez/linenoise/archive/c894b9e59f02203dbe4e2be657572cf88c4230c3.tar.gz",
],
- strip_prefix = "linenoise-c894b9e59f02203dbe4e2be657572cf88c4230c3",
- build_file = clean_dep("//third_party:linenoise.BUILD"),
)
# TODO(phawkins): currently, this rule uses an unofficial LLVM mirror.
# Switch to an official source of snapshots if/when possible.
tf_http_archive(
name = "llvm",
- urls = [
- "https://mirror.bazel.build/github.com/llvm-mirror/llvm/archive/7167e4d196a50f78abe8af6553c943d50b757a13.tar.gz",
- "https://github.com/llvm-mirror/llvm/archive/7167e4d196a50f78abe8af6553c943d50b757a13.tar.gz",
- ],
- sha256 = "11d933232b27531abc83592fc9f03e7f928e504c7d478eeaba51efa929a3d9df",
- strip_prefix = "llvm-7167e4d196a50f78abe8af6553c943d50b757a13",
build_file = clean_dep("//third_party/llvm:llvm.autogenerated.BUILD"),
+ sha256 = "a4f8bfe7e3e69069934a87e612a1d4d3b8b6af13e0f1213a42a6046e1bcd50d8",
+ strip_prefix = "llvm-d3429e96fe1e45b1dc0106463832523f37faf271",
+ urls = [
+ "https://mirror.bazel.build/github.com/llvm-mirror/llvm/archive/d3429e96fe1e45b1dc0106463832523f37faf271.tar.gz",
+ "https://github.com/llvm-mirror/llvm/archive/d3429e96fe1e45b1dc0106463832523f37faf271.tar.gz",
+ ],
)
tf_http_archive(
name = "lmdb",
+ build_file = clean_dep("//third_party:lmdb.BUILD"),
+ sha256 = "f3927859882eb608868c8c31586bb7eb84562a40a6bf5cc3e13b6b564641ea28",
+ strip_prefix = "lmdb-LMDB_0.9.22/libraries/liblmdb",
+ system_build_file = clean_dep("//third_party/systemlibs:lmdb.BUILD"),
urls = [
"https://mirror.bazel.build/github.com/LMDB/lmdb/archive/LMDB_0.9.22.tar.gz",
"https://github.com/LMDB/lmdb/archive/LMDB_0.9.22.tar.gz",
],
- sha256 = "f3927859882eb608868c8c31586bb7eb84562a40a6bf5cc3e13b6b564641ea28",
- strip_prefix = "lmdb-LMDB_0.9.22/libraries/liblmdb",
- build_file = clean_dep("//third_party:lmdb.BUILD"),
- system_build_file = clean_dep("//third_party/systemlibs:lmdb.BUILD"),
)
tf_http_archive(
name = "jsoncpp_git",
+ build_file = clean_dep("//third_party:jsoncpp.BUILD"),
+ sha256 = "c49deac9e0933bcb7044f08516861a2d560988540b23de2ac1ad443b219afdb6",
+ strip_prefix = "jsoncpp-1.8.4",
+ system_build_file = clean_dep("//third_party/systemlibs:jsoncpp.BUILD"),
urls = [
"https://mirror.bazel.build/github.com/open-source-parsers/jsoncpp/archive/1.8.4.tar.gz",
"https://github.com/open-source-parsers/jsoncpp/archive/1.8.4.tar.gz",
],
- sha256 = "c49deac9e0933bcb7044f08516861a2d560988540b23de2ac1ad443b219afdb6",
- strip_prefix = "jsoncpp-1.8.4",
- build_file = clean_dep("//third_party:jsoncpp.BUILD"),
- system_build_file = clean_dep("//third_party/systemlibs:jsoncpp.BUILD"),
)
tf_http_archive(
name = "boringssl",
+ sha256 = "1188e29000013ed6517168600fc35a010d58c5d321846d6a6dfee74e4c788b45",
+ strip_prefix = "boringssl-7f634429a04abc48e2eb041c81c5235816c96514",
+ system_build_file = clean_dep("//third_party/systemlibs:boringssl.BUILD"),
urls = [
"https://mirror.bazel.build/github.com/google/boringssl/archive/7f634429a04abc48e2eb041c81c5235816c96514.tar.gz",
"https://github.com/google/boringssl/archive/7f634429a04abc48e2eb041c81c5235816c96514.tar.gz",
],
- sha256 = "1188e29000013ed6517168600fc35a010d58c5d321846d6a6dfee74e4c788b45",
- strip_prefix = "boringssl-7f634429a04abc48e2eb041c81c5235816c96514",
- system_build_file = clean_dep("//third_party/systemlibs:boringssl.BUILD"),
)
tf_http_archive(
name = "zlib_archive",
+ build_file = clean_dep("//third_party:zlib.BUILD"),
+ sha256 = "c3e5e9fdd5004dcb542feda5ee4f0ff0744628baf8ed2dd5d66f8ca1197cb1a1",
+ strip_prefix = "zlib-1.2.11",
+ system_build_file = clean_dep("//third_party/systemlibs:zlib.BUILD"),
urls = [
"https://mirror.bazel.build/zlib.net/zlib-1.2.11.tar.gz",
"https://zlib.net/zlib-1.2.11.tar.gz",
],
- sha256 = "c3e5e9fdd5004dcb542feda5ee4f0ff0744628baf8ed2dd5d66f8ca1197cb1a1",
- strip_prefix = "zlib-1.2.11",
- build_file = clean_dep("//third_party:zlib.BUILD"),
- system_build_file = clean_dep("//third_party/systemlibs:zlib.BUILD"),
)
tf_http_archive(
name = "fft2d",
+ build_file = clean_dep("//third_party/fft2d:fft2d.BUILD"),
+ sha256 = "52bb637c70b971958ec79c9c8752b1df5ff0218a4db4510e60826e0cb79b5296",
urls = [
"https://mirror.bazel.build/www.kurims.kyoto-u.ac.jp/~ooura/fft.tgz",
"http://www.kurims.kyoto-u.ac.jp/~ooura/fft.tgz",
],
- sha256 = "52bb637c70b971958ec79c9c8752b1df5ff0218a4db4510e60826e0cb79b5296",
- build_file = clean_dep("//third_party/fft2d:fft2d.BUILD"),
)
tf_http_archive(
name = "snappy",
+ build_file = clean_dep("//third_party:snappy.BUILD"),
+ sha256 = "3dfa02e873ff51a11ee02b9ca391807f0c8ea0529a4924afa645fbf97163f9d4",
+ strip_prefix = "snappy-1.1.7",
+ system_build_file = clean_dep("//third_party/systemlibs:snappy.BUILD"),
urls = [
"https://mirror.bazel.build/github.com/google/snappy/archive/1.1.7.tar.gz",
"https://github.com/google/snappy/archive/1.1.7.tar.gz",
],
- sha256 = "3dfa02e873ff51a11ee02b9ca391807f0c8ea0529a4924afa645fbf97163f9d4",
- strip_prefix = "snappy-1.1.7",
- build_file = clean_dep("//third_party:snappy.BUILD"),
- system_build_file = clean_dep("//third_party/systemlibs:snappy.BUILD"),
)
tf_http_archive(
name = "nccl_archive",
+ build_file = clean_dep("//third_party:nccl/nccl_archive.BUILD"),
+ sha256 = "2ca86fb6179ecbff789cc67c836139c1bbc0324ed8c04643405a30bf26325176",
+ strip_prefix = "nccl-03d856977ecbaac87e598c0c4bafca96761b9ac7",
urls = [
"https://mirror.bazel.build/github.com/nvidia/nccl/archive/03d856977ecbaac87e598c0c4bafca96761b9ac7.tar.gz",
"https://github.com/nvidia/nccl/archive/03d856977ecbaac87e598c0c4bafca96761b9ac7.tar.gz",
],
- sha256 = "2ca86fb6179ecbff789cc67c836139c1bbc0324ed8c04643405a30bf26325176",
- strip_prefix = "nccl-03d856977ecbaac87e598c0c4bafca96761b9ac7",
- build_file = clean_dep("//third_party:nccl/nccl_archive.BUILD"),
)
tf_http_archive(
name = "kafka",
+ build_file = clean_dep("//third_party:kafka/BUILD"),
+ patch_file = clean_dep("//third_party/kafka:config.patch"),
+ sha256 = "cc6ebbcd0a826eec1b8ce1f625ffe71b53ef3290f8192b6cae38412a958f4fd3",
+ strip_prefix = "librdkafka-0.11.5",
urls = [
"https://mirror.bazel.build/github.com/edenhill/librdkafka/archive/v0.11.5.tar.gz",
"https://github.com/edenhill/librdkafka/archive/v0.11.5.tar.gz",
],
- sha256 = "cc6ebbcd0a826eec1b8ce1f625ffe71b53ef3290f8192b6cae38412a958f4fd3",
- strip_prefix = "librdkafka-0.11.5",
- build_file = clean_dep("//third_party:kafka/BUILD"),
- patch_file = clean_dep("//third_party/kafka:config.patch"),
)
tf_http_archive(
name = "aws",
+ build_file = clean_dep("//third_party:aws.BUILD"),
+ sha256 = "b888d8ce5fc10254c3dd6c9020c7764dd53cf39cf011249d0b4deda895de1b7c",
+ strip_prefix = "aws-sdk-cpp-1.3.15",
urls = [
"https://mirror.bazel.build/github.com/aws/aws-sdk-cpp/archive/1.3.15.tar.gz",
"https://github.com/aws/aws-sdk-cpp/archive/1.3.15.tar.gz",
],
- sha256 = "b888d8ce5fc10254c3dd6c9020c7764dd53cf39cf011249d0b4deda895de1b7c",
- strip_prefix = "aws-sdk-cpp-1.3.15",
- build_file = clean_dep("//third_party:aws.BUILD"),
)
java_import_external(
@@ -638,18 +642,6 @@
testonly_ = True,
)
- tf_http_archive(
- name = "jemalloc",
- urls = [
- "https://mirror.bazel.build/github.com/jemalloc/jemalloc/archive/4.4.0.tar.gz",
- "https://github.com/jemalloc/jemalloc/archive/4.4.0.tar.gz",
- ],
- sha256 = "3c8f25c02e806c3ce0ab5fb7da1817f89fc9732709024e2a81b6b82f7cc792a8",
- strip_prefix = "jemalloc-4.4.0",
- build_file = clean_dep("//third_party:jemalloc.BUILD"),
- system_build_file = clean_dep("//third_party/systemlibs:jemalloc.BUILD"),
- )
-
java_import_external(
name = "com_google_testing_compile",
jar_sha256 = "edc180fdcd9f740240da1a7a45673f46f59c5578d8cd3fbc912161f74b5aebb8",
@@ -696,196 +688,196 @@
tf_http_archive(
name = "com_google_pprof",
+ build_file = clean_dep("//third_party:pprof.BUILD"),
+ sha256 = "e0928ca4aa10ea1e0551e2d7ce4d1d7ea2d84b2abbdef082b0da84268791d0c4",
+ strip_prefix = "pprof-c0fb62ec88c411cc91194465e54db2632845b650",
urls = [
"https://mirror.bazel.build/github.com/google/pprof/archive/c0fb62ec88c411cc91194465e54db2632845b650.tar.gz",
"https://github.com/google/pprof/archive/c0fb62ec88c411cc91194465e54db2632845b650.tar.gz",
],
- sha256 = "e0928ca4aa10ea1e0551e2d7ce4d1d7ea2d84b2abbdef082b0da84268791d0c4",
- strip_prefix = "pprof-c0fb62ec88c411cc91194465e54db2632845b650",
- build_file = clean_dep("//third_party:pprof.BUILD"),
)
tf_http_archive(
name = "cub_archive",
+ build_file = clean_dep("//third_party:cub.BUILD"),
+ sha256 = "6bfa06ab52a650ae7ee6963143a0bbc667d6504822cbd9670369b598f18c58c3",
+ strip_prefix = "cub-1.8.0",
urls = [
"https://mirror.bazel.build/github.com/NVlabs/cub/archive/1.8.0.zip",
"https://github.com/NVlabs/cub/archive/1.8.0.zip",
],
- sha256 = "6bfa06ab52a650ae7ee6963143a0bbc667d6504822cbd9670369b598f18c58c3",
- strip_prefix = "cub-1.8.0",
- build_file = clean_dep("//third_party:cub.BUILD"),
)
tf_http_archive(
name = "cython",
+ build_file = clean_dep("//third_party:cython.BUILD"),
+ delete = ["BUILD.bazel"],
sha256 = "bccc9aa050ea02595b2440188813b936eaf345e85fb9692790cecfe095cf91aa",
+ strip_prefix = "cython-0.28.4",
+ system_build_file = clean_dep("//third_party/systemlibs:cython.BUILD"),
urls = [
"https://mirror.bazel.build/github.com/cython/cython/archive/0.28.4.tar.gz",
"https://github.com/cython/cython/archive/0.28.4.tar.gz",
],
- strip_prefix = "cython-0.28.4",
- build_file = clean_dep("//third_party:cython.BUILD"),
- delete = ["BUILD.bazel"],
- system_build_file = clean_dep("//third_party/systemlibs:cython.BUILD"),
)
tf_http_archive(
name = "bazel_toolchains",
+ sha256 = "3b604699685c5c65dd3f6f17425570a4b2f00ddba2f750db15acc72e55bb098b",
+ strip_prefix = "bazel-toolchains-37acf1841ab1475c98a152cb9e446460c8ae29e1",
urls = [
"https://mirror.bazel.build/github.com/bazelbuild/bazel-toolchains/archive/37acf1841ab1475c98a152cb9e446460c8ae29e1.tar.gz",
"https://github.com/bazelbuild/bazel-toolchains/archive/37acf1841ab1475c98a152cb9e446460c8ae29e1.tar.gz",
],
- strip_prefix = "bazel-toolchains-37acf1841ab1475c98a152cb9e446460c8ae29e1",
- sha256 = "3b604699685c5c65dd3f6f17425570a4b2f00ddba2f750db15acc72e55bb098b",
)
tf_http_archive(
name = "arm_neon_2_x86_sse",
+ build_file = clean_dep("//third_party:arm_neon_2_x86_sse.BUILD"),
sha256 = "c8d90aa4357f8079d427e87a6f4c493da1fa4140aee926c05902d7ec1533d9a5",
strip_prefix = "ARM_NEON_2_x86_SSE-0f77d9d182265259b135dad949230ecbf1a2633d",
urls = [
"https://mirror.bazel.build/github.com/intel/ARM_NEON_2_x86_SSE/archive/0f77d9d182265259b135dad949230ecbf1a2633d.tar.gz",
"https://github.com/intel/ARM_NEON_2_x86_SSE/archive/0f77d9d182265259b135dad949230ecbf1a2633d.tar.gz",
],
- build_file = clean_dep("//third_party:arm_neon_2_x86_sse.BUILD"),
)
tf_http_archive(
name = "double_conversion",
+ build_file = clean_dep("//third_party:double_conversion.BUILD"),
+ sha256 = "2f7fbffac0d98d201ad0586f686034371a6d152ca67508ab611adc2386ad30de",
+ strip_prefix = "double-conversion-3992066a95b823efc8ccc1baf82a1cfc73f6e9b8",
+ system_build_file = clean_dep("//third_party/systemlibs:double_conversion.BUILD"),
urls = [
"https://mirror.bazel.build/github.com/google/double-conversion/archive/3992066a95b823efc8ccc1baf82a1cfc73f6e9b8.zip",
"https://github.com/google/double-conversion/archive/3992066a95b823efc8ccc1baf82a1cfc73f6e9b8.zip",
],
- sha256 = "2f7fbffac0d98d201ad0586f686034371a6d152ca67508ab611adc2386ad30de",
- strip_prefix = "double-conversion-3992066a95b823efc8ccc1baf82a1cfc73f6e9b8",
- build_file = clean_dep("//third_party:double_conversion.BUILD"),
- system_build_file = clean_dep("//third_party/systemlibs:double_conversion.BUILD"),
)
tf_http_archive(
name = "tflite_mobilenet",
+ build_file = clean_dep("//third_party:tflite_mobilenet.BUILD"),
sha256 = "23f814d1c076bdf03715dfb6cab3713aa4fbdf040fd5448c43196bd2e97a4c1b",
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",
],
- build_file = clean_dep("//third_party:tflite_mobilenet.BUILD"),
)
tf_http_archive(
name = "tflite_mobilenet_ssd",
+ build_file = str(Label("//third_party:tflite_mobilenet.BUILD")),
sha256 = "767057f2837a46d97882734b03428e8dd640b93236052b312b2f0e45613c1cf0",
urls = [
"https://mirror.bazel.build/storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_ssd_tflite_v1.zip",
"https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_ssd_tflite_v1.zip",
],
- build_file = str(Label("//third_party:tflite_mobilenet.BUILD")),
)
tf_http_archive(
name = "tflite_mobilenet_ssd_quant",
+ build_file = str(Label("//third_party:tflite_mobilenet.BUILD")),
sha256 = "a809cd290b4d6a2e8a9d5dad076e0bd695b8091974e0eed1052b480b2f21b6dc",
urls = [
"https://mirror.bazel.build/storage.googleapis.com/download.tensorflow.org/models/tflite/coco_ssd_mobilenet_v1_0.75_quant_2018_06_29.zip",
"https://storage.googleapis.com/download.tensorflow.org/models/tflite/coco_ssd_mobilenet_v1_0.75_quant_2018_06_29.zip",
],
- build_file = str(Label("//third_party:tflite_mobilenet.BUILD")),
)
tf_http_archive(
name = "tflite_mobilenet_ssd_quant_protobuf",
+ build_file = str(Label("//third_party:tflite_mobilenet.BUILD")),
sha256 = "09280972c5777f1aa775ef67cb4ac5d5ed21970acd8535aeca62450ef14f0d79",
+ strip_prefix = "ssd_mobilenet_v1_quantized_300x300_coco14_sync_2018_07_18",
urls = [
"https://mirror.bazel.build/storage.googleapis.com/download.tensorflow.org/models/object_detection/ssd_mobilenet_v1_quantized_300x300_coco14_sync_2018_07_18.tar.gz",
"http://storage.googleapis.com/download.tensorflow.org/models/object_detection/ssd_mobilenet_v1_quantized_300x300_coco14_sync_2018_07_18.tar.gz",
],
- strip_prefix = "ssd_mobilenet_v1_quantized_300x300_coco14_sync_2018_07_18",
- build_file = str(Label("//third_party:tflite_mobilenet.BUILD")),
)
tf_http_archive(
name = "tflite_conv_actions_frozen",
+ build_file = str(Label("//third_party:tflite_mobilenet.BUILD")),
sha256 = "d947b38cba389b5e2d0bfc3ea6cc49c784e187b41a071387b3742d1acac7691e",
urls = [
"https://mirror.bazel.build/storage.googleapis.com/download.tensorflow.org/models/tflite/conv_actions_tflite.zip",
"https://storage.googleapis.com/download.tensorflow.org/models/tflite/conv_actions_tflite.zip",
],
- build_file = str(Label("//third_party:tflite_mobilenet.BUILD")),
)
tf_http_archive(
name = "tflite_smartreply",
+ build_file = clean_dep("//third_party:tflite_smartreply.BUILD"),
sha256 = "8980151b85a87a9c1a3bb1ed4748119e4a85abd3cb5744d83da4d4bd0fbeef7c",
urls = [
"https://mirror.bazel.build/storage.googleapis.com/download.tensorflow.org/models/tflite/smartreply_1.0_2017_11_01.zip",
"https://storage.googleapis.com/download.tensorflow.org/models/tflite/smartreply_1.0_2017_11_01.zip",
],
- build_file = clean_dep("//third_party:tflite_smartreply.BUILD"),
)
tf_http_archive(
name = "tflite_ovic_testdata",
+ build_file = clean_dep("//third_party:tflite_ovic_testdata.BUILD"),
sha256 = "a9a705d8d519220178e2e65d383fdb21da37fdb31d1e909b0a1acdac46479e9c",
+ strip_prefix = "ovic",
urls = [
"https://mirror.bazel.build/storage.googleapis.com/download.tensorflow.org/data/ovic.zip",
"https://storage.googleapis.com/download.tensorflow.org/data/ovic.zip",
],
- build_file = clean_dep("//third_party:tflite_ovic_testdata.BUILD"),
- strip_prefix = "ovic",
)
tf_http_archive(
name = "build_bazel_rules_android",
sha256 = "cd06d15dd8bb59926e4d65f9003bfc20f9da4b2519985c27e190cddc8b7a7806",
+ strip_prefix = "rules_android-0.1.1",
urls = [
"https://mirror.bazel.build/github.com/bazelbuild/rules_android/archive/v0.1.1.zip",
"https://github.com/bazelbuild/rules_android/archive/v0.1.1.zip",
],
- strip_prefix = "rules_android-0.1.1",
)
tf_http_archive(
name = "tbb",
+ build_file = clean_dep("//third_party/ngraph:tbb.BUILD"),
+ sha256 = "724686f90bcda78f13b76f297d964008737ccd6399328143c1c0093e73ae6a13",
+ strip_prefix = "tbb-tbb_2018",
urls = [
"https://mirror.bazel.build/github.com/01org/tbb/archive/tbb_2018.zip",
"https://github.com/01org/tbb/archive/tbb_2018.zip",
],
- sha256 = "724686f90bcda78f13b76f297d964008737ccd6399328143c1c0093e73ae6a13",
- strip_prefix = "tbb-tbb_2018",
- build_file = clean_dep("//third_party/ngraph:tbb.BUILD"),
)
tf_http_archive(
name = "ngraph",
- urls = [
- "https://mirror.bazel.build/github.com/NervanaSystems/ngraph/archive/v0.8.0.tar.gz",
- "https://github.com/NervanaSystems/ngraph/archive/v0.8.0.tar.gz",
- ],
- sha256 = "a8cf3ef2d0e6d31b54eb33f6a9e795f562195ce5c2a857e729ca9c35241cc45c",
- strip_prefix = "ngraph-0.8.0",
build_file = clean_dep("//third_party/ngraph:ngraph.BUILD"),
+ sha256 = "bf9dcc88e5c66021e3aac80491a231711211540d613bf9b6bd28db3f5bb86b62",
+ strip_prefix = "ngraph-0.8.1",
+ urls = [
+ "https://mirror.bazel.build/github.com/NervanaSystems/ngraph/archive/v0.8.1.tar.gz",
+ "https://github.com/NervanaSystems/ngraph/archive/v0.8.1.tar.gz",
+ ],
)
tf_http_archive(
name = "nlohmann_json_lib",
+ build_file = clean_dep("//third_party/ngraph:nlohmann_json.BUILD"),
+ sha256 = "9f3549824af3ca7e9707a2503959886362801fb4926b869789d6929098a79e47",
+ strip_prefix = "json-3.1.1",
urls = [
"https://mirror.bazel.build/github.com/nlohmann/json/archive/v3.1.1.tar.gz",
"https://github.com/nlohmann/json/archive/v3.1.1.tar.gz",
],
- sha256 = "9f3549824af3ca7e9707a2503959886362801fb4926b869789d6929098a79e47",
- strip_prefix = "json-3.1.1",
- build_file = clean_dep("//third_party/ngraph:nlohmann_json.BUILD"),
)
tf_http_archive(
name = "ngraph_tf",
- urls = [
- "https://mirror.bazel.build/github.com/NervanaSystems/ngraph-tf/archive/v0.6.0.tar.gz",
- "https://github.com/NervanaSystems/ngraph-tf/archive/v0.6.0.tar.gz",
- ],
- sha256 = "1f49391c02bef24872e9f85591e60e0e7eef12a337db71390444118049fe451f",
- strip_prefix = "ngraph-tf-0.6.0",
build_file = clean_dep("//third_party/ngraph:ngraph_tf.BUILD"),
+ sha256 = "402f84c748c113780a60f35f39aab118435285543aee4900d712b76fbf8a21ee",
+ strip_prefix = "ngraph-tf-0.6.1",
+ urls = [
+ "https://mirror.bazel.build/github.com/NervanaSystems/ngraph-tf/archive/v0.6.1.tar.gz",
+ "https://github.com/NervanaSystems/ngraph-tf/archive/v0.6.1.tar.gz",
+ ],
)
##############################################################################
diff --git a/third_party/gpus/crosstool/BUILD.tpl b/third_party/gpus/crosstool/BUILD.tpl
index f638756..c8812fa 100644
--- a/third_party/gpus/crosstool/BUILD.tpl
+++ b/third_party/gpus/crosstool/BUILD.tpl
@@ -2,6 +2,20 @@
package(default_visibility = ["//visibility:public"])
+toolchain(
+ name = "toolchain-linux-x86_64",
+ exec_compatible_with = [
+ "@bazel_tools//platforms:linux",
+ "@bazel_tools//platforms:x86_64",
+ ],
+ target_compatible_with = [
+ "@bazel_tools//platforms:linux",
+ "@bazel_tools//platforms:x86_64",
+ ],
+ toolchain = ":cc-compiler-local",
+ toolchain_type = "@bazel_tools//tools/cpp:toolchain_type",
+)
+
cc_toolchain_suite(
name = "toolchain",
toolchains = {
diff --git a/third_party/gpus/crosstool/CROSSTOOL_hipcc.tpl b/third_party/gpus/crosstool/CROSSTOOL_hipcc.tpl
new file mode 100644
index 0000000..0e175b3
--- /dev/null
+++ b/third_party/gpus/crosstool/CROSSTOOL_hipcc.tpl
@@ -0,0 +1,158 @@
+major_version: "local"
+minor_version: ""
+default_target_cpu: "same_as_host"
+
+default_toolchain {
+ cpu: "k8"
+ toolchain_identifier: "local_linux"
+}
+default_toolchain {
+ cpu: "piii"
+ toolchain_identifier: "local_linux"
+}
+default_toolchain {
+ cpu: "arm"
+ toolchain_identifier: "local_linux"
+}
+default_toolchain {
+ cpu: "ppc"
+ toolchain_identifier: "local_linux"
+}
+
+toolchain {
+ abi_version: "local"
+ abi_libc_version: "local"
+ builtin_sysroot: ""
+ compiler: "compiler"
+ host_system_name: "local"
+ needsPic: true
+ supports_gold_linker: false
+ supports_incremental_linker: false
+ supports_fission: false
+ supports_interface_shared_objects: false
+ supports_normalizing_ar: false
+ supports_start_end_lib: false
+ supports_thin_archives: false
+ target_libc: "local"
+ target_cpu: "local"
+ target_system_name: "local"
+ toolchain_identifier: "local_linux"
+
+ tool_path { name: "ar" path: "/usr/bin/ar" }
+ tool_path { name: "compat-ld" path: "/usr/bin/ld" }
+ tool_path { name: "cpp" path: "/usr/bin/cpp" }
+ tool_path { name: "dwp" path: "/usr/bin/dwp" }
+ # As part of the TensorFlow release, we place some ROCm-related compilation
+ # files in @local_config_rocm//crosstool/clang/bin, and this relative
+ # path, combined with the rest of our Bazel configuration causes our
+ # compilation to use those files.
+ tool_path { name: "gcc" path: "clang/bin/crosstool_wrapper_driver_rocm" }
+ # Use "-std=c++11" for hipcc. For consistency, force both the host compiler
+ # and the device compiler to use "-std=c++11".
+ cxx_flag: "-std=c++11"
+ linker_flag: "-Wl,-no-as-needed"
+ linker_flag: "-lstdc++"
+ #linker_flag: "-B/usr/bin/"
+ linker_flag: "-B/opt/rocm/hcc/compiler/bin"
+
+%{host_compiler_includes}
+ tool_path { name: "gcov" path: "/usr/bin/gcov" }
+
+ # C(++) compiles invoke the compiler (as that is the one knowing where
+ # to find libraries), but we provide LD so other rules can invoke the linker.
+ tool_path { name: "ld" path: "/usr/bin/ld" }
+
+ tool_path { name: "nm" path: "/usr/bin/nm" }
+ tool_path { name: "objcopy" path: "/usr/bin/objcopy" }
+ objcopy_embed_flag: "-I"
+ objcopy_embed_flag: "binary"
+ tool_path { name: "objdump" path: "/usr/bin/objdump" }
+ tool_path { name: "strip" path: "/usr/bin/strip" }
+
+ # Anticipated future default.
+ unfiltered_cxx_flag: "-no-canonical-prefixes"
+
+ # Make C++ compilation deterministic. Use linkstamping instead of these
+ # compiler symbols.
+ unfiltered_cxx_flag: "-Wno-builtin-macro-redefined"
+ unfiltered_cxx_flag: "-D__DATE__=\"redacted\""
+ unfiltered_cxx_flag: "-D__TIMESTAMP__=\"redacted\""
+ unfiltered_cxx_flag: "-D__TIME__=\"redacted\""
+ unfiltered_cxx_flag: "-D__HIP_PLATFORM_HCC__"
+ # The macro EIGEN_USE_HIP is used to tell Eigen to use the HIP platform headers
+ # It needs to be always set when compiling Eigen headers
+ # (irrespective of whether the source file is being compiled via HIPCC)
+ # so adding -DEIGEN_USE_HIP as a default CXX flag here
+ unfiltered_cxx_flag: "-DEIGEN_USE_HIP"
+
+
+ # Security hardening on by default.
+ # Conservative choice; -D_FORTIFY_SOURCE=2 may be unsafe in some cases.
+ # We need to undef it before redefining it as some distributions now have
+ # it enabled by default.
+ #compiler_flag: "-U_FORTIFY_SOURCE"
+ #compiler_flag: "-D_FORTIFY_SOURCE=1"
+ #compiler_flag: "-fstack-protector"
+ #compiler_flag: "-fPIE"
+ #linker_flag: "-pie"
+ #linker_flag: "-Wl,-z,relro,-z,now"
+
+ # Enable coloring even if there's no attached terminal. Bazel removes the
+ # escape sequences if --nocolor is specified. This isn't supported by gcc
+ # on Ubuntu 14.04.
+ # compiler_flag: "-fcolor-diagnostics"
+
+ # All warnings are enabled. Maybe enable -Werror as well?
+ compiler_flag: "-Wall"
+ # Enable a few more warnings that aren't part of -Wall.
+ compiler_flag: "-Wunused-but-set-parameter"
+ # But disable some that are problematic.
+ compiler_flag: "-Wno-free-nonheap-object" # has false positives
+
+ # Keep stack frames for debugging, even in opt mode.
+ compiler_flag: "-fno-omit-frame-pointer"
+
+ # Anticipated future default.
+ linker_flag: "-no-canonical-prefixes"
+ unfiltered_cxx_flag: "-fno-canonical-system-headers"
+ # Have gcc return the exit code from ld.
+ linker_flag: "-pass-exit-codes"
+ # Stamp the binary with a unique identifier.
+ linker_flag: "-Wl,--build-id=md5"
+ linker_flag: "-Wl,--hash-style=gnu"
+ # Gold linker only? Can we enable this by default?
+ # linker_flag: "-Wl,--warn-execstack"
+ # linker_flag: "-Wl,--detect-odr-violations"
+
+ # Include directory for ROCm headers.
+%{rocm_include_path}
+
+ compilation_mode_flags {
+ mode: DBG
+ # Enable debug symbols.
+ compiler_flag: "-g"
+ }
+ compilation_mode_flags {
+ mode: OPT
+
+ # No debug symbols.
+ # Maybe we should enable https://gcc.gnu.org/wiki/DebugFission for opt or
+ # even generally? However, that can't happen here, as it requires special
+ # handling in Bazel.
+ compiler_flag: "-g0"
+
+ # Conservative choice for -O
+ # -O3 can increase binary size and even slow down the resulting binaries.
+ # Profile first and / or use FDO if you need better performance than this.
+ compiler_flag: "-O2"
+
+ # Disable assertions
+ compiler_flag: "-DNDEBUG"
+
+ # Removal of unused code and data at link time (can this increase binary size in some cases?).
+ compiler_flag: "-ffunction-sections"
+ compiler_flag: "-fdata-sections"
+ linker_flag: "-Wl,--gc-sections"
+ }
+ linking_mode_flags { mode: DYNAMIC }
+}
diff --git a/third_party/gpus/crosstool/clang/bin/crosstool_wrapper_driver_rocm.tpl b/third_party/gpus/crosstool/clang/bin/crosstool_wrapper_driver_rocm.tpl
new file mode 100755
index 0000000..8242380
--- /dev/null
+++ b/third_party/gpus/crosstool/clang/bin/crosstool_wrapper_driver_rocm.tpl
@@ -0,0 +1,241 @@
+#!/usr/bin/env python
+"""Crosstool wrapper for compiling ROCm programs.
+
+SYNOPSIS:
+ crosstool_wrapper_driver_rocm [options passed in by cc_library()
+ or cc_binary() rule]
+
+DESCRIPTION:
+ This script is expected to be called by the cc_library() or cc_binary() bazel
+ rules. When the option "-x rocm" is present in the list of arguments passed
+ to this script, it invokes the hipcc compiler. Most arguments are passed
+ as is as a string to --compiler-options of hipcc. When "-x rocm" is not
+ present, this wrapper invokes gcc with the input arguments as is.
+"""
+
+from __future__ import print_function
+
+__author__ = 'whchung@gmail.com (Wen-Heng (Jack) Chung)'
+
+from argparse import ArgumentParser
+import os
+import subprocess
+import re
+import sys
+import pipes
+
+# Template values set by rocm_configure.bzl.
+CPU_COMPILER = ('%{cpu_compiler}')
+GCC_HOST_COMPILER_PATH = ('%{gcc_host_compiler_path}')
+
+HIPCC_PATH = '%{hipcc_path}'
+PREFIX_DIR = os.path.dirname(GCC_HOST_COMPILER_PATH)
+
+def Log(s):
+ print('gpus/crosstool: {0}'.format(s))
+
+
+def GetOptionValue(argv, option):
+ """Extract the list of values for option from the argv list.
+
+ Args:
+ argv: A list of strings, possibly the argv passed to main().
+ option: The option whose value to extract, without the leading '-'.
+
+ Returns:
+ A list of values, either directly following the option,
+ (eg., -opt val1 val2) or values collected from multiple occurrences of
+ the option (eg., -opt val1 -opt val2).
+ """
+
+ parser = ArgumentParser()
+ parser.add_argument('-' + option, nargs='*', action='append')
+ args, _ = parser.parse_known_args(argv)
+ if not args or not vars(args)[option]:
+ return []
+ else:
+ return sum(vars(args)[option], [])
+
+
+def GetHostCompilerOptions(argv):
+ """Collect the -isystem, -iquote, and --sysroot option values from argv.
+
+ Args:
+ argv: A list of strings, possibly the argv passed to main().
+
+ Returns:
+ The string that can be used as the --compiler-options to hipcc.
+ """
+
+ parser = ArgumentParser()
+ parser.add_argument('-isystem', nargs='*', action='append')
+ parser.add_argument('-iquote', nargs='*', action='append')
+ parser.add_argument('--sysroot', nargs=1)
+ parser.add_argument('-g', nargs='*', action='append')
+ parser.add_argument('-fno-canonical-system-headers', action='store_true')
+
+ args, _ = parser.parse_known_args(argv)
+
+ opts = ''
+
+ if args.isystem:
+ opts += ' -isystem ' + ' -isystem '.join(sum(args.isystem, []))
+ if args.iquote:
+ opts += ' -iquote ' + ' -iquote '.join(sum(args.iquote, []))
+ if args.g:
+ opts += ' -g' + ' -g'.join(sum(args.g, []))
+ #if args.fno_canonical_system_headers:
+ # opts += ' -fno-canonical-system-headers'
+ if args.sysroot:
+ opts += ' --sysroot ' + args.sysroot[0]
+
+ return opts
+
+def GetHipccOptions(argv):
+ """Collect the -hipcc_options values from argv.
+
+ Args:
+ argv: A list of strings, possibly the argv passed to main().
+
+ Returns:
+ The string that can be passed directly to hipcc.
+ """
+
+ parser = ArgumentParser()
+ parser.add_argument('-hipcc_options', nargs='*', action='append')
+
+ args, _ = parser.parse_known_args(argv)
+
+ if args.hipcc_options:
+ options = _update_options(sum(args.hipcc_options, []))
+ return ' '.join(['--'+a for a in options])
+ return ''
+
+
+def InvokeHipcc(argv, log=False):
+ """Call hipcc with arguments assembled from argv.
+
+ Args:
+ argv: A list of strings, possibly the argv passed to main().
+ log: True if logging is requested.
+
+ Returns:
+ The return value of calling os.system('hipcc ' + args)
+ """
+
+ host_compiler_options = GetHostCompilerOptions(argv)
+ hipcc_compiler_options = GetHipccOptions(argv)
+ opt_option = GetOptionValue(argv, 'O')
+ m_options = GetOptionValue(argv, 'm')
+ m_options = ''.join([' -m' + m for m in m_options if m in ['32', '64']])
+ include_options = GetOptionValue(argv, 'I')
+ out_file = GetOptionValue(argv, 'o')
+ depfiles = GetOptionValue(argv, 'MF')
+ defines = GetOptionValue(argv, 'D')
+ defines = ''.join([' -D' + define for define in defines])
+ undefines = GetOptionValue(argv, 'U')
+ undefines = ''.join([' -U' + define for define in undefines])
+ std_options = GetOptionValue(argv, 'std')
+ hipcc_allowed_std_options = ["c++11"]
+ std_options = ''.join([' -std=' + define
+ for define in std_options if define in hipcc_allowed_std_options])
+
+ # The list of source files get passed after the -c option. I don't know of
+ # any other reliable way to just get the list of source files to be compiled.
+ src_files = GetOptionValue(argv, 'c')
+
+ if len(src_files) == 0:
+ return 1
+ if len(out_file) != 1:
+ return 1
+
+ opt = (' -O2' if (len(opt_option) > 0 and int(opt_option[0]) > 0)
+ else ' -g')
+
+ includes = (' -I ' + ' -I '.join(include_options)
+ if len(include_options) > 0
+ else '')
+
+ # Unfortunately, there are other options that have -c prefix too.
+ # So allowing only those look like C/C++ files.
+ src_files = [f for f in src_files if
+ re.search('\.cpp$|\.cc$|\.c$|\.cxx$|\.C$', f)]
+ srcs = ' '.join(src_files)
+ out = ' -o ' + out_file[0]
+
+ hipccopts = ' '
+ hipccopts += ' ' + hipcc_compiler_options
+ hipccopts += undefines
+ hipccopts += defines
+ hipccopts += std_options
+ hipccopts += m_options
+
+ if depfiles:
+ # Generate the dependency file
+ depfile = depfiles[0]
+ cmd = (HIPCC_PATH + ' ' + hipccopts +
+ host_compiler_options +
+ ' ' + GCC_HOST_COMPILER_PATH +
+ ' -I .' + includes + ' ' + srcs + ' -M -o ' + depfile)
+ if log: Log(cmd)
+ exit_status = os.system(cmd)
+ if exit_status != 0:
+ return exit_status
+
+ cmd = (HIPCC_PATH + ' ' + hipccopts +
+ host_compiler_options + ' -fPIC' +
+ ' ' + GCC_HOST_COMPILER_PATH +
+ ' -I .' + opt + includes + ' -c ' + srcs + out)
+
+ # TODO(zhengxq): for some reason, 'gcc' needs this help to find 'as'.
+ # Need to investigate and fix.
+ cmd = 'PATH=' + PREFIX_DIR + ':$PATH ' + cmd
+ if log: Log(cmd)
+ return os.system(cmd)
+
+
+def main():
+ # ignore PWD env var
+ os.environ['PWD']=''
+
+ parser = ArgumentParser()
+ parser.add_argument('-x', nargs=1)
+ parser.add_argument('--rocm_log', action='store_true')
+ parser.add_argument('-pass-exit-codes', action='store_true')
+ args, leftover = parser.parse_known_args(sys.argv[1:])
+
+ if args.x and args.x[0] == 'rocm':
+ if args.rocm_log: Log('-x rocm')
+ leftover = [pipes.quote(s) for s in leftover]
+ if args.rocm_log: Log('using hipcc')
+ return InvokeHipcc(leftover, log=args.rocm_log)
+
+ # XXX use hipcc to link
+ if args.pass_exit_codes:
+ gpu_compiler_flags = [flag for flag in sys.argv[1:]
+ if not flag.startswith(('-pass-exit-codes'))]
+
+ # special handling for $ORIGIN
+ # - guard every argument with ''
+ modified_gpu_compiler_flags = []
+ for flag in gpu_compiler_flags:
+ modified_gpu_compiler_flags.append("'" + flag + "'")
+
+ if args.rocm_log: Log('Link with hipcc: %s' % (' '.join([HIPCC_PATH] + modified_gpu_compiler_flags)))
+ return subprocess.call([HIPCC_PATH] + modified_gpu_compiler_flags)
+
+ # Strip our flags before passing through to the CPU compiler for files which
+ # are not -x rocm. We can't just pass 'leftover' because it also strips -x.
+ # We not only want to pass -x to the CPU compiler, but also keep it in its
+ # relative location in the argv list (the compiler is actually sensitive to
+ # this).
+ cpu_compiler_flags = [flag for flag in sys.argv[1:]
+ if not flag.startswith(('--rocm_log'))]
+
+ # XXX: SE codes need to be built with gcc, but need this macro defined
+ cpu_compiler_flags.append("-D__HIP_PLATFORM_HCC__")
+
+ return subprocess.call([CPU_COMPILER] + cpu_compiler_flags)
+
+if __name__ == '__main__':
+ sys.exit(main())
diff --git a/third_party/gpus/cuda_configure.bzl b/third_party/gpus/cuda_configure.bzl
index f5fdd3a..69f4599 100644
--- a/third_party/gpus/cuda_configure.bzl
+++ b/third_party/gpus/cuda_configure.bzl
@@ -1107,8 +1107,8 @@
# $(@D) will include the full path to the file.
dest = "$(@D)/" + dest_dir + dest_files[i] if len(dest_files) != 1 else "$(@D)/" + dest_files[i]
- # On Windows, symlink is not supported, so we just copy all the files.
- cmd = "cp -f" if _is_windows(repository_ctx) else "ln -s"
+ # Copy the headers to create a sandboxable setup.
+ cmd = "cp -f"
command.append(cmd + ' "%s" "%s"' % (src_files[i], dest))
outs.append(' "' + dest_dir + dest_files[i] + '",')
genrule = _genrule(
@@ -1334,27 +1334,14 @@
cuda_defines["%{host_compiler_path}"] = "clang/bin/crosstool_wrapper_driver_is_not_gcc"
cuda_defines["%{host_compiler_warnings}"] = ""
- # TODO(klimek): We currently need to inject "/" as builtin directory path
- # to disable bazel's dependency checks.
- # The problem is that:
- # - the python rules symlink the python headers into the bazel root
- # - the rules use 'includes' in the BUILD file to redirect includes of the
- # python headers through those paths
- # - bazel currently uses -isystem for include paths specified via 'includes'
- # - gcc follows symlinks when resolving files via -isystem paths, and puts
- # the resolved paths into the .d file, which makes the dependency check
- # fail for bazel
- # There are multiple possible ways to solve this:
- # 1. make bazel not use -isystem for paths specified via 'includes'
- # 2. cp the headers instead of symlinking them
- #
- # Once this is fixed, the right builtin directory path is:
- # (host_compiler_includes +
- # "\n cxx_builtin_include_directory: \"%s\"" % cuda_include_path)
- # The cuda directory needs to be passed, as there is currently no rule
- # providing the cuda headers in the same way the python headers are
- # provided.
- cuda_defines["%{host_compiler_includes}"] = "\n cxx_builtin_include_directory: \"/\""
+ # nvcc has the system include paths built in and will automatically
+ # search them; we cannot work around that, so we add the relevant cuda
+ # system paths to the allowed compiler specific include paths.
+ cuda_defines["%{host_compiler_includes}"] = (
+ host_compiler_includes + "\n" +
+ _cuda_include_path(repository_ctx, cuda_config) +
+ "\n cxx_builtin_include_directory: \"%s\"" % cupti_header_dir +
+ "\n cxx_builtin_include_directory: \"%s\"" % cudnn_header_dir)
nvcc_path = str(repository_ctx.path("%s/bin/nvcc%s" %
(
cuda_config.cuda_toolkit_path,
diff --git a/tensorflow/docs_src/__init__.py b/third_party/gpus/rocm/BUILD
similarity index 100%
rename from tensorflow/docs_src/__init__.py
rename to third_party/gpus/rocm/BUILD
diff --git a/third_party/gpus/rocm/BUILD.tpl b/third_party/gpus/rocm/BUILD.tpl
new file mode 100644
index 0000000..8258bb3
--- /dev/null
+++ b/third_party/gpus/rocm/BUILD.tpl
@@ -0,0 +1,99 @@
+licenses(["restricted"]) # MPL2, portions GPL v3, LGPL v3, BSD-like
+
+package(default_visibility = ["//visibility:public"])
+
+config_setting(
+ name = "using_hipcc",
+ values = {
+ "define": "using_rocm_hipcc=true",
+ },
+)
+
+cc_library(
+ name = "rocm_headers",
+ hdrs = [
+ "rocm/rocm_config.h",
+ %{rocm_headers}
+ ],
+ includes = [
+ ".",
+ "rocm/include",
+ ],
+ visibility = ["//visibility:public"],
+)
+
+cc_library(
+ name = "hip",
+ srcs = ["rocm/lib/%{hip_lib}"],
+ data = ["rocm/lib/%{hip_lib}"],
+ includes = [
+ ".",
+ "rocm/include",
+ ],
+ linkstatic = 1,
+ visibility = ["//visibility:public"],
+)
+
+cc_library(
+ name = "rocblas",
+ srcs = ["rocm/lib/%{rocblas_lib}"],
+ data = ["rocm/lib/%{rocblas_lib}"],
+ includes = [
+ ".",
+ "rocm/include",
+ ],
+ linkstatic = 1,
+ visibility = ["//visibility:public"],
+)
+
+cc_library(
+ name = "rocfft",
+ srcs = ["rocm/lib/%{rocfft_lib}"],
+ data = ["rocm/lib/%{rocfft_lib}"],
+ includes = [
+ ".",
+ "rocm/include",
+ ],
+ linkstatic = 1,
+ visibility = ["//visibility:public"],
+)
+
+cc_library(
+ name = "hiprand",
+ srcs = ["rocm/lib/%{hiprand_lib}"],
+ data = ["rocm/lib/%{hiprand_lib}"],
+ includes = [
+ ".",
+ "rocm/include",
+ "rocm/include/rocrand",
+ ],
+ linkstatic = 1,
+ visibility = ["//visibility:public"],
+)
+
+cc_library(
+ name = "miopen",
+ srcs = ["rocm/lib/%{miopen_lib}"],
+ data = ["rocm/lib/%{miopen_lib}"],
+ includes = [
+ ".",
+ "rocm/include",
+ ],
+ linkstatic = 1,
+ visibility = ["//visibility:public"],
+)
+
+cc_library(
+ name = "rocm",
+ visibility = ["//visibility:public"],
+ deps = [
+ ":rocm_headers",
+ ":hip",
+ ":rocblas",
+ ":rocfft",
+ ":hiprand",
+ ":miopen",
+ ],
+)
+
+%{rocm_include_genrules}
diff --git a/third_party/gpus/rocm/build_defs.bzl.tpl b/third_party/gpus/rocm/build_defs.bzl.tpl
new file mode 100644
index 0000000..08c59f9
--- /dev/null
+++ b/third_party/gpus/rocm/build_defs.bzl.tpl
@@ -0,0 +1,45 @@
+# Macros for building ROCm code.
+def if_rocm(if_true, if_false = []):
+ """Shorthand for select()'ing on whether we're building with ROCm.
+
+ Returns a select statement which evaluates to if_true if we're building
+ with ROCm enabled. Otherwise, the select statement evaluates to if_false.
+
+ """
+ return select({
+ "@local_config_rocm//rocm:using_hipcc": if_true,
+ "//conditions:default": if_false
+ })
+
+
+def rocm_default_copts():
+ """Default options for all ROCm compilations."""
+ return if_rocm(["-x", "rocm"] + %{rocm_extra_copts})
+
+def rocm_copts(opts = []):
+ """Gets the appropriate set of copts for (maybe) ROCm compilation.
+
+ If we're doing ROCm compilation, returns copts for our particular ROCm
+ compiler. If we're not doing ROCm compilation, returns an empty list.
+
+ """
+ return rocm_default_copts() + select({
+ "//conditions:default": [],
+ "@local_config_rocm//rocm:using_hipcc": ([
+ "",
+ ]),
+ }) + if_rocm_is_configured(opts)
+
+def rocm_is_configured():
+ """Returns true if ROCm was enabled during the configure process."""
+ return %{rocm_is_configured}
+
+def if_rocm_is_configured(x):
+ """Tests if the ROCm was enabled during the configure process.
+
+ Unlike if_rocm(), this does not require that we are building with
+ --config=rocm. Used to allow non-ROCm code to depend on ROCm libraries.
+ """
+ if rocm_is_configured():
+ return x
+ return []
diff --git a/third_party/gpus/rocm/rocm_config.h.tpl b/third_party/gpus/rocm/rocm_config.h.tpl
new file mode 100644
index 0000000..c5f25a8
--- /dev/null
+++ b/third_party/gpus/rocm/rocm_config.h.tpl
@@ -0,0 +1,21 @@
+/* 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.
+==============================================================================*/
+
+#ifndef ROCM_ROCM_CONFIG_H_
+#define ROCM_ROCM_CONFIG_H_
+
+#define TF_ROCM_TOOLKIT_PATH "/opt/rocm"
+
+#endif // ROCM_ROCM_CONFIG_H_
diff --git a/third_party/gpus/rocm_configure.bzl b/third_party/gpus/rocm_configure.bzl
new file mode 100644
index 0000000..9108639
--- /dev/null
+++ b/third_party/gpus/rocm_configure.bzl
@@ -0,0 +1,784 @@
+# -*- Python -*-
+"""Repository rule for ROCm autoconfiguration.
+
+`rocm_configure` depends on the following environment variables:
+
+ * `TF_NEED_ROCM`: Whether to enable building with ROCm.
+ * `GCC_HOST_COMPILER_PATH`: The GCC host compiler path
+ * `ROCM_TOOLKIT_PATH`: The path to the ROCm toolkit. Default is
+ `/opt/rocm`.
+ * `TF_ROCM_VERSION`: The version of the ROCm toolkit. If this is blank, then
+ use the system default.
+ * `TF_MIOPEN_VERSION`: The version of the MIOpen library.
+ * `TF_ROCM_AMDGPU_TARGETS`: The AMDGPU targets. Default is
+ `gfx803,gfx900`.
+"""
+
+_GCC_HOST_COMPILER_PATH = "GCC_HOST_COMPILER_PATH"
+_ROCM_TOOLKIT_PATH = "ROCM_TOOLKIT_PATH"
+_TF_ROCM_VERSION = "TF_ROCM_VERSION"
+_TF_MIOPEN_VERSION = "TF_MIOPEN_VERSION"
+_TF_ROCM_AMDGPU_TARGETS = "TF_ROCM_AMDGPU_TARGETS"
+_TF_ROCM_CONFIG_REPO = "TF_ROCM_CONFIG_REPO"
+
+_DEFAULT_ROCM_VERSION = ""
+_DEFAULT_MIOPEN_VERSION = ""
+_DEFAULT_ROCM_TOOLKIT_PATH = "/opt/rocm"
+_DEFAULT_ROCM_AMDGPU_TARGETS = ["gfx803", "gfx900"]
+
+def find_cc(repository_ctx):
+ """Find the C++ compiler."""
+
+ # Return a dummy value for GCC detection here to avoid error
+ target_cc_name = "gcc"
+ cc_path_envvar = _GCC_HOST_COMPILER_PATH
+ cc_name = target_cc_name
+
+ if cc_path_envvar in repository_ctx.os.environ:
+ cc_name_from_env = repository_ctx.os.environ[cc_path_envvar].strip()
+ if cc_name_from_env:
+ cc_name = cc_name_from_env
+ if cc_name.startswith("/"):
+ # Absolute path, maybe we should make this supported by our which function.
+ return cc_name
+ cc = repository_ctx.which(cc_name)
+ if cc == None:
+ fail(("Cannot find {}, either correct your path or set the {}" +
+ " environment variable").format(target_cc_name, cc_path_envvar))
+ return cc
+
+_INC_DIR_MARKER_BEGIN = "#include <...>"
+
+def _cxx_inc_convert(path):
+ """Convert path returned by cc -E xc++ in a complete path."""
+ path = path.strip()
+ return path
+
+def _get_cxx_inc_directories_impl(repository_ctx, cc, lang_is_cpp):
+ """Compute the list of default C or C++ include directories."""
+ if lang_is_cpp:
+ lang = "c++"
+ else:
+ lang = "c"
+
+ # TODO: We pass -no-canonical-prefixes here to match the compiler flags,
+ # but in rocm_clang CROSSTOOL file that is a `feature` and we should
+ # handle the case when it's disabled and no flag is passed
+ result = repository_ctx.execute([
+ cc,
+ "-no-canonical-prefixes",
+ "-E",
+ "-x" + lang,
+ "-",
+ "-v",
+ ])
+ index1 = result.stderr.find(_INC_DIR_MARKER_BEGIN)
+ if index1 == -1:
+ return []
+ index1 = result.stderr.find("\n", index1)
+ if index1 == -1:
+ return []
+ index2 = result.stderr.rfind("\n ")
+ if index2 == -1 or index2 < index1:
+ return []
+ index2 = result.stderr.find("\n", index2 + 1)
+ if index2 == -1:
+ inc_dirs = result.stderr[index1 + 1:]
+ else:
+ inc_dirs = result.stderr[index1 + 1:index2].strip()
+
+ return [
+ str(repository_ctx.path(_cxx_inc_convert(p)))
+ for p in inc_dirs.split("\n")
+ ]
+
+def get_cxx_inc_directories(repository_ctx, cc):
+ """Compute the list of default C and C++ include directories."""
+
+ # For some reason `clang -xc` sometimes returns include paths that are
+ # different from the ones from `clang -xc++`. (Symlink and a dir)
+ # So we run the compiler with both `-xc` and `-xc++` and merge resulting lists
+ includes_cpp = _get_cxx_inc_directories_impl(repository_ctx, cc, True)
+ includes_c = _get_cxx_inc_directories_impl(repository_ctx, cc, False)
+
+ includes_cpp_set = depset(includes_cpp)
+ return includes_cpp + [
+ inc
+ for inc in includes_c
+ if inc not in includes_cpp_set
+ ]
+
+def auto_configure_fail(msg):
+ """Output failure message when rocm configuration fails."""
+ red = "\033[0;31m"
+ no_color = "\033[0m"
+ fail("\n%sROCm Configuration Error:%s %s\n" % (red, no_color, msg))
+
+# END cc_configure common functions (see TODO above).
+
+def _host_compiler_includes(repository_ctx, cc):
+ """Generates the cxx_builtin_include_directory entries for gcc inc dirs.
+
+ Args:
+ repository_ctx: The repository context.
+ cc: The path to the gcc host compiler.
+
+ Returns:
+ A string containing the cxx_builtin_include_directory for each of the gcc
+ host compiler include directories, which can be added to the CROSSTOOL
+ file.
+ """
+ inc_dirs = get_cxx_inc_directories(repository_ctx, cc)
+
+ # Add numpy headers
+ inc_dirs.append("/usr/lib/python2.7/dist-packages/numpy/core/include")
+
+ entries = []
+ for inc_dir in inc_dirs:
+ entries.append(" cxx_builtin_include_directory: \"%s\"" % inc_dir)
+
+ # define TENSORFLOW_USE_ROCM
+ entries.append(" unfiltered_cxx_flag: \"-DTENSORFLOW_USE_ROCM\"")
+
+ return "\n".join(entries)
+
+def _rocm_include_path(repository_ctx, rocm_config):
+ """Generates the cxx_builtin_include_directory entries for rocm inc dirs.
+
+ Args:
+ repository_ctx: The repository context.
+ cc: The path to the gcc host compiler.
+
+ Returns:
+ A string containing the cxx_builtin_include_directory for each of the gcc
+ host compiler include directories, which can be added to the CROSSTOOL
+ file.
+ """
+ inc_dirs = []
+
+ # general ROCm include path
+ inc_dirs.append(rocm_config.rocm_toolkit_path + "/include")
+
+ # Add HSA headers
+ inc_dirs.append("/opt/rocm/hsa/include")
+
+ # Add HIP headers
+ inc_dirs.append("/opt/rocm/include/hip")
+ inc_dirs.append("/opt/rocm/include/hip/hcc_detail")
+
+ # Add rocrand and hiprand headers
+ inc_dirs.append("/opt/rocm/rocrand/include")
+ inc_dirs.append("/opt/rocm/hiprand/include")
+
+ # Add rocfft headers
+ inc_dirs.append("/opt/rocm/rocfft/include")
+
+ # Add rocBLAS headers
+ inc_dirs.append("/opt/rocm/rocblas/include")
+
+ # Add MIOpen headers
+ inc_dirs.append("/opt/rocm/miopen/include")
+
+ # Add hcc headers
+ inc_dirs.append("/opt/rocm/hcc/include")
+ inc_dirs.append("/opt/rocm/hcc/compiler/lib/clang/7.0.0/include/")
+ inc_dirs.append("/opt/rocm/hcc/lib/clang/7.0.0/include")
+
+ # Newer hcc builds use/are based off of clang 8.0.0.
+ inc_dirs.append("/opt/rocm/hcc/compiler/lib/clang/8.0.0/include/")
+ inc_dirs.append("/opt/rocm/hcc/lib/clang/8.0.0/include")
+
+ inc_entries = []
+ for inc_dir in inc_dirs:
+ inc_entries.append(" cxx_builtin_include_directory: \"%s\"" % inc_dir)
+ return "\n".join(inc_entries)
+
+def _enable_rocm(repository_ctx):
+ if "TF_NEED_ROCM" in repository_ctx.os.environ:
+ enable_rocm = repository_ctx.os.environ["TF_NEED_ROCM"].strip()
+ return enable_rocm == "1"
+ return False
+
+def _rocm_toolkit_path(repository_ctx):
+ """Finds the rocm toolkit directory.
+
+ Args:
+ repository_ctx: The repository context.
+
+ Returns:
+ A speculative real path of the rocm toolkit install directory.
+ """
+ rocm_toolkit_path = _DEFAULT_ROCM_TOOLKIT_PATH
+ if _ROCM_TOOLKIT_PATH in repository_ctx.os.environ:
+ rocm_toolkit_path = repository_ctx.os.environ[_ROCM_TOOLKIT_PATH].strip()
+ if not repository_ctx.path(rocm_toolkit_path).exists:
+ auto_configure_fail("Cannot find rocm toolkit path.")
+ return str(repository_ctx.path(rocm_toolkit_path).realpath)
+
+def _amdgpu_targets(repository_ctx):
+ """Returns a list of strings representing AMDGPU targets."""
+ if _TF_ROCM_AMDGPU_TARGETS not in repository_ctx.os.environ:
+ return _DEFAULT_ROCM_AMDGPU_TARGETS
+ amdgpu_targets_str = repository_ctx.os.environ[_TF_ROCM_AMDGPU_TARGETS]
+ amdgpu_targets = amdgpu_targets_str.split(",")
+ for amdgpu_target in amdgpu_targets:
+ if amdgpu_target[:3] != "gfx" or not amdgpu_target[3:].isdigit():
+ auto_configure_fail("Invalid AMDGPU target: %s" % amdgpu_target)
+ return amdgpu_targets
+
+def _cpu_value(repository_ctx):
+ """Returns the name of the host operating system.
+
+ Args:
+ repository_ctx: The repository context.
+
+ Returns:
+ A string containing the name of the host operating system.
+ """
+ os_name = repository_ctx.os.name.lower()
+ if os_name.startswith("mac os"):
+ return "Darwin"
+ if os_name.find("windows") != -1:
+ return "Windows"
+ result = repository_ctx.execute(["uname", "-s"])
+ return result.stdout.strip()
+
+def _lib_name(lib, cpu_value, version = "", static = False):
+ """Constructs the platform-specific name of a library.
+
+ Args:
+ lib: The name of the library, such as "hip"
+ cpu_value: The name of the host operating system.
+ version: The version of the library.
+ static: True the library is static or False if it is a shared object.
+
+ Returns:
+ The platform-specific name of the library.
+ """
+ if cpu_value in ("Linux"):
+ if static:
+ return "lib%s.a" % lib
+ else:
+ if version:
+ version = ".%s" % version
+ return "lib%s.so%s" % (lib, version)
+ elif cpu_value == "Windows":
+ return "%s.lib" % lib
+ elif cpu_value == "Darwin":
+ if static:
+ return "lib%s.a" % lib
+ elif version:
+ version = ".%s" % version
+ return "lib%s%s.dylib" % (lib, version)
+ else:
+ auto_configure_fail("Invalid cpu_value: %s" % cpu_value)
+
+def _find_rocm_lib(
+ lib,
+ repository_ctx,
+ cpu_value,
+ basedir,
+ version = "",
+ static = False):
+ """Finds the given ROCm libraries on the system.
+
+ Args:
+ lib: The name of the library, such as "hip"
+ repository_ctx: The repository context.
+ cpu_value: The name of the host operating system.
+ basedir: The install directory of ROCm.
+ version: The version of the library.
+ static: True if static library, False if shared object.
+
+ Returns:
+ Returns a struct with the following fields:
+ file_name: The basename of the library found on the system.
+ path: The full path to the library.
+ """
+ file_name = _lib_name(lib, cpu_value, version, static)
+ if cpu_value == "Linux":
+ path = repository_ctx.path("%s/lib64/%s" % (basedir, file_name))
+ if path.exists:
+ return struct(file_name = file_name, path = str(path.realpath))
+ path = repository_ctx.path("%s/lib64/stubs/%s" % (basedir, file_name))
+ if path.exists:
+ return struct(file_name = file_name, path = str(path.realpath))
+ path = repository_ctx.path(
+ "%s/lib/x86_64-linux-gnu/%s" % (basedir, file_name),
+ )
+ if path.exists:
+ return struct(file_name = file_name, path = str(path.realpath))
+
+ path = repository_ctx.path("%s/lib/%s" % (basedir, file_name))
+ if path.exists:
+ return struct(file_name = file_name, path = str(path.realpath))
+ path = repository_ctx.path("%s/%s" % (basedir, file_name))
+ if path.exists:
+ return struct(file_name = file_name, path = str(path.realpath))
+
+ auto_configure_fail("Cannot find rocm library %s" % file_name)
+
+def _find_libs(repository_ctx, rocm_config):
+ """Returns the ROCm libraries on the system.
+
+ Args:
+ repository_ctx: The repository context.
+ rocm_config: The ROCm config as returned by _get_rocm_config
+
+ Returns:
+ Map of library names to structs of filename and path as returned by
+ _find_rocm_lib.
+ """
+ cpu_value = rocm_config.cpu_value
+ return {
+ "hip": _find_rocm_lib(
+ "hip_hcc",
+ repository_ctx,
+ cpu_value,
+ rocm_config.rocm_toolkit_path,
+ ),
+ "rocblas": _find_rocm_lib(
+ "rocblas",
+ repository_ctx,
+ cpu_value,
+ rocm_config.rocm_toolkit_path + "/rocblas",
+ ),
+ "rocfft": _find_rocm_lib(
+ "rocfft",
+ repository_ctx,
+ cpu_value,
+ rocm_config.rocm_toolkit_path + "/rocfft",
+ ),
+ "hiprand": _find_rocm_lib(
+ "hiprand",
+ repository_ctx,
+ cpu_value,
+ rocm_config.rocm_toolkit_path + "/hiprand",
+ ),
+ "miopen": _find_rocm_lib(
+ "MIOpen",
+ repository_ctx,
+ cpu_value,
+ rocm_config.rocm_toolkit_path + "/miopen",
+ ),
+ }
+
+def _get_rocm_config(repository_ctx):
+ """Detects and returns information about the ROCm installation on the system.
+
+ Args:
+ repository_ctx: The repository context.
+
+ Returns:
+ A struct containing the following fields:
+ rocm_toolkit_path: The ROCm toolkit installation directory.
+ amdgpu_targets: A list of the system's AMDGPU targets.
+ cpu_value: The name of the host operating system.
+ """
+ cpu_value = _cpu_value(repository_ctx)
+ rocm_toolkit_path = _rocm_toolkit_path(repository_ctx)
+ return struct(
+ rocm_toolkit_path = rocm_toolkit_path,
+ amdgpu_targets = _amdgpu_targets(repository_ctx),
+ cpu_value = cpu_value,
+ )
+
+def _tpl(repository_ctx, tpl, substitutions = {}, out = None):
+ if not out:
+ out = tpl.replace(":", "/")
+ repository_ctx.template(
+ out,
+ Label("//third_party/gpus/%s.tpl" % tpl),
+ substitutions,
+ )
+
+def _file(repository_ctx, label):
+ repository_ctx.template(
+ label.replace(":", "/"),
+ Label("//third_party/gpus/%s.tpl" % label),
+ {},
+ )
+
+_DUMMY_CROSSTOOL_BZL_FILE = """
+def error_gpu_disabled():
+ fail("ERROR: Building with --config=rocm but TensorFlow is not configured " +
+ "to build with GPU support. Please re-run ./configure and enter 'Y' " +
+ "at the prompt to build with GPU support.")
+
+ native.genrule(
+ name = "error_gen_crosstool",
+ outs = ["CROSSTOOL"],
+ cmd = "echo 'Should not be run.' && exit 1",
+ )
+
+ native.filegroup(
+ name = "crosstool",
+ srcs = [":CROSSTOOL"],
+ output_licenses = ["unencumbered"],
+ )
+"""
+
+_DUMMY_CROSSTOOL_BUILD_FILE = """
+load("//crosstool:error_gpu_disabled.bzl", "error_gpu_disabled")
+
+error_gpu_disabled()
+"""
+
+def _create_dummy_repository(repository_ctx):
+ cpu_value = _cpu_value(repository_ctx)
+
+ # Set up BUILD file for rocm/.
+ _tpl(
+ repository_ctx,
+ "rocm:build_defs.bzl",
+ {
+ "%{rocm_is_configured}": "False",
+ "%{rocm_extra_copts}": "[]",
+ },
+ )
+ _tpl(
+ repository_ctx,
+ "rocm:BUILD",
+ {
+ "%{hip_lib}": _lib_name("hip", cpu_value),
+ "%{rocblas_lib}": _lib_name("rocblas", cpu_value),
+ "%{miopen_lib}": _lib_name("miopen", cpu_value),
+ "%{rocfft_lib}": _lib_name("rocfft", cpu_value),
+ "%{hiprand_lib}": _lib_name("hiprand", cpu_value),
+ "%{rocm_include_genrules}": "",
+ "%{rocm_headers}": "",
+ },
+ )
+
+ # Create dummy files for the ROCm toolkit since they are still required by
+ # tensorflow/core/platform/default/build_config:rocm.
+ repository_ctx.file("rocm/hip/include/hip/hip_runtime.h", "")
+
+ # Set up rocm_config.h, which is used by
+ # tensorflow/stream_executor/dso_loader.cc.
+ _tpl(
+ repository_ctx,
+ "rocm:rocm_config.h",
+ {
+ "%{rocm_toolkit_path}": _DEFAULT_ROCM_TOOLKIT_PATH,
+ },
+ "rocm/rocm/rocm_config.h",
+ )
+
+ # If rocm_configure is not configured to build with GPU support, and the user
+ # attempts to build with --config=rocm, add a dummy build rule to intercept
+ # this and fail with an actionable error message.
+ repository_ctx.file(
+ "crosstool/error_gpu_disabled.bzl",
+ _DUMMY_CROSSTOOL_BZL_FILE,
+ )
+ repository_ctx.file("crosstool/BUILD", _DUMMY_CROSSTOOL_BUILD_FILE)
+
+def _execute(
+ repository_ctx,
+ cmdline,
+ error_msg = None,
+ error_details = None,
+ empty_stdout_fine = False):
+ """Executes an arbitrary shell command.
+
+ Args:
+ repository_ctx: the repository_ctx object
+ cmdline: list of strings, the command to execute
+ error_msg: string, a summary of the error if the command fails
+ error_details: string, details about the error or steps to fix it
+ empty_stdout_fine: bool, if True, an empty stdout result is fine, otherwise
+ it's an error
+ Return:
+ the result of repository_ctx.execute(cmdline)
+ """
+ result = repository_ctx.execute(cmdline)
+ if result.stderr or not (empty_stdout_fine or result.stdout):
+ auto_configure_fail(
+ "\n".join([
+ error_msg.strip() if error_msg else "Repository command failed",
+ result.stderr.strip(),
+ error_details if error_details else "",
+ ]),
+ )
+ return result
+
+def _norm_path(path):
+ """Returns a path with '/' and remove the trailing slash."""
+ path = path.replace("\\", "/")
+ if path[-1] == "/":
+ path = path[:-1]
+ return path
+
+def _symlink_genrule_for_dir(
+ repository_ctx,
+ src_dir,
+ dest_dir,
+ genrule_name,
+ src_files = [],
+ dest_files = []):
+ """Returns a genrule to symlink(or copy if on Windows) a set of files.
+
+ If src_dir is passed, files will be read from the given directory; otherwise
+ we assume files are in src_files and dest_files
+ """
+ if src_dir != None:
+ src_dir = _norm_path(src_dir)
+ dest_dir = _norm_path(dest_dir)
+ files = _read_dir(repository_ctx, src_dir)
+
+ # Create a list with the src_dir stripped to use for outputs.
+ dest_files = files.replace(src_dir, "").splitlines()
+ src_files = files.splitlines()
+ command = []
+
+ # We clear folders that might have been generated previously to avoid
+ # undesired inclusions
+ command.append('if [ -d "$(@D)/include" ]; then rm $(@D)/include -drf; fi')
+ command.append('if [ -d "$(@D)/lib" ]; then rm $(@D)/lib -drf; fi')
+ outs = []
+ for i in range(len(dest_files)):
+ if dest_files[i] != "":
+ # If we have only one file to link we do not want to use the dest_dir, as
+ # $(@D) will include the full path to the file.
+ dest = "$(@D)/" + dest_dir + dest_files[i] if len(dest_files) != 1 else "$(@D)/" + dest_files[i]
+
+ # On Windows, symlink is not supported, so we just copy all the files.
+ cmd = "ln -s"
+ command.append(cmd + ' "%s" "%s"' % (src_files[i], dest))
+ outs.append(' "' + dest_dir + dest_files[i] + '",')
+ genrule = _genrule(
+ src_dir,
+ genrule_name,
+ " && ".join(command),
+ "\n".join(outs),
+ )
+ return genrule
+
+def _genrule(src_dir, genrule_name, command, outs):
+ """Returns a string with a genrule.
+
+ Genrule executes the given command and produces the given outputs.
+ """
+ return (
+ "genrule(\n" +
+ ' name = "' +
+ genrule_name + '",\n' +
+ " outs = [\n" +
+ outs +
+ "\n ],\n" +
+ ' cmd = """\n' +
+ command +
+ '\n """,\n' +
+ ")\n"
+ )
+
+def _read_dir(repository_ctx, src_dir):
+ """Returns a string with all files in a directory.
+
+ Finds all files inside a directory, traversing subfolders and following
+ symlinks. The returned string contains the full path of all files
+ separated by line breaks.
+ """
+ find_result = _execute(
+ repository_ctx,
+ ["find", src_dir, "-follow", "-type", "f"],
+ empty_stdout_fine = True,
+ )
+ result = find_result.stdout
+ return result
+
+def _compute_rocm_extra_copts(repository_ctx, amdgpu_targets):
+ if False:
+ amdgpu_target_flags = ["--amdgpu-target=" +
+ amdgpu_target for amdgpu_target in amdgpu_targets]
+ else:
+ # AMDGPU targets are handled in the "crosstool_wrapper_driver_is_not_gcc"
+ amdgpu_target_flags = []
+ return str(amdgpu_target_flags)
+
+def _create_local_rocm_repository(repository_ctx):
+ """Creates the repository containing files set up to build with ROCm."""
+ rocm_config = _get_rocm_config(repository_ctx)
+
+ # Set up symbolic links for the rocm toolkit by creating genrules to do
+ # symlinking. We create one genrule for each directory we want to track under
+ # rocm_toolkit_path
+ rocm_toolkit_path = rocm_config.rocm_toolkit_path
+ rocm_include_path = rocm_toolkit_path + "/include"
+ genrules = [_symlink_genrule_for_dir(
+ repository_ctx,
+ rocm_include_path,
+ "rocm/include",
+ "rocm-include",
+ )]
+ genrules.append(_symlink_genrule_for_dir(
+ repository_ctx,
+ rocm_toolkit_path + "/rocfft/include",
+ "rocm/include/rocfft",
+ "rocfft-include",
+ ))
+ genrules.append(_symlink_genrule_for_dir(
+ repository_ctx,
+ rocm_toolkit_path + "/rocblas/include",
+ "rocm/include/rocblas",
+ "rocblas-include",
+ ))
+ genrules.append(_symlink_genrule_for_dir(
+ repository_ctx,
+ rocm_toolkit_path + "/miopen/include",
+ "rocm/include/miopen",
+ "miopen-include",
+ ))
+
+ rocm_libs = _find_libs(repository_ctx, rocm_config)
+ rocm_lib_src = []
+ rocm_lib_dest = []
+ for lib in rocm_libs.values():
+ rocm_lib_src.append(lib.path)
+ rocm_lib_dest.append("rocm/lib/" + lib.file_name)
+ genrules.append(_symlink_genrule_for_dir(
+ repository_ctx,
+ None,
+ "",
+ "rocm-lib",
+ rocm_lib_src,
+ rocm_lib_dest,
+ ))
+
+ included_files = _read_dir(repository_ctx, rocm_include_path).replace(
+ rocm_include_path,
+ "",
+ ).splitlines()
+
+ # Set up BUILD file for rocm/
+ _tpl(
+ repository_ctx,
+ "rocm:build_defs.bzl",
+ {
+ "%{rocm_is_configured}": "True",
+ "%{rocm_extra_copts}": _compute_rocm_extra_copts(
+ repository_ctx,
+ rocm_config.amdgpu_targets,
+ ),
+ },
+ )
+ _tpl(
+ repository_ctx,
+ "rocm:BUILD",
+ {
+ "%{hip_lib}": rocm_libs["hip"].file_name,
+ "%{rocblas_lib}": rocm_libs["rocblas"].file_name,
+ "%{rocfft_lib}": rocm_libs["rocfft"].file_name,
+ "%{hiprand_lib}": rocm_libs["hiprand"].file_name,
+ "%{miopen_lib}": rocm_libs["miopen"].file_name,
+ "%{rocm_include_genrules}": "\n".join(genrules),
+ "%{rocm_headers}": ('":rocm-include",\n' +
+ '":rocfft-include",\n' +
+ '":rocblas-include",\n' +
+ '":miopen-include",'),
+ },
+ )
+
+ # Set up crosstool/
+ _tpl(repository_ctx, "crosstool:BUILD", {"%{linker_files}": ":empty", "%{win_linker_files}": ":empty"})
+ cc = find_cc(repository_ctx)
+ host_compiler_includes = _host_compiler_includes(repository_ctx, cc)
+ rocm_defines = {
+ "%{rocm_include_path}": _rocm_include_path(
+ repository_ctx,
+ rocm_config,
+ ),
+ "%{host_compiler_includes}": host_compiler_includes,
+ "%{clang_path}": str(cc),
+ }
+
+ _tpl(repository_ctx, "crosstool:CROSSTOOL_hipcc", rocm_defines, out = "crosstool/CROSSTOOL")
+
+ _tpl(
+ repository_ctx,
+ "crosstool:clang/bin/crosstool_wrapper_driver_rocm",
+ {
+ "%{cpu_compiler}": str(cc),
+ "%{hipcc_path}": "/opt/rocm/bin/hipcc",
+ "%{gcc_host_compiler_path}": str(cc),
+ "%{rocm_amdgpu_targets}": ",".join(
+ ["\"%s\"" % c for c in rocm_config.amdgpu_targets],
+ ),
+ },
+ )
+
+ # Set up rocm_config.h, which is used by
+ # tensorflow/stream_executor/dso_loader.cc.
+ _tpl(
+ repository_ctx,
+ "rocm:rocm_config.h",
+ {
+ "%{rocm_amdgpu_targets}": ",".join(
+ ["\"%s\"" % c for c in rocm_config.amdgpu_targets],
+ ),
+ "%{rocm_toolkit_path}": rocm_config.rocm_toolkit_path,
+ },
+ "rocm/rocm/rocm_config.h",
+ )
+
+def _create_remote_rocm_repository(repository_ctx, remote_config_repo):
+ """Creates pointers to a remotely configured repo set up to build with ROCm."""
+ _tpl(
+ repository_ctx,
+ "rocm:build_defs.bzl",
+ {
+ "%{rocm_is_configured}": "True",
+ "%{rocm_extra_copts}": _compute_rocm_extra_copts(
+ repository_ctx, #_compute_capabilities(repository_ctx)
+ ),
+ },
+ )
+ _tpl(
+ repository_ctx,
+ "rocm:remote.BUILD",
+ {
+ "%{remote_rocm_repo}": remote_config_repo,
+ },
+ "rocm/BUILD",
+ )
+ _tpl(repository_ctx, "crosstool:remote.BUILD", {
+ "%{remote_rocm_repo}": remote_config_repo,
+ }, "crosstool/BUILD")
+
+def _rocm_autoconf_impl(repository_ctx):
+ """Implementation of the rocm_autoconf repository rule."""
+ if not _enable_rocm(repository_ctx):
+ _create_dummy_repository(repository_ctx)
+ elif _TF_ROCM_CONFIG_REPO in repository_ctx.os.environ:
+ _create_remote_rocm_repository(
+ repository_ctx,
+ repository_ctx.os.environ[_TF_ROCM_CONFIG_REPO],
+ )
+ else:
+ _create_local_rocm_repository(repository_ctx)
+
+rocm_configure = repository_rule(
+ implementation = _rocm_autoconf_impl,
+ environ = [
+ _GCC_HOST_COMPILER_PATH,
+ "TF_NEED_ROCM",
+ _ROCM_TOOLKIT_PATH,
+ _TF_ROCM_VERSION,
+ _TF_MIOPEN_VERSION,
+ _TF_ROCM_AMDGPU_TARGETS,
+ _TF_ROCM_CONFIG_REPO,
+ ],
+)
+
+"""Detects and configures the local ROCm toolchain.
+
+Add the following to your WORKSPACE FILE:
+
+```python
+rocm_configure(name = "local_config_rocm")
+```
+
+Args:
+ name: A unique name for this workspace rule.
+"""
diff --git a/third_party/icu/BUILD b/third_party/icu/BUILD
new file mode 100644
index 0000000..82bab3f
--- /dev/null
+++ b/third_party/icu/BUILD
@@ -0,0 +1 @@
+# This empty BUILD file is required to make Bazel treat this directory as a package.
diff --git a/third_party/icu/BUILD.bazel b/third_party/icu/BUILD.bazel
new file mode 100644
index 0000000..36d6b90
--- /dev/null
+++ b/third_party/icu/BUILD.bazel
@@ -0,0 +1,88 @@
+package(
+ default_visibility = ["//visibility:public"],
+)
+
+licenses(["notice"]) # Apache 2.0
+
+exports_files([
+ "icu4c/LICENSE",
+ "icu4j/main/shared/licenses/LICENSE",
+])
+
+cc_library(
+ name = "headers",
+ hdrs = glob(["icu4c/source/common/unicode/*.h"]),
+ includes = [
+ "icu4c/source/common",
+ ],
+ deps = [
+ ],
+)
+
+cc_library(
+ name = "common",
+ hdrs = glob(["icu4c/source/common/unicode/*.h"]),
+ includes = [
+ "icu4c/source/common",
+ ],
+ deps = [
+ ":icuuc",
+ ],
+)
+
+cc_library(
+ name = "icuuc",
+ srcs = glob(
+ [
+ "icu4c/source/common/*.c",
+ "icu4c/source/common/*.cpp",
+ "icu4c/source/stubdata/*.cpp",
+ ],
+ ),
+ hdrs = glob([
+ "icu4c/source/common/*.h",
+ ]),
+ copts = [
+ "-DU_COMMON_IMPLEMENTATION",
+ "-DU_HAVE_STD_ATOMICS",
+ ] + select({
+ ":android": [
+ "-fdata-sections",
+ "-DGOOGLE_VENDOR_SRC_BRANCH",
+ "-DU_HAVE_NL_LANGINFO_CODESET=0",
+ "-Wno-deprecated-declarations",
+ ],
+ ":apple": [
+ "-DGOOGLE_VENDOR_SRC_BRANCH",
+ "-Wno-shorten-64-to-32",
+ "-Wno-unused-variable",
+ ],
+ ":windows": [
+ "/utf-8",
+ "/DLOCALE_ALLOW_NEUTRAL_NAMES=0",
+ ],
+ "//conditions:default": [],
+ }),
+ tags = ["requires-rtti"],
+ visibility = [
+ "//visibility:private",
+ ],
+ deps = [
+ ":headers",
+ ],
+)
+
+config_setting(
+ name = "android",
+ values = {"crosstool_top": "//external:android/crosstool"},
+)
+
+config_setting(
+ name = "apple",
+ values = {"cpu": "darwin"},
+)
+
+config_setting(
+ name = "windows",
+ values = {"cpu": "x64_windows"},
+)
diff --git a/third_party/icu/workspace.bzl b/third_party/icu/workspace.bzl
new file mode 100644
index 0000000..bfebf42
--- /dev/null
+++ b/third_party/icu/workspace.bzl
@@ -0,0 +1,15 @@
+"""Loads a lightweight subset of the ICU library for Unicode processing."""
+
+load("//third_party:repo.bzl", "third_party_http_archive")
+
+def repo():
+ third_party_http_archive(
+ name = "icu",
+ strip_prefix = "icu-release-62-1",
+ sha256 = "e15ffd84606323cbad5515bf9ecdf8061cc3bf80fb883b9e6aa162e485aa9761",
+ urls = [
+ "https://mirror.bazel.build/github.com/unicode-org/icu/archive/release-62-1.tar.gz",
+ "https://github.com/unicode-org/icu/archive/release-62-1.tar.gz",
+ ],
+ build_file = "//third_party/icu:BUILD.bazel",
+ )
diff --git a/third_party/jemalloc.BUILD b/third_party/jemalloc.BUILD
deleted file mode 100644
index 1b0829b..0000000
--- a/third_party/jemalloc.BUILD
+++ /dev/null
@@ -1,356 +0,0 @@
-# Description:
-# jemalloc - a general-purpose scalable concurrent malloc implementation
-
-licenses(["notice"]) # BSD
-
-exports_files(["COPYING"])
-
-load("@org_tensorflow//third_party:common.bzl", "template_rule")
-
-cc_library(
- name = "jemalloc_headers",
- hdrs = ["include/jemalloc/jemalloc.h"],
- includes = ["include"],
- visibility = ["//visibility:public"],
-)
-
-cc_library(
- name = "jemalloc_impl",
- srcs = [
- "src/arena.c",
- "src/atomic.c",
- "src/base.c",
- "src/bitmap.c",
- "src/chunk.c",
- "src/chunk_dss.c",
- "src/chunk_mmap.c",
- "src/ckh.c",
- "src/ctl.c",
- "src/extent.c",
- "src/hash.c",
- "src/huge.c",
- "src/jemalloc.c",
- "src/mb.c",
- "src/mutex.c",
- "src/nstime.c",
- "src/pages.c",
- "src/prng.c",
- "src/prof.c",
- "src/quarantine.c",
- "src/rtree.c",
- "src/spin.c",
- "src/stats.c",
- "src/tcache.c",
- "src/tsd.c",
- "src/util.c",
- "src/witness.c",
- ],
- hdrs = [
- "include/jemalloc/internal/arena.h",
- "include/jemalloc/internal/assert.h",
- "include/jemalloc/internal/atomic.h",
- "include/jemalloc/internal/base.h",
- "include/jemalloc/internal/bitmap.h",
- "include/jemalloc/internal/chunk.h",
- "include/jemalloc/internal/chunk_dss.h",
- "include/jemalloc/internal/chunk_mmap.h",
- "include/jemalloc/internal/ckh.h",
- "include/jemalloc/internal/ctl.h",
- "include/jemalloc/internal/extent.h",
- "include/jemalloc/internal/hash.h",
- "include/jemalloc/internal/huge.h",
- "include/jemalloc/internal/jemalloc_internal.h",
- "include/jemalloc/internal/jemalloc_internal_decls.h",
- "include/jemalloc/internal/jemalloc_internal_defs.h",
- "include/jemalloc/internal/jemalloc_internal_macros.h",
- "include/jemalloc/internal/mb.h",
- "include/jemalloc/internal/mutex.h",
- "include/jemalloc/internal/nstime.h",
- "include/jemalloc/internal/pages.h",
- "include/jemalloc/internal/ph.h",
- "include/jemalloc/internal/private_namespace.h",
- "include/jemalloc/internal/prng.h",
- "include/jemalloc/internal/prof.h",
- "include/jemalloc/internal/ql.h",
- "include/jemalloc/internal/qr.h",
- "include/jemalloc/internal/quarantine.h",
- "include/jemalloc/internal/rb.h",
- "include/jemalloc/internal/rtree.h",
- "include/jemalloc/internal/size_classes.h",
- "include/jemalloc/internal/smoothstep.h",
- "include/jemalloc/internal/spin.h",
- "include/jemalloc/internal/stats.h",
- "include/jemalloc/internal/tcache.h",
- "include/jemalloc/internal/ticker.h",
- "include/jemalloc/internal/tsd.h",
- "include/jemalloc/internal/util.h",
- "include/jemalloc/internal/valgrind.h",
- "include/jemalloc/internal/witness.h",
- ],
- # Same flags that jemalloc uses to build.
- copts = [
- "-O3",
- "-funroll-loops",
- "-D_GNU_SOURCE",
- "-D_REENTRANT",
- ],
- includes = ["include"],
- # pthread_atfork() is called for PPC.
- linkopts = select({
- "@org_tensorflow//tensorflow:linux_ppc64le": [
- "-lpthread",
- ],
- "@org_tensorflow//tensorflow:linux_x86_64": [
- "-lpthread",
- ],
- "//conditions:default": [
- ],
- }),
- visibility = ["//visibility:public"],
- deps = [":jemalloc_headers"],
-)
-
-sh_binary(
- name = "jemalloc_sh",
- srcs = ["include/jemalloc/jemalloc.sh"],
-)
-
-genrule(
- name = "jemalloc_h",
- srcs = [
- ":jemalloc_defs_h",
- ":jemalloc_macros_h",
- ":jemalloc_mangle_h",
- ":jemalloc_protos_h",
- ":jemalloc_rename_h",
- ":jemalloc_typedefs_h",
- ],
- outs = ["include/jemalloc/jemalloc.h"],
- cmd = "$(location :jemalloc_sh) $$(dirname $(location :jemalloc_defs_h))/../../ >$@",
- tools = [":jemalloc_sh"],
-)
-
-# Add to this list if you want to export more symbols from jemalloc.
-genrule(
- name = "public_symbols_txt",
- outs = ["include/jemalloc/internal/public_symbols.txt"],
- cmd = "\n".join([
- "cat <<'EOF' > $@",
- "free:jemalloc_free",
- "malloc:jemalloc_malloc",
- "posix_memalign:jemalloc_posix_memalign",
- "realloc:jemalloc_realloc",
- "EOF",
- ]),
-)
-
-sh_binary(
- name = "jemalloc_mangle_sh",
- srcs = ["include/jemalloc/jemalloc_mangle.sh"],
-)
-
-genrule(
- name = "jemalloc_mangle_h",
- srcs = [":public_symbols_txt"],
- outs = ["include/jemalloc/jemalloc_mangle.h"],
- cmd = "$(location :jemalloc_mangle_sh) $(location :public_symbols_txt) je_ >$@",
- tools = [":jemalloc_mangle_sh"],
-)
-
-sh_binary(
- name = "jemalloc_rename_sh",
- srcs = ["include/jemalloc/jemalloc_rename.sh"],
-)
-
-genrule(
- name = "jemalloc_rename_h",
- srcs = [":public_symbols_txt"],
- outs = ["include/jemalloc/jemalloc_rename.h"],
- cmd = "$(location :jemalloc_rename_sh) $(location :public_symbols_txt) >$@",
- tools = [":jemalloc_rename_sh"],
-)
-
-sh_binary(
- name = "private_namespace_sh",
- srcs = ["include/jemalloc/internal/private_namespace.sh"],
-)
-
-genrule(
- name = "private_namespace_h",
- srcs = ["include/jemalloc/internal/private_symbols.txt"],
- outs = ["include/jemalloc/internal/private_namespace.h"],
- cmd = "$(location :private_namespace_sh) $(location include/jemalloc/internal/private_symbols.txt) >$@",
- tools = [":private_namespace_sh"],
-)
-
-sh_binary(
- name = "public_namespace_sh",
- srcs = ["include/jemalloc/internal/public_namespace.sh"],
-)
-
-genrule(
- name = "public_namespace_h",
- srcs = [":public_symbols_txt"],
- outs = ["include/jemalloc/internal/public_namespace.h"],
- cmd = "$(location :public_namespace_sh) $(location :public_symbols_txt) >$@",
- tools = [":public_namespace_sh"],
-)
-
-sh_binary(
- name = "size_classes_sh",
- srcs = ["include/jemalloc/internal/size_classes.sh"],
-)
-
-# Size classes for Linux x86_64 and ppc64le. Update if adding builds for other
-# architectures. See size_classes.sh for details on the arguments.
-# For default case, kept the arguments same as that of x86_64 for now.
-genrule(
- name = "size_classes_h",
- outs = ["include/jemalloc/internal/size_classes.h"],
- cmd = select({
- "@org_tensorflow//tensorflow:linux_ppc64le": "$(location :size_classes_sh) \"3 4\" 3 16 2 >$@",
- "@org_tensorflow//tensorflow:linux_x86_64": "$(location :size_classes_sh) \"3 4\" 3 12 2 >$@",
- "//conditions:default": "$(location :size_classes_sh) \"3 4\" 3 12 2 >$@",
- }),
- tools = [":size_classes_sh"],
-)
-
-template_rule(
- name = "jemalloc_internal_h",
- src = "include/jemalloc/internal/jemalloc_internal.h.in",
- out = "include/jemalloc/internal/jemalloc_internal.h",
- substitutions = {
- "@private_namespace@": "je_",
- "@install_suffix@": "",
- },
-)
-
-template_rule(
- name = "jemalloc_internal_defs_h",
- src = "include/jemalloc/internal/jemalloc_internal_defs.h.in",
- out = "include/jemalloc/internal/jemalloc_internal_defs.h",
- substitutions = {
- "#undef JEMALLOC_PREFIX": "#define JEMALLOC_PREFIX \"jemalloc_\"",
- "#undef JEMALLOC_CPREFIX": "#define JEMALLOC_CPREFIX \"JEMALLOC_\"",
- "#undef JEMALLOC_PRIVATE_NAMESPACE": "#define JEMALLOC_PRIVATE_NAMESPACE je_",
- "#undef CPU_SPINWAIT": "\n".join([
- "#if defined(__powerpc64__) || defined(__powerpc__)",
- "#define CPU_SPINWAIT __asm__ volatile(\"or 27,27,27\")",
- "#else",
- "#define CPU_SPINWAIT __asm__ volatile(\"pause\")",
- "#endif",
- ]),
- "#undef JEMALLOC_HAVE_BUILTIN_CLZ": "#define JEMALLOC_HAVE_BUILTIN_CLZ",
- "#undef JEMALLOC_USE_SYSCALL": "#define JEMALLOC_USE_SYSCALL",
- "#undef JEMALLOC_HAVE_SECURE_GETENV": "#define JEMALLOC_HAVE_SECURE_GETENV",
- "#undef JEMALLOC_HAVE_PTHREAD_ATFORK": "#define JEMALLOC_HAVE_PTHREAD_ATFORK",
- "#undef JEMALLOC_HAVE_CLOCK_MONOTONIC_COARSE": "#define JEMALLOC_HAVE_CLOCK_MONOTONIC_COARSE 1",
- # Newline required because of substitution conflicts.
- "#undef JEMALLOC_HAVE_CLOCK_MONOTONIC\n": "#define JEMALLOC_HAVE_CLOCK_MONOTONIC 1\n",
- "#undef JEMALLOC_THREADED_INIT": "#define JEMALLOC_THREADED_INIT",
- "#undef JEMALLOC_TLS_MODEL": "#define JEMALLOC_TLS_MODEL __attribute__((tls_model(\"initial-exec\")))",
- "#undef JEMALLOC_CC_SILENCE": "#define JEMALLOC_CC_SILENCE",
- "#undef JEMALLOC_STATS": "#define JEMALLOC_STATS",
- "#undef JEMALLOC_TCACHE": "#define JEMALLOC_TCACHE",
- "#undef JEMALLOC_DSS": "#define JEMALLOC_DSS",
- "#undef JEMALLOC_FILL": "#define JEMALLOC_FILL",
- "#undef LG_TINY_MIN": "#define LG_TINY_MIN 3",
- "#undef LG_PAGE": "\n".join([
- "#if defined(__powerpc64__) || defined(__powerpc__)",
- "#define LG_PAGE 16",
- "#else",
- "#define LG_PAGE 12",
- "#endif",
- ]),
- "#undef JEMALLOC_MAPS_COALESCE": "#define JEMALLOC_MAPS_COALESCE",
- "#undef JEMALLOC_TLS": "#define JEMALLOC_TLS",
- "#undef JEMALLOC_INTERNAL_UNREACHABLE": "#define JEMALLOC_INTERNAL_UNREACHABLE __builtin_unreachable",
- "#undef JEMALLOC_INTERNAL_FFSLL": "#define JEMALLOC_INTERNAL_FFSLL __builtin_ffsll",
- # Newline required because of substitution conflicts.
- "#undef JEMALLOC_INTERNAL_FFSL\n": "#define JEMALLOC_INTERNAL_FFSL __builtin_ffsl\n",
- "#undef JEMALLOC_INTERNAL_FFS\n": "#define JEMALLOC_INTERNAL_FFS __builtin_ffs\n",
- "#undef JEMALLOC_CACHE_OBLIVIOUS": "#define JEMALLOC_CACHE_OBLIVIOUS",
- "#undef JEMALLOC_PROC_SYS_VM_OVERCOMMIT_MEMORY": "#define JEMALLOC_PROC_SYS_VM_OVERCOMMIT_MEMORY",
- "#undef JEMALLOC_HAVE_MADVISE": "#define JEMALLOC_HAVE_MADVISE",
- "#undef JEMALLOC_PURGE_MADVISE_DONTNEED": "#define JEMALLOC_PURGE_MADVISE_DONTNEED",
- "#undef JEMALLOC_THP": "#define JEMALLOC_THP",
- "#undef JEMALLOC_HAS_ALLOCA_H": "#define JEMALLOC_HAS_ALLOCA_H 1",
- # Newline required because of substitution conflicts.
- "#undef LG_SIZEOF_INT\n": "#define LG_SIZEOF_INT 2\n",
- "#undef LG_SIZEOF_LONG\n": "#define LG_SIZEOF_LONG 3\n",
- "#undef LG_SIZEOF_LONG_LONG": "#define LG_SIZEOF_LONG_LONG 3",
- "#undef LG_SIZEOF_INTMAX_T": "#define LG_SIZEOF_INTMAX_T 3",
- "#undef JEMALLOC_GLIBC_MALLOC_HOOK": "#define JEMALLOC_GLIBC_MALLOC_HOOK",
- "#undef JEMALLOC_GLIBC_MEMALIGN_HOOK": "#define JEMALLOC_GLIBC_MEMALIGN_HOOK",
- "#undef JEMALLOC_HAVE_PTHREAD_MUTEX_ADAPTIVE_NP": "#define JEMALLOC_HAVE_PTHREAD_MUTEX_ADAPTIVE_NP",
- "#undef JEMALLOC_CONFIG_MALLOC_CONF": "#define JEMALLOC_CONFIG_MALLOC_CONF \"\"",
- },
-)
-
-template_rule(
- name = "jemalloc_defs_h",
- src = "include/jemalloc/jemalloc_defs.h.in",
- out = "include/jemalloc/jemalloc_defs.h",
- substitutions = {
- "#undef JEMALLOC_HAVE_ATTR": "#define JEMALLOC_HAVE_ATTR",
- "#undef JEMALLOC_HAVE_ATTR_ALLOC_SIZE": "#define JEMALLOC_HAVE_ATTR_ALLOC_SIZE",
- "#undef JEMALLOC_HAVE_ATTR_FORMAT_GNU_PRINTF": "#define JEMALLOC_HAVE_ATTR_FORMAT_GNU_PRINTF",
- "#undef JEMALLOC_HAVE_ATTR_FORMAT_PRINTF": "#define JEMALLOC_HAVE_ATTR_FORMAT_PRINTF",
- "#undef JEMALLOC_OVERRIDE_MEMALIGN": "#define JEMALLOC_OVERRIDE_MEMALIGN",
- "#undef JEMALLOC_OVERRIDE_VALLOC": "#define JEMALLOC_OVERRIDE_VALLOC",
- "#undef JEMALLOC_USABLE_SIZE_CONST": "#define JEMALLOC_USABLE_SIZE_CONST",
- "#undef JEMALLOC_USE_CXX_THROW": "#define JEMALLOC_USE_CXX_THROW",
- "#undef LG_SIZEOF_PTR": "#define LG_SIZEOF_PTR 3",
- },
-)
-
-template_rule(
- name = "jemalloc_macros_h",
- src = "include/jemalloc/jemalloc_macros.h.in",
- out = "include/jemalloc/jemalloc_macros.h",
- substitutions = {
- "@jemalloc_version@": "0.0.0",
- "@jemalloc_version_major@": "0",
- "@jemalloc_version_minor@": "0",
- "@jemalloc_version_bugfix@": "0",
- "@jemalloc_version_nrev@": "0",
- "@jemalloc_version_gid@": "0000000000000000000000000000000000000000",
- },
-)
-
-template_rule(
- name = "jemalloc_protos_h",
- src = "include/jemalloc/jemalloc_protos.h.in",
- out = "include/jemalloc/jemalloc_protos.h",
- substitutions = {
- "@aligned_alloc": "aligned_alloc",
- "@calloc": "calloc",
- "@cbopaque": "cbopaque",
- "@dallocx": "dallocx",
- "@free": "free",
- "@je": "je",
- "@mallctl": "mallctl",
- "@mallctlnametomib": "mallctlnametomib",
- "@mallctlbymib": "mallctlbymib",
- "@malloc_stats_print": "malloc_stats_print",
- "@malloc_usable_size": "malloc_usable_size",
- "@malloc": "malloc",
- "@mallocx": "mallocx",
- "@memalign": "memalign",
- "@nallocx": "nallocx",
- "@posix_memalign": "posix_memalign",
- "@rallocx": "rallocx",
- "@realloc": "realloc",
- "@sallocx": "sallocx",
- "@sdallocx": "sdallocx",
- "@valloc": "valloc",
- "@xallocx": "xallocx",
- },
-)
-
-template_rule(
- name = "jemalloc_typedefs_h",
- src = "include/jemalloc/jemalloc_typedefs.h.in",
- out = "include/jemalloc/jemalloc_typedefs.h",
- substitutions = {},
-)
diff --git a/third_party/mkl/build_defs.bzl b/third_party/mkl/build_defs.bzl
index bb798e7..10c2d90 100644
--- a/third_party/mkl/build_defs.bzl
+++ b/third_party/mkl/build_defs.bzl
@@ -92,7 +92,7 @@
A select evaluating to either if_true or if_false as appropriate.
"""
return select({
- "//third_party/mkl:enable_mkl": if_true,
+ str(Label("//third_party/mkl:enable_mkl")): if_true,
"//conditions:default": if_false,
})
diff --git a/third_party/ngraph/ngraph.BUILD b/third_party/ngraph/ngraph.BUILD
index 71b2187..6602a48 100644
--- a/third_party/ngraph/ngraph.BUILD
+++ b/third_party/ngraph/ngraph.BUILD
@@ -110,7 +110,7 @@
"-I external/ngraph/src",
"-I external/nlohmann_json_lib/include/",
'-D SHARED_LIB_EXT=\\".so\\"',
- '-D NGRAPH_VERSION=\\"0.8.0\\"',
+ '-D NGRAPH_VERSION=\\"0.8.1\\"',
"-D NGRAPH_DEX_ONLY",
],
visibility = ["//visibility:public"],
@@ -144,7 +144,7 @@
"-I external/ngraph/src",
"-I external/nlohmann_json_lib/include/",
'-D SHARED_LIB_EXT=\\".so\\"',
- '-D NGRAPH_VERSION=\\"0.8.0\\"',
+ '-D NGRAPH_VERSION=\\"0.8.1\\"',
],
visibility = ["//visibility:public"],
alwayslink = 1,
diff --git a/third_party/py/python_configure.bzl b/third_party/py/python_configure.bzl
index 3c7e5c8..5326463 100644
--- a/third_party/py/python_configure.bzl
+++ b/third_party/py/python_configure.bzl
@@ -130,8 +130,8 @@
# If we have only one file to link we do not want to use the dest_dir, as
# $(@D) will include the full path to the file.
dest = '$(@D)/' + dest_dir + dest_files[i] if len(dest_files) != 1 else '$(@D)/' + dest_files[i]
- # On Windows, symlink is not supported, so we just copy all the files.
- cmd = 'cp -f' if _is_windows(repository_ctx) else 'ln -s'
+ # Copy the headers to create a sandboxable setup.
+ cmd = 'cp -f'
command.append(cmd + ' "%s" "%s"' % (src_files[i] , dest))
outs.append(' "' + dest_dir + dest_files[i] + '",')
genrule = _genrule(src_dir, genrule_name, " && ".join(command),
diff --git a/third_party/systemlibs/jemalloc.BUILD b/third_party/systemlibs/jemalloc.BUILD
deleted file mode 100644
index 6a48d58..0000000
--- a/third_party/systemlibs/jemalloc.BUILD
+++ /dev/null
@@ -1,30 +0,0 @@
-licenses(["notice"]) # BSD
-
-filegroup(
- name = "COPYING",
- visibility = ["//visibility:public"],
-)
-
-cc_library(
- name = "jemalloc_headers",
- defines = [
- "jemalloc_posix_memalign=posix_memalign",
- "jemalloc_malloc=malloc",
- "jemalloc_realloc=realloc",
- "jemalloc_free=free",
- ],
- visibility = ["//visibility:public"],
-)
-
-cc_library(
- name = "jemalloc_impl",
- linkopts = ["-ljemalloc"],
- defines = [
- "jemalloc_posix_memalign=posix_memalign",
- "jemalloc_malloc=malloc",
- "jemalloc_realloc=realloc",
- "jemalloc_free=free",
- ],
- visibility = ["//visibility:public"],
- deps = [":jemalloc_headers"],
-)
diff --git a/third_party/systemlibs/syslibs_configure.bzl b/third_party/systemlibs/syslibs_configure.bzl
index 8b0ab39..b03d338 100644
--- a/third_party/systemlibs/syslibs_configure.bzl
+++ b/third_party/systemlibs/syslibs_configure.bzl
@@ -23,7 +23,6 @@
"gast_archive",
"gif_archive",
"grpc",
- "jemalloc",
"jpeg",
"jsoncpp_git",
"lmdb",
diff --git a/third_party/toolchains/BUILD b/third_party/toolchains/BUILD
index 7256a7d..bcbc4dd 100644
--- a/third_party/toolchains/BUILD
+++ b/third_party/toolchains/BUILD
@@ -26,12 +26,10 @@
constraint_values = [
"@bazel_tools//platforms:x86_64",
"@bazel_tools//platforms:linux",
- "@bazel_tools//tools/cpp:clang",
- "@bazel_toolchains//constraints:xenial",
],
remote_execution_properties = """
properties: {
name: "container-image"
- value:"docker://gcr.io/asci-toolchain/nosla-cuda9.0-cudnn7-ubuntu14.04@sha256:06b585f42eed3b2030e9566b8f88f48d7472fa0f47e59765bc115376c8801bdf"
+ value:"docker://gcr.io/asci-toolchain/nosla-cuda9.0-cudnn7-ubuntu14.04@sha256:e5099ff15650986e268a43ee99e2d2b7ffe2459b8b6935385078d1d3b2ed4d02"
}""",
)
diff --git a/third_party/toolchains/preconfig/ubuntu14.04/cuda9.0-cudnn7/cuda/BUILD b/third_party/toolchains/preconfig/ubuntu14.04/cuda9.0-cudnn7/cuda/BUILD
index 2d3e411..05abcb5 100755
--- a/third_party/toolchains/preconfig/ubuntu14.04/cuda9.0-cudnn7/cuda/BUILD
+++ b/third_party/toolchains/preconfig/ubuntu14.04/cuda9.0-cudnn7/cuda/BUILD
@@ -1253,7 +1253,7 @@
"cuda/lib/libcupti.so.9.0",
],
cmd = """
-if [ -d "$(@D)/extras" ]; then rm $(@D)/extras -drf; fi && if [ -d "$(@D)/include" ]; then rm $(@D)/include -drf; fi && if [ -d "$(@D)/lib" ]; then rm $(@D)/lib -drf; fi && if [ -d "$(@D)/nvvm" ]; then rm $(@D)/nvvm -drf; fi && cp "/usr/local/cuda-9.0/targets/x86_64-linux/lib/stubs/libcuda.so" "$(@D)/cuda/lib/libcuda.so" && cp "/usr/local/cuda-9.0/targets/x86_64-linux/lib/libcudart.so.9.0.176" "$(@D)/cuda/lib/libcudart.so.9.0" && cp "/usr/local/cuda-9.0/targets/x86_64-linux/lib/libcudart_static.a" "$(@D)/cuda/lib/libcudart_static.a" && cp "/usr/local/cuda-9.0/targets/x86_64-linux/lib/libcublas.so.9.0.480" "$(@D)/cuda/lib/libcublas.so.9.0" && cp "/usr/local/cuda-9.0/targets/x86_64-linux/lib/libcusolver.so.9.0.176" "$(@D)/cuda/lib/libcusolver.so.9.0" && cp "/usr/local/cuda-9.0/targets/x86_64-linux/lib/libcurand.so.9.0.176" "$(@D)/cuda/lib/libcurand.so.9.0" && cp "/usr/local/cuda-9.0/targets/x86_64-linux/lib/libcufft.so.9.0.176" "$(@D)/cuda/lib/libcufft.so.9.0" && cp "/usr/lib/x86_64-linux-gnu/libcudnn.so.7.2.1" "$(@D)/cuda/lib/libcudnn.so.7" && cp "/usr/local/cuda-9.0/extras/CUPTI/lib64/libcupti.so.9.0.176" "$(@D)/cuda/lib/libcupti.so.9.0"
+if [ -d "$(@D)/extras" ]; then rm $(@D)/extras -drf; fi && if [ -d "$(@D)/include" ]; then rm $(@D)/include -drf; fi && if [ -d "$(@D)/lib" ]; then rm $(@D)/lib -drf; fi && if [ -d "$(@D)/nvvm" ]; then rm $(@D)/nvvm -drf; fi && cp "/usr/local/cuda-9.0/targets/x86_64-linux/lib/stubs/libcuda.so" "$(@D)/cuda/lib/libcuda.so" && cp "/usr/local/cuda-9.0/targets/x86_64-linux/lib/libcudart.so.9.0.176" "$(@D)/cuda/lib/libcudart.so.9.0" && cp "/usr/local/cuda-9.0/targets/x86_64-linux/lib/libcudart_static.a" "$(@D)/cuda/lib/libcudart_static.a" && cp "/usr/local/cuda-9.0/targets/x86_64-linux/lib/libcublas.so.9.0.480" "$(@D)/cuda/lib/libcublas.so.9.0" && cp "/usr/local/cuda-9.0/targets/x86_64-linux/lib/libcusolver.so.9.0.176" "$(@D)/cuda/lib/libcusolver.so.9.0" && cp "/usr/local/cuda-9.0/targets/x86_64-linux/lib/libcurand.so.9.0.176" "$(@D)/cuda/lib/libcurand.so.9.0" && cp "/usr/local/cuda-9.0/targets/x86_64-linux/lib/libcufft.so.9.0.176" "$(@D)/cuda/lib/libcufft.so.9.0" && cp "/usr/lib/x86_64-linux-gnu/libcudnn.so.7.1.4" "$(@D)/cuda/lib/libcudnn.so.7" && cp "/usr/local/cuda-9.0/extras/CUPTI/lib64/libcupti.so.9.0.176" "$(@D)/cuda/lib/libcupti.so.9.0"
""",
)
diff --git a/third_party/toolchains/preconfig/ubuntu14.04/gcc-nvcc/BUILD b/third_party/toolchains/preconfig/ubuntu14.04/gcc-nvcc/BUILD
index a56b451..6442e76 100755
--- a/third_party/toolchains/preconfig/ubuntu14.04/gcc-nvcc/BUILD
+++ b/third_party/toolchains/preconfig/ubuntu14.04/gcc-nvcc/BUILD
@@ -2,6 +2,20 @@
package(default_visibility = ["//visibility:public"])
+toolchain(
+ name = "toolchain-linux-x86_64",
+ exec_compatible_with = [
+ "@bazel_tools//platforms:linux",
+ "@bazel_tools//platforms:x86_64",
+ ],
+ target_compatible_with = [
+ "@bazel_tools//platforms:linux",
+ "@bazel_tools//platforms:x86_64",
+ ],
+ toolchain = ":cc-compiler-local",
+ toolchain_type = "@bazel_tools//tools/cpp:toolchain_type",
+)
+
cc_toolchain_suite(
name = "toolchain",
toolchains = {
diff --git a/tools/bazel.rc b/tools/bazel.rc
index 6747c7e..3734fab 100644
--- a/tools/bazel.rc
+++ b/tools/bazel.rc
@@ -43,6 +43,9 @@
build:cuda --crosstool_top=@local_config_cuda//crosstool:toolchain
build:cuda --define=using_cuda=true --define=using_cuda_nvcc=true
+build:rocm --crosstool_top=@local_config_rocm//crosstool:toolchain
+build:rocm --define=using_rocm=true --define=using_rocm_hipcc=true
+
build:cuda_clang --crosstool_top=@local_config_cuda//crosstool:toolchain
build:cuda_clang --define=using_cuda=true --define=using_cuda_clang=true --define=using_clang=true
@@ -58,6 +61,11 @@
build:sycl_trisycl --crosstool_top=@local_config_sycl//crosstool:toolchain
build:sycl_trisycl --define=using_sycl=true --define=using_trisycl=true
+# Options extracted from configure script
+build:gdr --define=with_gdr_support=true
+build:ngraph --define=with_ngraph_support=true
+build:verbs --define=with_verbs_support=true
+
build --define=use_fast_cpp_protos=true
build --define=allow_oversize_protos=true
build --define=grpc_no_ares=true
@@ -66,6 +74,9 @@
build --genrule_strategy=standalone
build -c opt
+# Other build flags.
+build --define=grpc_no_ares=true
+
# Modular TF build options
build:dynamic_kernels --define=dynamic_loaded_kernels=true