wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 3 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/python/ops/parallel_for/BUILD
+337
View File
@@ -0,0 +1,337 @@
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:tensorflow.bzl", "py_test")
load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//tensorflow:internal",
],
licenses = ["notice"],
)
py_library(
name = "parallel_for",
srcs = [
"__init__.py",
],
strict_deps = True,
deps = [
":control_flow_ops",
":gradients",
],
)
py_library(
name = "pfor_lib",
srcs = ["pfor.py"],
strict_deps = True,
deps = [
"//tensorflow/compiler/tf2xla/python:xla",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/eager:execute",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:func_graph",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:smart_cond",
"//tensorflow/python/framework:sparse_tensor",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/framework:tensor_util",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_gen",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:cond",
"//tensorflow/python/ops:control_flow_assert",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/ops:control_flow_switch_case",
"//tensorflow/python/ops:data_flow_ops",
"//tensorflow/python/ops:handle_data_util",
"//tensorflow/python/ops:image_ops_gen",
"//tensorflow/python/ops:linalg_ops",
"//tensorflow/python/ops:linalg_ops_gen",
"//tensorflow/python/ops:list_ops",
"//tensorflow/python/ops:list_ops_gen",
"//tensorflow/python/ops:manip_ops",
"//tensorflow/python/ops:map_fn",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:math_ops_gen",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/ops:nn_ops_gen",
"//tensorflow/python/ops:optional_ops_gen",
"//tensorflow/python/ops:parsing_ops_gen",
"//tensorflow/python/ops:random_ops_gen",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:sparse_ops",
"//tensorflow/python/ops:sparse_ops_gen",
"//tensorflow/python/ops:special_math_ops",
"//tensorflow/python/ops:spectral_ops_gen",
"//tensorflow/python/ops:tensor_array_ops",
"//tensorflow/python/ops:while_loop",
"//tensorflow/python/platform:flags",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/util:compat",
"//tensorflow/python/util:nest",
"//tensorflow/python/util:numpy_compat",
"//tensorflow/python/util:object_identity",
],
)
py_library(
name = "control_flow_ops",
srcs = ["control_flow_ops.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
":pfor_lib",
"//tensorflow/python/autograph/core:ag_ctx",
"//tensorflow/python/autograph/impl:api",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:composite_tensor",
"//tensorflow/python/framework:indexed_slices",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:sparse_tensor",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/framework:tensor_util",
"//tensorflow/python/framework:type_spec",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:cond",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:tensor_array_ops",
"//tensorflow/python/ops:while_loop",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/util:nest",
"//tensorflow/python/util:tf_decorator_py",
"//tensorflow/python/util:tf_export",
"//tensorflow/python/util:tf_inspect",
"//tensorflow/python/util:variable_utils",
],
)
py_library(
name = "test_util",
srcs = ["test_util.py"],
strict_deps = True,
deps = [
":control_flow_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/util:nest",
"//third_party/py/numpy",
],
)
cuda_py_strict_test(
name = "control_flow_ops_test",
srcs = ["control_flow_ops_test.py"],
shard_count = 16,
tags = [
"no_oss",
"no_rocm",
],
deps = [
":control_flow_ops",
":test_util",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/client:session",
"//tensorflow/python/eager:backprop",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:composite_tensor",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:importer",
"//tensorflow/python/framework:indexed_slices",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:sparse_tensor",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/framework:type_spec",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:bitwise_ops",
"//tensorflow/python/ops:cond",
"//tensorflow/python/ops:cond_v2",
"//tensorflow/python/ops:control_flow_assert",
"//tensorflow/python/ops:control_flow_v2_toggles",
"//tensorflow/python/ops:data_flow_ops",
"//tensorflow/python/ops:functional_ops",
"//tensorflow/python/ops:gradient_checker_v2",
"//tensorflow/python/ops:gradients",
"//tensorflow/python/ops:image_ops",
"//tensorflow/python/ops:list_ops",
"//tensorflow/python/ops:list_ops_gen",
"//tensorflow/python/ops:logging_ops",
"//tensorflow/python/ops:manip_ops",
"//tensorflow/python/ops:map_fn",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nn",
"//tensorflow/python/ops:nn_ops_gen",
"//tensorflow/python/ops:optional_ops_gen",
"//tensorflow/python/ops:parsing_ops",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:rnn_cell",
"//tensorflow/python/ops:stateless_random_ops",
"//tensorflow/python/ops:tensor_array_grad",
"//tensorflow/python/ops:tensor_array_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/ops:while_loop",
"//tensorflow/python/ops/ragged:ragged_tensor",
"//tensorflow/python/ops/signal:fft_ops",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/util:nest",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
cuda_py_strict_test(
name = "xla_control_flow_ops_test",
srcs = ["xla_control_flow_ops_test.py"],
tags = [
# XLA is not enabled by default on Mac or Windows.
"no_mac",
"no_windows",
"no_gpu", # TODO(b/155761551): Flaky on GPU on TAP
"no_tfrt", # Note: Legacy XLA test, which depends on EncapsulateXlaComputationsPass.
"notsan", # TODO(b/271486347): Fails TSAN on TAP
],
xla_enabled = True,
deps = [
":control_flow_ops",
":test_util",
"//tensorflow/compiler/tf2xla/python:xla",
"//tensorflow/python/compiler/xla",
"//tensorflow/python/compiler/xla:compiler_py",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:control_flow_v2_toggles",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:while_loop",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "array_test",
srcs = ["array_test.py"],
tags = [
"nogpu", # b/217374776
],
deps = [
":control_flow_ops",
":test_util",
"//tensorflow/python/eager:backprop",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nn",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/ops:tensor_array_grad",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "math_test",
srcs = ["math_test.py"],
shard_count = 5,
tags = ["optonly"], # Too slow in non-opt mode
deps = [
":control_flow_ops",
":test_util",
"//tensorflow/python/eager:backprop",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:clip_ops",
"//tensorflow/python/ops:linalg_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nn",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/ops:special_math_ops",
"//tensorflow/python/ops:tensor_array_grad",
"//tensorflow/python/platform:client_testlib",
"@absl_py//absl/testing:parameterized",
],
)
py_library(
name = "gradients",
srcs = ["gradients.py"],
strict_deps = True,
deps = [
":control_flow_ops",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:check_ops",
"//tensorflow/python/ops:gradients_impl",
"//tensorflow/python/util:nest",
],
)
cuda_py_strict_test(
name = "gradients_test",
srcs = ["gradients_test.py"],
tags = ["optonly"], # Too slow in non-opt mode
deps = [
":control_flow_ops",
":gradients",
"//tensorflow/python/client:session",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/layers",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:functional_ops",
"//tensorflow/python/ops:gradients",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nn",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/ops:rnn_cell",
"//tensorflow/python/ops:variables",
"//tensorflow/python/ops:while_loop",
"//tensorflow/python/ops/losses",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/util:nest",
"//third_party/py/numpy",
],
)
py_test(
name = "pfor_test",
srcs = ["pfor_test.py"],
strict_deps = True,
deps = [
":pfor_lib",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/platform:client_testlib",
],
)
tensorflow/python/ops/parallel_for/__init__.py
+15
View File
@@ -0,0 +1,15 @@
# 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.
# ==============================================================================
"""Ops for pfor, for_loop, jacobian."""
tensorflow/python/ops/parallel_for/array_test.py
+572
View File
@@ -0,0 +1,572 @@
# 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 vectorization of array kernels."""
from tensorflow.python.eager import backprop
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import array_ops_stack
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn
from tensorflow.python.ops import random_ops
from tensorflow.python.ops import tensor_array_grad # pylint: disable=unused-import
from tensorflow.python.ops.parallel_for import control_flow_ops as pfor_control_flow_ops
from tensorflow.python.ops.parallel_for.test_util import PForTestCase
from tensorflow.python.platform import test
@test_util.with_eager_op_as_function
@test_util.run_all_in_graph_and_eager_modes
class ArrayTest(PForTestCase):
def test_gather(self):
x = random_ops.random_uniform([3, 3, 3, 3])
x2 = array_ops.placeholder_with_default(x, shape=None) # Has dynamic shape.
def loop_fn(i):
outputs = []
x_i = array_ops.gather(x, i)
for y in [x, x2, x_i]:
for axis in [0, 2, -1]:
outputs.append(array_ops.gather(y, 2, axis=axis))
outputs.append(
array_ops.gather(y, math_ops.cast(2, dtypes.int64), axis=axis))
outputs.append(
array_ops.gather(y, 2, axis=math_ops.cast(axis, dtypes.int64)))
outputs.append(
array_ops.gather(y, math_ops.cast(i, dtypes.int64), axis=axis))
outputs.append(array_ops.gather(y, [i], axis=axis))
outputs.append(array_ops.gather(y, [i, 2], axis=axis))
outputs.append(array_ops.gather(y, [[2, i], [i, 1]], axis=axis))
outputs.append(array_ops.gather(y, [0, 1, 2], axis=1, batch_dims=1))
outputs.append(array_ops.gather(y, [i, 1, 2], axis=2, batch_dims=1))
outputs.append(array_ops.gather(y, [[2, i], [i, 1], [2, 1]],
axis=-1, batch_dims=1))
outputs.append(
array_ops.gather(y, [[0, 1, 2]] * 3, axis=2, batch_dims=2))
outputs.append(array_ops.gather(y, [0, 1, 2], axis=1, batch_dims=-1))
outputs.append(
array_ops.gather(y, [[0, 1, 2]] * 3, axis=2, batch_dims=-2))
return outputs
self._test_loop_fn(loop_fn, 3)
def test_gather_nd(self):
x = random_ops.random_uniform([3, 3, 3])
def loop_fn(i):
outputs = []
x_i = array_ops.gather(x, i)
outputs.append(array_ops.gather_nd(x_i, [0], batch_dims=0))
outputs.append(array_ops.gather_nd(x_i, [i], batch_dims=0))
outputs.append(array_ops.gather_nd(x_i, [[i], [i], [i]], batch_dims=1))
return outputs
self._test_loop_fn(loop_fn, 3)
@test_util.run_v2_only
def test_gather_pfor_grad(self):
x = array_ops.zeros([1, 2])
with backprop.GradientTape() as tape:
tape.watch(x)
r = pfor_control_flow_ops.vectorized_map(
lambda t: array_ops.gather(x, t, axis=-1), math_ops.range(2))
self.assertAllClose([[1., 1.]], tape.gradient(r, x))
def test_shape(self):
x = random_ops.random_uniform([3, 2, 3])
def loop_fn(i):
x_i = array_ops.gather(x, i)
return array_ops.shape(x_i), array_ops.shape(x_i, out_type=dtypes.int64)
self._test_loop_fn(loop_fn, 3)
def test_size(self):
x = random_ops.random_uniform([3, 2, 3])
def loop_fn(i):
x_i = array_ops.gather(x, i)
return array_ops.size(x_i), array_ops.size(x_i, out_type=dtypes.int64)
self._test_loop_fn(loop_fn, 3)
def test_rank(self):
x = random_ops.random_uniform([3, 2, 3])
def loop_fn(i):
x_i = array_ops.gather(x, i)
return array_ops.rank(x_i)
self._test_loop_fn(loop_fn, 3)
def test_shape_n(self):
x = random_ops.random_uniform([3, 2, 3])
y = random_ops.random_uniform([3])
def loop_fn(i):
x_i = array_ops.gather(x, i)
y_i = array_ops.gather(y, i)
return array_ops.shape_n([x_i, x, y,
y_i]), array_ops.shape_n([x_i, x, y, y_i],
out_type=dtypes.int64)
self._test_loop_fn(loop_fn, 3)
def test_reshape(self):
x = random_ops.random_uniform([3, 2, 3])
def loop_fn(i):
x1 = array_ops.gather(x, i)
return array_ops.reshape(x1, [-1]), array_ops.reshape(x1, [1, 3, 1, -1])
self._test_loop_fn(loop_fn, 3)
def test_fill(self):
def loop_fn(i):
return array_ops.fill((2, 3), i)
self._test_loop_fn(loop_fn, 3)
def test_broadcast_to(self):
x = random_ops.random_uniform([3, 2, 1, 3])
def loop_fn(i):
x1 = array_ops.gather(x, i)
return (array_ops.broadcast_to(x1, [2, 2, 3]),
array_ops.broadcast_to(x1, [1, 2, 1, 3]))
self._test_loop_fn(loop_fn, 3)
def test_expand_dims(self):
x = random_ops.random_uniform([3, 2, 3])
def loop_fn(i):
x1 = array_ops.gather(x, i)
return [
array_ops.expand_dims(x1, axis=-1),
array_ops.expand_dims(x1, axis=1),
array_ops.expand_dims(
x1, axis=constant_op.constant(1, dtype=dtypes.int64))
]
self._test_loop_fn(loop_fn, 3)
def test_one_hot(self):
indices = random_ops.random_uniform([3, 2, 3],
minval=0,
maxval=4,
dtype=dtypes.int32)
def loop_fn(i):
indices_i = array_ops.gather(indices, i)
return (array_ops.one_hot(indices_i, depth=4, on_value=2., off_value=-2.),
array_ops.one_hot(indices_i, depth=4, axis=1))
self._test_loop_fn(loop_fn, 3)
def test_searchsorted(self):
sorted_inputs = math_ops.cumsum(
random_ops.random_uniform([3, 2, 4]), axis=-1)
values = random_ops.random_uniform([2, 3], minval=-1, maxval=4.5)
def loop_fn(i):
inputs_i = array_ops.gather(sorted_inputs, i)
return [
array_ops.searchsorted(
inputs_i, values, out_type=dtypes.int32,
side="left"), # creates LowerBound op.
array_ops.searchsorted(
inputs_i, values, out_type=dtypes.int64, side="right")
] # creates UpperBound op.
self._test_loop_fn(loop_fn, 3)
def test_slice(self):
x = random_ops.random_uniform([3, 2, 3])
def loop_fn(i):
x1 = array_ops.gather(x, i)
return array_ops.slice(x1, begin=(0, 1), size=(2, 1))
self._test_loop_fn(loop_fn, 3)
def test_slice_loop_variant_begin(self):
x = random_ops.random_uniform([3, 2, 5, 3])
def loop_fn(i):
x1 = array_ops.gather(x, i)
return array_ops.slice(x1, begin=(0, 2 - i, i), size=(-1, 2, 1))
self._test_loop_fn(loop_fn, 3)
def test_tile(self):
x = random_ops.random_uniform([3, 2, 3])
def loop_fn(i):
x1 = array_ops.gather(x, i)
return array_ops.tile(x1, [2, 1])
self._test_loop_fn(loop_fn, 3)
def test_tile_loop_dependent(self):
x = random_ops.random_uniform([3, 2, 3])
def loop_fn(i):
x1 = array_ops.gather(x, i)
return array_ops.tile(x1, [i, 1])
with self.assertRaisesRegex(ValueError, "expected to be loop invariant"):
pfor_control_flow_ops.pfor(loop_fn, 2, fallback_to_while_loop=False)
def test_pack(self):
x = random_ops.random_uniform([3, 2, 3])
y = random_ops.random_uniform([2, 3])
def loop_fn(i):
x1 = array_ops.gather(x, i)
return array_ops_stack.stack([x1, y], axis=-1)
self._test_loop_fn(loop_fn, 1)
def test_unpack(self):
x = random_ops.random_uniform([3, 2, 3, 4])
def loop_fn(i):
x_i = array_ops.gather(x, i)
return array_ops_stack.unstack(
x_i, 4, axis=-1), array_ops_stack.unstack(
x_i, 3, axis=1)
self._test_loop_fn(loop_fn, 3)
def test_pad(self):
x = random_ops.random_uniform([3, 2, 3])
padding = constant_op.constant([[1, 2], [3, 4]])
def loop_fn(i):
x1 = array_ops.gather(x, i)
return array_ops.pad(x1, padding, mode="CONSTANT")
self._test_loop_fn(loop_fn, 3)
def test_pad_v2(self):
x = random_ops.random_uniform([3, 2, 3])
padding = constant_op.constant([[1, 2], [3, 4]])
def loop_fn(i):
x1 = array_ops.gather(x, i)
return array_ops.pad_v2(x1, padding, mode="CONSTANT", constant_values=1.5)
self._test_loop_fn(loop_fn, 3)
def test_split(self):
x = random_ops.random_uniform([3, 2, 3])
def loop_fn(i):
x1 = array_ops.gather(x, i)
return array_ops.split(x1, 2, axis=0), array_ops.split(x1, 3, axis=-1)
self._test_loop_fn(loop_fn, 3)
def test_split_v(self):
x = random_ops.random_uniform([3, 6, 3])
def loop_fn(i):
x1 = array_ops.gather(x, i)
return (array_ops.split(x1, [2, 1, 3],
axis=0), array_ops.split(x1, [3], axis=-1))
self._test_loop_fn(loop_fn, 3)
def test_squeeze(self):
x = random_ops.random_uniform([5, 1, 2, 1])
def loop_fn(i):
x1 = array_ops.gather(x, i)
return (array_ops.squeeze(x1, axis=0), array_ops.squeeze(x1, axis=-1),
array_ops.squeeze(x1))
self._test_loop_fn(loop_fn, 3)
def test_reverse(self):
x = random_ops.random_uniform([3, 4, 2, 3])
def loop_fn(i):
x1 = array_ops.gather(x, i)
return (array_ops.reverse(x1, axis=[0]),
array_ops.reverse(x1, axis=[-1]),
array_ops.reverse(x1, axis=[1, -1]))
self._test_loop_fn(loop_fn, 3)
def test_transpose(self):
x = random_ops.random_uniform([3, 2, 3, 4])
def loop_fn(i):
x1 = array_ops.gather(x, i)
return array_ops.transpose(x1, [2, 1, 0])
self._test_loop_fn(loop_fn, 3)
def test_top_k(self):
x = random_ops.random_uniform([3, 2, 3, 4])
def loop_fn(i):
x1 = array_ops.gather(x, i)
return nn.top_k(x1, k=2)
self._test_loop_fn(loop_fn, 3)
def test_conjugate_transpose(self):
x = math_ops.complex(
random_ops.random_uniform([3, 2, 3, 4]),
random_ops.random_uniform([3, 2, 3, 4]))
def loop_fn(i):
x_i = array_ops.gather(x, i)
return array_ops.conjugate_transpose(x_i, [2, 1, 0])
self._test_loop_fn(loop_fn, 3)
def test_zeros_like(self):
x = random_ops.random_uniform([3, 2, 3])
def loop_fn(i):
x1 = array_ops.gather(x, i)
z = array_ops.zeros_like(x1),
return z, z + x1
self._test_loop_fn(loop_fn, 3)
def test_ones_like(self):
x = random_ops.random_uniform([3, 2, 3])
def loop_fn(i):
x1 = array_ops.gather(x, i)
z = array_ops.ones_like(x1),
return z, z + x1
self._test_loop_fn(loop_fn, 3)
def test_concat_v2(self):
x = random_ops.random_uniform([3, 2, 3])
y = random_ops.random_uniform([2, 3])
def loop_fn(i):
x1 = array_ops.gather(x, i)
return [
array_ops.concat([x1, x1, y], axis=0),
array_ops.concat([x1, x1, y], axis=-1),
array_ops.concat([x1, x1, y],
axis=constant_op.constant(0, dtype=dtypes.int64))
]
self._test_loop_fn(loop_fn, 3)
def test_unary_cwise_ops(self):
for op in [array_ops.identity, array_ops.stop_gradient]:
with backprop.GradientTape(persistent=True) as g:
x = random_ops.random_uniform([3, 5])
g.watch(x)
# pylint: disable=cell-var-from-loop
def loop_fn(i):
with g:
x1 = array_ops.gather(x, i)
y = op(x1) + x1
loss = nn.l2_loss(y)
return op(x), y, g.gradient(loss, x1)
# pylint: enable=cell-var-from-loop
self._test_loop_fn(loop_fn, 3)
def test_identity_n(self):
x = random_ops.random_uniform([3, 4])
def loop_fn(i):
return array_ops.identity_n([x, array_ops.gather(x, i)])
self._test_loop_fn(loop_fn, 3)
def test_matrix_band_part(self):
x = random_ops.random_uniform([3, 4, 2, 2])
for num_lower, num_upper in ((0, -1), (-1, 0), (1, 1)):
# pylint: disable=cell-var-from-loop
def loop_fn(i):
return array_ops.matrix_band_part(
array_ops.gather(x, i), num_lower=num_lower, num_upper=num_upper)
# pylint: enable=cell-var-from-loop
self._test_loop_fn(loop_fn, 3)
def test_matrix_diag(self):
x = random_ops.random_uniform([3, 2, 4])
def loop_fn(i):
diagonal = array_ops.gather(x, i)
return array_ops.matrix_diag(
diagonal, k=(0, 1), num_rows=4, num_cols=5, align="RIGHT_LEFT")
self._test_loop_fn(loop_fn, 3)
def test_matrix_diag_part(self):
x = random_ops.random_uniform([3, 4, 6])
def loop_fn(i):
input = array_ops.gather(x, i) # pylint: disable=redefined-builtin
return array_ops.matrix_diag_part(
input, k=(-2, 0), padding_value=3, align="RIGHT_LEFT")
self._test_loop_fn(loop_fn, 3)
def test_diag(self):
for x in (random_ops.random_uniform([3, 4]),
random_ops.random_uniform([3, 4, 2])):
# pylint: disable=cell-var-from-loop
def loop_fn(i):
inp = array_ops.gather(x, i)
return array_ops.diag(inp)
# pylint: disable=cell-var-from-loop
self._test_loop_fn(loop_fn, 3)
def test_diag_part(self):
for x in (random_ops.random_uniform([3, 2, 2]),
random_ops.random_uniform([3, 4, 2, 4, 2])):
# pylint: disable=cell-var-from-loop
def loop_fn(i):
inp = array_ops.gather(x, i) # pylint: disable=redefined-builtin
return array_ops.diag_part(inp)
# pylint: disable=cell-var-from-loop
self._test_loop_fn(loop_fn, 3)
def test_matrix_set_diag(self):
matrices = random_ops.random_uniform([3, 4, 4])
diags = random_ops.random_uniform([3, 4])
bands = random_ops.random_uniform([3, 3, 4])
def loop_fn(i):
matrix_i = array_ops.gather(matrices, i)
diag_i = array_ops.gather(diags, i)
results = [
array_ops.matrix_set_diag(matrix_i, diag_i),
array_ops.matrix_set_diag(matrices[0, ...], diag_i),
array_ops.matrix_set_diag(matrix_i, diags[0, ...]),
]
k = (-1, 1)
band_i = array_ops.gather(bands, i)
for align in ["RIGHT_LEFT", "LEFT_RIGHT"]:
results.extend([
array_ops.matrix_set_diag(matrix_i, band_i, k=k, align=align),
array_ops.matrix_set_diag(
matrices[0, ...], band_i, k=k, align=align),
array_ops.matrix_set_diag(
matrix_i, bands[0, ...], k=k, align=align)
])
return results
self._test_loop_fn(loop_fn, 3)
def test_strided_slice(self):
with backprop.GradientTape(persistent=True) as g:
x = random_ops.random_uniform([3, 3, 4, 4, 2, 2, 2])
g.watch(x)
def loop_fn(i):
with g:
x_i = array_ops.gather(x, i)
y = x_i[:2, ::2, 1::3, ..., array_ops.newaxis, 1]
loss = nn.l2_loss(y)
return y, g.gradient(loss, x_i)
self._test_loop_fn(loop_fn, 3)
def test_strided_slice_loop_variant(self):
x = random_ops.random_uniform([3, 3, 4, 4, 2, 2, 2])
def loop_fn(i):
x_i = array_ops.gather(x, i)
return x_i[i:i+1, ...]
# Test the fallback to while loop for a ConversionNotImplementedError is
# handled.
self._test_loop_fn(loop_fn, 3, fallback_to_while_loop=True)
# Without fallback, ValueError is thrown.
with self.assertRaisesRegex(ValueError, "expected to be loop invariant"):
self._test_loop_fn(loop_fn, 3, fallback_to_while_loop=False)
def test_depth_to_space(self):
x = random_ops.random_uniform([2, 3, 2, 2, 12])
def loop_fn(i):
x1 = array_ops.gather(x, i)
return array_ops.depth_to_space(x1, 2, data_format="NHWC")
self._test_loop_fn(loop_fn, 2)
def test_space_to_depth(self):
x = random_ops.random_uniform([2, 3, 12, 12, 3])
def loop_fn(i):
x1 = array_ops.gather(x, i)
return array_ops.space_to_depth(x1, 2, data_format="NHWC")
self._test_loop_fn(loop_fn, 2)
def test_batch_to_space_nd(self):
x = random_ops.random_uniform([7, 5 * 2 * 3, 2, 2, 3, 2])
block_shapes = [2, 3]
crops = [[1, 2], [1, 0]]
def loop_fn(i):
x1 = array_ops.gather(x, i)
return array_ops.batch_to_space_nd(x1, block_shapes, crops)
self._test_loop_fn(loop_fn, 7)
def test_space_to_batch_nd(self):
x = random_ops.random_uniform([7, 5, 2 * 2 - 3, 2 * 3 - 1, 3, 2])
block_shapes = [2, 3]
paddings = [[1, 2], [1, 0]]
def loop_fn(i):
x1 = array_ops.gather(x, i)
return array_ops.space_to_batch_nd(x1, block_shapes, paddings)
self._test_loop_fn(loop_fn, 7)
def test_check_numerics(self):
x = random_ops.random_uniform([2, 3, 4])
def loop_fn(i):
x_i = array_ops.gather(x, i)
return array_ops.check_numerics(x_i, "test_message")
self._test_loop_fn(loop_fn, 2)
if __name__ == "__main__":
test.main()
tensorflow/python/ops/parallel_for/control_flow_ops.py
+582
View File
@@ -0,0 +1,582 @@
# 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.
# ==============================================================================
"""for_loop and pfor ops."""
# pylint: disable=g-direct-tensorflow-import
import functools
from tensorflow.python.eager import context
from tensorflow.python.eager import def_function
from tensorflow.python.autograph.core import ag_ctx as autograph_ctx
from tensorflow.python.autograph.impl import api as autograph
from tensorflow.python.framework import composite_tensor
from tensorflow.python.framework import indexed_slices
from tensorflow.python.framework import ops
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.framework import tensor
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import tensor_util
from tensorflow.python.framework import type_spec
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import cond
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import tensor_array_ops
from tensorflow.python.ops import while_loop
from tensorflow.python.ops.parallel_for.pfor import PFor
from tensorflow.python.ops.parallel_for.pfor import PForConfig
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.util import nest
from tensorflow.python.util import tf_decorator
from tensorflow.python.util import tf_inspect
from tensorflow.python.util import variable_utils
from tensorflow.python.util.tf_export import tf_export
def for_loop(loop_fn, loop_fn_dtypes, iters, parallel_iterations=None):
"""Runs `loop_fn` `iters` times and stacks the outputs.
Runs `loop_fn` `iters` times, with input values from 0 to `iters - 1`, and
stacks corresponding outputs of the different runs.
Args:
loop_fn: A function that takes an int32 scalar tf.Tensor object representing
the iteration number, and returns a possibly nested structure of tensor
objects. The shape of these outputs should not depend on the input.
loop_fn_dtypes: dtypes for the outputs of `loop_fn`.
iters: Number of iterations for which to run `loop_fn`.
parallel_iterations: The number of iterations that can be dispatched in
parallel. This knob can be used to control the total memory usage.
Returns:
Returns a nested structure of stacked output tensor objects with the same
nested structure as the output of `loop_fn`.
"""
flat_loop_fn_dtypes = nest.flatten(loop_fn_dtypes)
is_none_list = []
def while_body(i, *ta_list):
"""Body of while loop."""
fn_conv = autograph.tf_convert(loop_fn, autograph_ctx.control_status_ctx())
fn_output = nest.flatten(fn_conv(i))
if len(fn_output) != len(flat_loop_fn_dtypes):
raise ValueError(
f"Number of expected outputs {len(flat_loop_fn_dtypes)}, does not "
f"match the number of actual outputs {len(fn_output)} from loop_fn: "
f"{loop_fn} with output {fn_output}.")
outputs = []
del is_none_list[:]
is_none_list.extend(x is None for x in fn_output)
for out, ta in zip(fn_output, ta_list):
# TODO(agarwal): support returning Operation objects from loop_fn.
if out is not None:
# out may be a ref tensor, wrap it in identity to get a non-ref tensor.
ta = ta.write(i, out)
outputs.append(ta)
return tuple([i + 1] + outputs)
if parallel_iterations is not None:
extra_args = {"parallel_iterations": parallel_iterations}
else:
extra_args = {}
ta_list = while_loop.while_loop(lambda i, *ta: i < iters, while_body, [0] + [
tensor_array_ops.TensorArray(dtype.base_dtype, iters)
for dtype in flat_loop_fn_dtypes
], **extra_args)[1:]
# TODO(rachelim): enable this for sparse tensors
output = [
None if is_none else ta.stack()
for ta, is_none in zip(ta_list, is_none_list)
]
assert len(output) in (0, len(flat_loop_fn_dtypes))
if not output:
# This may happen for the case where iters == 0.
# Pack a list of empty tensors with the proper ranks to match pfor output on 0 iters
loop_var = array_ops.placeholder_with_default(0, shape=[])
try:
loop_fn_out = loop_fn(loop_var)
out_shapes = [
[0] + ops.convert_to_tensor(x).shape
for x in nest.flatten(loop_fn_out)
]
output = [
array_ops.zeros(out_shapes[i], dt)
for i, dt in enumerate(flat_loop_fn_dtypes)
]
except Exception:
output = [array_ops.zeros([0])]
return nest.pack_sequence_as(loop_fn_dtypes, output)
def _flatten_first_two_dims(x):
"""Flattens the first two dimensions of x into a single dimension."""
old_shape = array_ops.shape(x)
new_shape = array_ops.concat([[old_shape[0] * old_shape[1]], old_shape[2:]],
axis=0)
return array_ops.reshape(x, new_shape)
PFOR_CONFIG_ARG = "pfor_config"
def _is_under_xla_context():
"""Check if we are currently inside an XLA compile context."""
g = ops.get_default_graph()
while g is not None:
control_flow_context = g._get_control_flow_context() # pylint: disable=protected-access
while control_flow_context is not None:
if control_flow_context.IsXLAContext():
return True
else:
control_flow_context = control_flow_context.outer_context
# If g is a FuncGraph, get its outer_graph.
g = getattr(g, "outer_graph", None)
return False
def pfor(loop_fn,
iters,
fallback_to_while_loop=True,
parallel_iterations=None,
warn=False):
"""Equivalent to running `loop_fn` `iters` times and stacking the outputs.
`pfor` has functionality similar to `for_loop`, i.e. running `loop_fn` `iters`
times, with input from 0 to `iters - 1`, and stacking corresponding output of
each iteration. However the implementation does not use a `tf.while_loop`.
Instead it adds new operations to the graph that collectively compute the same
value as what running `loop_fn` in a loop would compute.
This is an experimental feature and currently has a lot of limitations:
- There should be no data dependency between the different iterations. For
example, a future iteration should not depend on a value or side-effect of
a previous iteration.
- Stateful kernels may mostly not be supported since these often imply a
data dependency or ordering of the iterations. We do support a limited set
of such stateful kernels though (like RandomFoo, Variable operations like
reads, etc).
- Conversion works only on a limited set of kernels for which a converter
has been registered.
- `loop_fn` has limited support for control flow operations. `tf.cond` in
particular is not supported.
- `loop_fn` should return nested structure of Tensors or Operations. However
if an Operation is returned, it should have zero outputs.
- The shape and dtype of `loop_fn` outputs should not depend on the input
to loop_fn.
Args:
loop_fn: A function that takes an int32 scalar tf.Tensor object representing
the iteration number, and optionally a keyword argument `pfor_config` set
to a PForConfig object. It returns a possibly nested structure of Tensor
or Operation objects. Note that if setting `parallel_iterations` argument
to something other than None, `loop_fn` may be called more than once
during graph construction. So it may need to avoid mutating global state.
iters: Number of iterations for which to run `loop_fn`.
fallback_to_while_loop: If true, on failing to vectorize an operation, pfor
fallbacks to using a `tf.while_loop` to dispatch the iterations.
parallel_iterations: A knob to control how many iterations are vectorized
and dispatched in parallel. The default value of None corresponds to
vectorizing all the iterations. If `parallel_iterations` is smaller than
`iters`, then chunks of at most that many iterations are dispatched in
sequence. This knob can be used to control the total memory usage.
warn: Whether or not to warn when falling back to while loops.
Returns:
Returns a nested structure of stacked tensor objects with the same nested
structure as the output of `loop_fn`.
Raises:
ValueError: If parallel_iterations is not None and not an integer > 1.
"""
def f():
return _pfor_impl(
loop_fn,
iters,
fallback_to_while_loop=fallback_to_while_loop,
parallel_iterations=parallel_iterations,
warn=warn)
# Note that we wrap into a tf.function if in eager execution mode or under
# XLA compilation. The latter is so that we don't compile operations like
# tf.placeholder that are created by the loop body.
functions_run_eagerly = None
if context.executing_eagerly() or _is_under_xla_context():
functions_run_eagerly = def_function.functions_run_eagerly()
if functions_run_eagerly:
logging.warning(
"It looks like tf.function behavior was disabled, perhaps using "
"tf.config.run_functions_eagerly. Vectorization "
"primitives (e.g. tf.vectorized_map) require tf.function to work. "
"These primitives will override the disable.")
def_function.run_functions_eagerly(False)
f = def_function.function(f)
outputs = f()
if functions_run_eagerly is not None:
def_function.run_functions_eagerly(functions_run_eagerly)
return outputs
def _should_expand_composite(value):
return (isinstance(value, composite_tensor.CompositeTensor)
# Leave sparse tensors to be converted by `PFor._convert_sparse`.
and not isinstance(value, sparse_tensor.SparseTensor)
and not isinstance(value, indexed_slices.IndexedSlices))
# pylint: disable=protected-access
def _composite_to_tensors(value, is_batched=False):
"""Converts a CompositeTensor into a list of stackable tensors."""
if _should_expand_composite(value):
spec = value._type_spec
if not isinstance(spec, type_spec.BatchableTypeSpec):
raise ValueError(f"CompositeTensor instance {value} returned from "
"parallel_for or vectorized_map loop body must provide "
f"a `BatchableTypeSpec` (saw: {spec}).")
if is_batched:
return spec._to_batched_tensor_list(value)
return spec._to_tensor_list(value)
return value
# pylint: enable=protected-access
# pylint: disable=protected-access
def _composite_from_tensors(stacked_tensors,
preconverted_value,
batch_size):
"""Converts a list of stacked tensors to a batch CompositeTensor."""
if _should_expand_composite(preconverted_value):
batch_type_spec = preconverted_value._type_spec._batch(batch_size)
return batch_type_spec._from_compatible_tensor_list(stacked_tensors)
return stacked_tensors
# pylint: enable=protected-access
def _loop_fn_has_config(loop_fn):
"""Test if `loop_fn` has a `pfor_config` argument."""
if tf_inspect.isfunction(loop_fn):
argspec = tf_inspect.getargspec(loop_fn)
return PFOR_CONFIG_ARG in argspec.args
elif isinstance(loop_fn, functools.partial):
fn = loop_fn.func
argspec = tf_inspect.getargspec(fn)
return (PFOR_CONFIG_ARG in argspec.args and
PFOR_CONFIG_ARG not in loop_fn.keywords)
else:
loop_class = tf_decorator.unwrap(loop_fn)[1]
if not hasattr(loop_class, "__call__"):
raise ValueError("`loop_fn` object did not have a __call__ method")
argspec = tf_inspect.getargspec(loop_class.__call__)
return PFOR_CONFIG_ARG in argspec.args
def _pfor_impl(loop_fn,
iters,
fallback_to_while_loop,
parallel_iterations=None,
pfor_config=None,
warn=False):
"""Implementation of pfor."""
assert not context.executing_eagerly()
loop_fn_has_config = _loop_fn_has_config(loop_fn)
existing_ops = set(ops.get_default_graph().get_operations())
iters_value = tensor_util.constant_value(iters)
# Run the loop body
with ops.name_scope("loop_body"):
loop_var = array_ops.placeholder_with_default(0, shape=[])
if loop_fn_has_config:
if pfor_config is None:
pfor_config = PForConfig()
pfor_config._set_iters(iters) # pylint: disable=protected-access
loop_fn_outputs = loop_fn(loop_var, **{PFOR_CONFIG_ARG: pfor_config})
else:
assert pfor_config is None
f = autograph.tf_convert(loop_fn, autograph_ctx.control_status_ctx())
loop_fn_outputs = f(loop_var)
loop_fn_output_tensors = nest.map_structure(_composite_to_tensors,
loop_fn_outputs)
# Convert outputs to Tensor if needed.
tmp_loop_fn_outputs = []
for loop_fn_output in nest.flatten(loop_fn_output_tensors):
if (loop_fn_output is not None and not isinstance(
loop_fn_output,
(ops.Operation, tensor.Tensor, sparse_tensor.SparseTensor))):
if isinstance(loop_fn_output, indexed_slices.IndexedSlices):
logging.warn("Converting %s to a dense representation may make it slow."
" Alternatively, output the indices and values of the"
" IndexedSlices separately, and handle the vectorized"
" outputs directly." % loop_fn_output)
loop_fn_output = ops.convert_to_tensor(loop_fn_output)
else:
loop_fn_output = ops.convert_to_tensor(loop_fn_output)
tmp_loop_fn_outputs.append(loop_fn_output)
loop_fn_output_tensors = nest.pack_sequence_as(loop_fn_output_tensors,
tmp_loop_fn_outputs)
new_ops = set(ops.get_default_graph().get_operations()) - existing_ops
iters = ops.convert_to_tensor(iters)
if parallel_iterations is not None:
if parallel_iterations < 1:
raise ValueError(
"Argument `parallel_iterations` must be None or a positive integer. "
f"Received: {parallel_iterations}.")
if parallel_iterations == 1:
raise ValueError(
"Found `parallel_iterations == 1`. Use `for_loop` instead.")
if iters_value is not None and iters_value < parallel_iterations:
parallel_iterations = None
if parallel_iterations is None:
with ops.name_scope("pfor"):
converter = PFor(
loop_var,
iters,
new_ops,
fallback_to_while_loop=fallback_to_while_loop,
pfor_config=pfor_config,
warn=warn)
flattened_output_tensors = []
for loop_fn_output in nest.flatten(loop_fn_output_tensors):
output = converter.convert(loop_fn_output)
flattened_output_tensors.append(output)
else:
if pfor_config is not None and pfor_config._has_reductions(): # pylint: disable=protected-access
raise ValueError("Setting `parallel_iterations` currently unsupported if "
"reductions across iterations are performed.")
num_tiled_iterations = iters // parallel_iterations
num_remaining_iterations = iters % parallel_iterations
# TODO(agarwal): Avoid calling loop_fn twice. Generate the loop body inside
# a tf.function and extract the graph from there to vectorize it.
with ops.name_scope("pfor_untiled"):
converter = PFor(loop_var, num_remaining_iterations, new_ops,
fallback_to_while_loop=fallback_to_while_loop,
pfor_config=pfor_config)
remaining_output_tensors = []
flattened_output_tensors = nest.flatten(loop_fn_output_tensors)
for loop_fn_output in flattened_output_tensors:
output = converter.convert(loop_fn_output)
remaining_output_tensors.append(output)
with ops.name_scope("pfor_tiled"):
loop_fn_dtypes = [ops.convert_to_tensor(x).dtype
for x in flattened_output_tensors]
def tiled_loop_body(j):
offset = j * parallel_iterations + num_remaining_iterations
def tiled_loop_fn(i, pfor_config=None):
if loop_fn_has_config:
loop_fn_outputs = loop_fn(i + offset, pfor_config=pfor_config)
else:
loop_fn_outputs = loop_fn(i + offset)
return nest.flatten(
# Stacking across iterations requires explicit Tensors.
nest.map_structure(_composite_to_tensors, loop_fn_outputs))
return _pfor_impl(
tiled_loop_fn,
parallel_iterations,
fallback_to_while_loop=fallback_to_while_loop,
pfor_config=pfor_config)
tiled_output_tensors = for_loop(
tiled_loop_body, loop_fn_dtypes,
num_tiled_iterations, parallel_iterations=1)
tiled_output_tensors = [
_flatten_first_two_dims(y) for y in tiled_output_tensors]
with ops.name_scope("pfor"):
if iters_value is None or iters_value % parallel_iterations:
output_tensors = cond.cond(
math_ops.equal(num_remaining_iterations, 0),
lambda: tiled_output_tensors,
lambda: [array_ops.concat([x, y], axis=0) # pylint: disable=g-long-lambda
for x, y in zip(remaining_output_tensors,
tiled_output_tensors)])
else:
output_tensors = tiled_output_tensors
flattened_output_tensors = nest.flatten(output_tensors)
for output, original_output in zip(flattened_output_tensors,
nest.flatten(loop_fn_output_tensors)):
# Restore any shape information lost from tiling.
# TODO(b/174254748): this may not be correct for stacked `variant`s.
output.set_shape(
tensor_shape.TensorShape([iters_value]).concatenate(
original_output.shape))
return nest.map_structure_up_to(
loop_fn_outputs,
functools.partial(_composite_from_tensors, batch_size=iters_value),
nest.pack_sequence_as(loop_fn_output_tensors,
flattened_output_tensors),
loop_fn_outputs)
def _broadcasting_gather(x, i):
"""Wrapper for gather that implicitly broadcasts unit dimensions."""
static_first_dim = tensor_shape.dimension_value(x.shape[0])
if static_first_dim == 1:
i = 0
elif static_first_dim is None:
i = array_ops.where_v2(array_ops.shape(x)[0] > 1, i, 0)
result = array_ops.gather(x, i)
return result
# pylint: disable=protected-access
def _gather_from_tensor_or_composite(x, i):
"""Wrapper for gather that handles CompositeTensors."""
if _should_expand_composite(x):
spec = x._type_spec
gathered_tensors = [_broadcasting_gather(t, i)
for t in spec._to_batched_tensor_list(x)]
return spec._unbatch()._from_compatible_tensor_list(gathered_tensors)
return _broadcasting_gather(x, i)
# pylint: enable=protected-access
@tf_export("vectorized_map")
def vectorized_map(fn, elems, fallback_to_while_loop=True, warn=True):
"""Parallel map on the list of tensors unpacked from `elems` on dimension 0.
This method works similar to `tf.map_fn` but is optimized to run much faster,
possibly with a much larger memory footprint. The speedups are obtained by
vectorization (see [Auto-Vectorizing TensorFlow Graphs: Jacobians,
Auto-Batching and Beyond](https://arxiv.org/pdf/1903.04243.pdf)). The idea
behind vectorization is to semantically launch all the invocations of `fn` in
parallel and fuse corresponding operations across all these invocations. This
fusion is done statically at graph generation time and the generated code is
often similar in performance to a manually fused version.
Because `tf.vectorized_map` fully parallelizes the batch, this method will
generally be significantly faster than using `tf.map_fn`, especially in eager
mode. However this is an experimental feature and currently has a lot of
limitations:
- There should be no data dependency between the different semantic
invocations of `fn`, i.e. it should be safe to map the elements of the
inputs in any order.
- Stateful kernels may mostly not be supported since these often imply a
data dependency. We do support a limited set of such stateful kernels
though (like RandomFoo, Variable operations like reads, etc).
- `fn` has limited support for control flow operations.
- `fn` should return nested structure of Tensors or Operations. However
if an Operation is returned, it should have zero outputs.
- The shape and dtype of any intermediate or output tensors in the
computation of `fn` should not depend on the input to `fn`.
Examples:
```python
def outer_product(a):
return tf.tensordot(a, a, 0)
batch_size = 100
a = tf.ones((batch_size, 32, 32))
c = tf.vectorized_map(outer_product, a)
assert c.shape == (batch_size, 32, 32, 32, 32)
```
```python
# Computing per-example gradients
batch_size = 10
num_features = 32
layer = tf.keras.layers.Dense(1)
def model_fn(arg):
with tf.GradientTape() as g:
inp, label = arg
inp = tf.expand_dims(inp, 0)
label = tf.expand_dims(label, 0)
prediction = layer(inp)
loss = tf.nn.l2_loss(label - prediction)
return g.gradient(loss, (layer.kernel, layer.bias))
inputs = tf.random.uniform([batch_size, num_features])
labels = tf.random.uniform([batch_size, 1])
per_example_gradients = tf.vectorized_map(model_fn, (inputs, labels))
assert per_example_gradients[0].shape == (batch_size, num_features, 1)
assert per_example_gradients[1].shape == (batch_size, 1)
```
Args:
fn: The callable to be performed. It accepts one argument, which will have
the same (possibly nested) structure as `elems`, and returns a possibly
nested structure of Tensors and Operations, which may be different than
the structure of `elems`.
elems: A tensor or (possibly nested) sequence of tensors, each of which will
be unpacked along their first dimension. The nested sequence of the
resulting slices will be mapped over by `fn`. The first dimensions of all
elements must broadcast to a consistent value; equivalently, each
element tensor must have first dimension of either `B` or `1`, for some
common batch size `B >= 1`.
fallback_to_while_loop: If true, on failing to vectorize an operation,
the unsupported op is wrapped in a tf.while_loop to execute the map
iterations. Note that this fallback only happens for unsupported ops and
other parts of `fn` are still vectorized. If false, on encountering an
unsupported op, a ValueError is thrown. Note that the fallbacks can result
in slowdowns since vectorization often yields speedup of one to two orders
of magnitude.
warn: If set to `false`, this will supress any warnings due to operation
conversions in the provided `fn` falling back to while loops.
Returns:
A tensor or (possibly nested) sequence of tensors. Each tensor packs the
results of applying fn to tensors unpacked from elems along the first
dimension, from first to last.
Although they are less common as user-visible inputs and outputs, note that
tensors of type `tf.variant` which represent tensor lists (for example from
`tf.raw_ops.TensorListFromTensor`) are vectorized by stacking the list
contents rather than the variant itself, and so the container tensor will
have a scalar shape when returned rather than the usual stacked shape. This
improves the performance of control flow gradient vectorization.
Raises:
ValueError: If vectorization fails and fallback_to_while_loop is False.
"""
elems = variable_utils.convert_variables_to_tensors(elems)
elems = nest.map_structure(ops.convert_to_tensor,
elems,
expand_composites=True)
def loop_fn(i):
gathered_elems = nest.map_structure(
lambda x: _gather_from_tensor_or_composite(x, i), elems)
return fn(gathered_elems)
# Extract batch size from the maximum first dimension of any element.
flat_elems = nest.flatten(
nest.map_structure(
functools.partial(_composite_to_tensors,
is_batched=True),
elems))
def _get_shape(x):
if x.shape.rank is None:
return None
return x.shape.as_list()[0]
static_first_dims = [_get_shape(elem) for elem in flat_elems]
if any(s is None for s in static_first_dims):
batch_size = math_ops.reduce_max(
[array_ops.shape(elem)[0] for elem in flat_elems])
else:
batch_size = max(static_first_dims)
return pfor(
loop_fn,
batch_size,
fallback_to_while_loop=fallback_to_while_loop,
warn=warn)
tensorflow/python/ops/parallel_for/control_flow_ops_test.py
+2908
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/ops/parallel_for/gradients.py
+144
View File
@@ -0,0 +1,144 @@
# 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.
# ==============================================================================
"""Jacobian ops."""
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import check_ops
from tensorflow.python.ops import gradients_impl as gradient_ops
from tensorflow.python.ops.parallel_for import control_flow_ops
from tensorflow.python.util import nest
def jacobian(output, inputs, use_pfor=True, parallel_iterations=None):
"""Computes jacobian of `output` w.r.t. `inputs`.
Args:
output: A tensor.
inputs: A tensor or a nested structure of tensor objects.
use_pfor: If true, uses pfor for computing the jacobian. Else uses
tf.while_loop.
parallel_iterations: A knob to control how many iterations and dispatched in
parallel. This knob can be used to control the total memory usage.
Returns:
A tensor or a nested structure of tensors with the same structure as
`inputs`. Each entry is the jacobian of `output` w.r.t. to the corresponding
value in `inputs`. If output has shape [y_1, ..., y_n] and inputs_i has
shape [x_1, ..., x_m], the corresponding jacobian has shape
[y_1, ..., y_n, x_1, ..., x_m]. Note that in cases where the gradient is
sparse (IndexedSlices), jacobian function currently makes it dense and
returns a Tensor instead. This may change in the future.
"""
flat_inputs = nest.flatten(inputs)
output_tensor_shape = output.shape
output_shape = array_ops.shape(output)
output = array_ops.reshape(output, [-1])
def loop_fn(i):
y = array_ops.gather(output, i)
return gradient_ops.gradients(y, flat_inputs)
try:
output_size = int(output.shape[0])
except TypeError:
output_size = array_ops.shape(output)[0]
if use_pfor:
pfor_outputs = control_flow_ops.pfor(
loop_fn, output_size, parallel_iterations=parallel_iterations)
else:
pfor_outputs = control_flow_ops.for_loop(
loop_fn,
[output.dtype] * len(flat_inputs),
output_size,
parallel_iterations=parallel_iterations)
for i, out in enumerate(pfor_outputs):
if isinstance(out, tensor.Tensor):
new_shape = array_ops.concat(
[output_shape, array_ops.shape(out)[1:]], axis=0)
out = array_ops.reshape(out, new_shape)
out.set_shape(output_tensor_shape.concatenate(flat_inputs[i].shape))
pfor_outputs[i] = out
return nest.pack_sequence_as(inputs, pfor_outputs)
def batch_jacobian(output, inp, use_pfor=True, parallel_iterations=None):
"""Computes and stacks jacobians of `output[i,...]` w.r.t. `input[i,...]`.
e.g.
x = tf.constant([[1, 2], [3, 4]], dtype=tf.float32)
y = x * x
jacobian = batch_jacobian(y, x)
# => [[[2, 0], [0, 4]], [[6, 0], [0, 8]]]
Args:
output: A tensor with shape [b, y1, ..., y_n]. `output[i,...]` should
only depend on `inp[i,...]`.
inp: A tensor with shape [b, x1, ..., x_m]
use_pfor: If true, uses pfor for computing the Jacobian. Else uses a
tf.while_loop.
parallel_iterations: A knob to control how many iterations are vectorized
and dispatched in parallel. The default value of None, when use_pfor is
true, corresponds to vectorizing all the iterations. When use_pfor is
false, the default value of None corresponds to parallel_iterations=10.
This knob can be used to control the total memory usage.
Returns:
A tensor `t` with shape [b, y_1, ..., y_n, x1, ..., x_m] where `t[i, ...]`
is the jacobian of `output[i, ...]` w.r.t. `inp[i, ...]`, i.e. stacked
per-example jacobians.
Raises:
ValueError: if first dimension of `output` and `inp` do not match.
"""
output_shape = output.shape
if not output_shape[0].is_compatible_with(inp.shape[0]):
raise ValueError(f"Need first dimension of `output` shape ({output.shape}) "
f"and `inp` shape ({inp.shape}) to match.")
if output_shape.is_fully_defined():
batch_size = int(output_shape[0])
output_row_size = output_shape.num_elements() // batch_size
else:
output_shape = array_ops.shape(output)
batch_size = output_shape[0]
output_row_size = array_ops.size(output) // batch_size
inp_shape = array_ops.shape(inp)
# Flatten output to 2-D.
with ops.control_dependencies(
[check_ops.assert_equal(batch_size, inp_shape[0])]):
output = array_ops.reshape(output, [batch_size, output_row_size])
def loop_fn(i):
y = array_ops.gather(output, i, axis=1)
return gradient_ops.gradients(y, inp)[0]
if use_pfor:
pfor_output = control_flow_ops.pfor(loop_fn, output_row_size,
parallel_iterations=parallel_iterations)
else:
pfor_output = control_flow_ops.for_loop(
loop_fn, output.dtype,
output_row_size,
parallel_iterations=parallel_iterations)
if pfor_output is None:
return None
pfor_output = array_ops.reshape(pfor_output,
[output_row_size, batch_size, -1])
output = array_ops.transpose(pfor_output, [1, 0, 2])
new_shape = array_ops.concat([output_shape, inp_shape[1:]], axis=0)
return array_ops.reshape(output, new_shape)
tensorflow/python/ops/parallel_for/gradients_test.py
+684
View File
@@ -0,0 +1,684 @@
# 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 jacobian and batch_jacobian ops."""
import functools
import os
import time
import numpy as np
from tensorflow.python.client import session
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow.python.layers import layers as tf_layers
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import array_ops_stack
from tensorflow.python.ops import functional_ops
from tensorflow.python.ops import gradients as gradient_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn
from tensorflow.python.ops import random_ops
from tensorflow.python.ops import rnn
from tensorflow.python.ops import rnn_cell
from tensorflow.python.ops import variables
from tensorflow.python.ops import while_loop
from tensorflow.python.ops.losses import losses
from tensorflow.python.ops.parallel_for import control_flow_ops
from tensorflow.python.ops.parallel_for import gradients
from tensorflow.python.platform import test
from tensorflow.python.util import nest
class FullyConnectedModel:
def __init__(self, activation_size, num_layers):
self._layers = [
tf_layers.Dense(activation_size, activation=nn.relu)
for _ in range(num_layers)
]
def __call__(self, inp):
activation = inp
for layer in self._layers:
activation = layer(activation)
return activation
def fully_connected_model_fn(batch_size, activation_size, num_layers):
model = FullyConnectedModel(activation_size, num_layers)
inp = random_ops.random_normal([batch_size, activation_size])
return inp, model(inp)
def lstm_model_fn(batch_size, state_size, steps, inputs_size=None):
inputs_size = inputs_size or state_size
inputs = [
random_ops.random_normal([batch_size, inputs_size]) for _ in range(steps)
]
cell = rnn_cell.BasicLSTMCell(state_size)
init_state = cell.zero_state(batch_size, dtypes.float32)
state = init_state
for inp in inputs:
_, state = cell(inp, state)
return init_state.c, state.c
def dynamic_lstm_model_fn(batch_size, state_size, max_steps):
# We make inputs and sequence_length constant so that multiple session.run
# calls produce the same result.
inputs = constant_op.constant(
np.random.rand(batch_size, max_steps, state_size), dtype=dtypes.float32)
sequence_length = constant_op.constant(
np.random.randint(0, size=[batch_size], high=max_steps + 1),
dtype=dtypes.int32)
cell = rnn_cell.BasicLSTMCell(state_size)
initial_state = cell.zero_state(batch_size, dtypes.float32)
return inputs, rnn.dynamic_rnn(
cell,
inputs,
sequence_length=sequence_length,
initial_state=initial_state)
def create_fc_batch_jacobian(batch_size, activation_size, num_layers):
inp, output = fully_connected_model_fn(batch_size, activation_size,
num_layers)
pfor_jacobian = gradients.batch_jacobian(output, inp, use_pfor=True)
while_jacobian = gradients.batch_jacobian(output, inp, use_pfor=False)
return pfor_jacobian, while_jacobian
def create_lstm_batch_jacobian(batch_size, state_size, steps, inputs_size=None):
inp, output = lstm_model_fn(batch_size, state_size, steps,
inputs_size=inputs_size)
pfor_jacobian = gradients.batch_jacobian(output, inp, use_pfor=True)
while_jacobian = gradients.batch_jacobian(output, inp, use_pfor=False)
return pfor_jacobian, while_jacobian
def create_dynamic_lstm_batch_jacobian(batch_size, state_size, max_steps):
inp, (_, final_state) = dynamic_lstm_model_fn(batch_size, state_size,
max_steps)
pfor_jacobian = gradients.batch_jacobian(final_state.c, inp, use_pfor=True)
# Note that use_pfor=False does not work above given the current limitations
# on implementation of while_loop. So we statically unroll the looping in the
# jacobian computation.
while_gradients = [
gradient_ops.gradients(array_ops.gather(final_state.c, i, axis=1), inp)[0]
for i in range(state_size)
]
return pfor_jacobian, while_gradients
def create_lstm_batch_hessian(batch_size, state_size, steps):
inp, output = lstm_model_fn(batch_size, state_size, steps)
pfor_jacobian = gradients.batch_jacobian(output, inp, use_pfor=True)
pfor_jacobian = array_ops.reshape(pfor_jacobian, [batch_size, -1])
pfor_hessian = gradients.batch_jacobian(pfor_jacobian, inp, use_pfor=True)
# TODO(agarwal): using two nested while_loop doesn't seem to work here.
# Hence we use pfor_jacobian for computing while_hessian.
while_jacobian = pfor_jacobian
while_hessian = gradients.batch_jacobian(while_jacobian, inp, use_pfor=False)
return pfor_hessian, while_hessian
def create_lstm_hessian(batch_size, state_size, steps):
_, output = lstm_model_fn(batch_size, state_size, steps)
weights = variables.trainable_variables()
pfor_jacobians = gradients.jacobian(output, weights, use_pfor=True)
pfor_hessians = [
gradients.jacobian(x, weights, use_pfor=True) for x in pfor_jacobians
]
# TODO(agarwal): using two nested while_loop doesn't seem to work here.
# Hence we use pfor_jacobians for computing while_hessians.
while_jacobians = pfor_jacobians
while_hessians = [
gradients.jacobian(x, weights, use_pfor=False) for x in while_jacobians
]
return pfor_hessians, while_hessians
def create_fc_per_eg_grad(batch_size, activation_size, num_layers):
inp = random_ops.random_normal([batch_size, activation_size])
layers = [
tf_layers.Dense(activation_size, activation=nn.relu)
for _ in range(num_layers)
]
projection = tf_layers.Dense(1)
def model_fn(activation):
for layer in layers:
activation = layer(activation)
activation = projection(activation)
activation = nn.l2_loss(activation)
return gradient_ops.gradients(activation, variables.trainable_variables())
def loop_fn(i):
return model_fn(array_ops.expand_dims(array_ops.gather(inp, i), 0))
pfor_outputs = control_flow_ops.pfor(loop_fn, batch_size)
loop_fn_dtypes = [x.dtype for x in variables.trainable_variables()]
while_outputs = control_flow_ops.for_loop(loop_fn, loop_fn_dtypes, batch_size)
return pfor_outputs, while_outputs
def create_lstm_per_eg_grad(batch_size, state_size, steps, inputs_size=None):
inputs_size = inputs_size or state_size
inputs = [
random_ops.random_normal([batch_size, inputs_size]) for _ in range(steps)
]
cell = rnn_cell.BasicLSTMCell(state_size)
init_state = cell.zero_state(batch_size, dtypes.float32)
def model_fn(inps, init_state):
state = init_state
for inp in inps:
_, state = cell(inp, state)
output = nn.l2_loss(state.c)
return gradient_ops.gradients(output, variables.trainable_variables())
def loop_fn(i):
loop_inputs = [
array_ops.expand_dims(array_ops.gather(x, i), 0) for x in inputs
]
loop_init_state = rnn_cell.LSTMStateTuple(
*[array_ops.expand_dims(array_ops.gather(x, i), 0) for x in init_state])
return model_fn(loop_inputs, loop_init_state)
pfor_outputs = control_flow_ops.pfor(loop_fn, batch_size)
loop_fn_dtypes = [x.dtype for x in variables.trainable_variables()]
while_outputs = control_flow_ops.for_loop(loop_fn, loop_fn_dtypes, batch_size)
return pfor_outputs, while_outputs
# Importing the code from tensorflow_models seems to cause errors. Hence we
# duplicate the model definition here.
# TODO(agarwal): Use the version in tensorflow_models/official instead.
class Mnist(tf_layers.Layer):
def __init__(self, data_format):
"""Creates a model for classifying a hand-written digit.
Args:
data_format: Either 'channels_first' or 'channels_last'.
"""
super(Mnist, self).__init__()
if data_format == "channels_first":
self._input_shape = [-1, 1, 28, 28]
else:
assert data_format == "channels_last"
self._input_shape = [-1, 28, 28, 1]
self.conv1 = tf_layers.Conv2D(
32, 5, padding="same", data_format=data_format, activation=nn.relu)
self.conv2 = tf_layers.Conv2D(
64, 5, padding="same", data_format=data_format, activation=nn.relu)
self.fc1 = tf_layers.Dense(1024, activation=nn.relu)
self.fc2 = tf_layers.Dense(10)
self.dropout = tf_layers.Dropout(0.4)
self.max_pool2d = tf_layers.MaxPooling2D(
(2, 2), (2, 2), padding="same", data_format=data_format)
def __call__(self, inputs, training):
"""Add operations to classify a batch of input images.
Args:
inputs: A Tensor representing a batch of input images.
training: A boolean. Set to True to add operations required only when
training the classifier.
Returns:
A logits Tensor with shape [<batch_size>, 10].
"""
y = array_ops.reshape(inputs, self._input_shape)
y = self.conv1(y)
y = self.max_pool2d(y)
y = self.conv2(y)
y = self.max_pool2d(y)
y = tf_layers.flatten(y)
y = self.fc1(y)
y = self.dropout(y, training=training)
return self.fc2(y)
def create_mnist_autobatch(batch_size, data_format, training):
images = random_ops.random_uniform([batch_size, 28, 28])
model = Mnist(data_format)
manual = model(images, training=training)
def loop_fn(i):
image = array_ops.gather(images, i)
return model(image, training=training)
pfor_outputs = control_flow_ops.pfor(loop_fn, batch_size)
while_outputs = control_flow_ops.for_loop(
loop_fn, dtypes.float32, batch_size)
return pfor_outputs, while_outputs, manual
def create_mnist_per_eg_grad(batch_size, data_format, training):
images = random_ops.random_uniform([batch_size, 28, 28])
sparse_labels = np.random.randint(
low=0, high=10, size=[batch_size]).astype(np.int32)
labels = np.zeros((batch_size, 10)).astype(np.float32)
labels[np.arange(batch_size), sparse_labels] = 1.
model = Mnist(data_format)
def loop_fn(i):
image = array_ops.gather(images, i)
label = array_ops.gather(labels, i)
logits = array_ops.reshape(model(image, training=training), [-1])
loss = losses.softmax_cross_entropy(
logits=logits, onehot_labels=label, reduction=losses.Reduction.NONE)
return gradient_ops.gradients(loss, variables.trainable_variables())
pfor_outputs = control_flow_ops.pfor(loop_fn, batch_size)
while_outputs = control_flow_ops.for_loop(
loop_fn, [dtypes.float32] * len(variables.trainable_variables()),
batch_size)
return pfor_outputs, while_outputs
def create_mnist_batch_jacobian(batch_size, data_format, training):
images = random_ops.random_uniform([batch_size, 28, 28])
model = Mnist(data_format)
logits = model(images, training=training)
pfor_jacobian = gradients.batch_jacobian(logits, images, use_pfor=True)
while_jacobian = gradients.batch_jacobian(logits, images, use_pfor=False)
return pfor_jacobian, while_jacobian
def create_mnist_per_eg_jacobian(batch_size, data_format, training):
images = random_ops.random_uniform([batch_size, 28, 28])
model = Mnist(data_format)
def loop_fn(i, use_pfor):
image = array_ops.gather(images, i)
logits = array_ops.reshape(model(image, training=training), [-1])
return gradients.jacobian(
logits, variables.trainable_variables(), use_pfor=use_pfor)
pfor_outputs = control_flow_ops.pfor(
functools.partial(loop_fn, use_pfor=True),
batch_size)
while_outputs = control_flow_ops.for_loop(
functools.partial(loop_fn, use_pfor=False),
[dtypes.float32] * len(variables.trainable_variables()), batch_size)
return pfor_outputs, while_outputs
def create_fc_per_eg_jacobians(batch_size, activation_size, num_layers):
model = FullyConnectedModel(activation_size=activation_size,
num_layers=num_layers)
inp = random_ops.random_normal([batch_size, activation_size])
output = model(inp)
jacobians = gradients.jacobian(output, variables.trainable_variables())
def loop_fn(i, use_pfor):
inp_i = array_ops.expand_dims(array_ops.gather(inp, i), 0)
output = array_ops.reshape(model(inp_i), [-1])
return gradients.jacobian(
output, variables.trainable_variables(), use_pfor=use_pfor)
per_eg_jacobians_pfor = control_flow_ops.pfor(
functools.partial(loop_fn, use_pfor=True),
batch_size)
per_eg_jacobians_while = control_flow_ops.for_loop(
functools.partial(loop_fn, use_pfor=False),
[dtypes.float32] * len(variables.trainable_variables()), batch_size)
return jacobians, per_eg_jacobians_pfor, per_eg_jacobians_while
@test_util.run_v1_only("b/122612051")
class GradientsTest(test.TestCase):
def run_and_assert_equal(self, targets1, targets2, atol=1e-4, rtol=1e-4):
targets1 = nest.flatten(targets1)
targets2 = nest.flatten(targets2)
assert len(targets1) == len(targets2)
init = variables.global_variables_initializer()
self.evaluate(init)
outputs = self.evaluate(targets1 + targets2)
n = len(outputs) // 2
for i in range(n):
self.assertAllClose(outputs[i], outputs[i + n], rtol=rtol, atol=atol)
def test_no_path(self):
for grad_func in [gradients.jacobian, gradients.batch_jacobian]:
for use_pfor in [True, False]:
x = constant_op.constant([[1.0]])
y = constant_op.constant([[2.0]])
self.assertIsNone(grad_func(y, x, use_pfor=use_pfor))
def test_jacobian_fixed_shape(self):
x = random_ops.random_uniform([2, 2])
y = math_ops.matmul(x, x, transpose_a=True)
jacobian_pfor = gradients.jacobian(y, x, use_pfor=True)
jacobian_while = gradients.jacobian(y, x, use_pfor=False)
answer = ops.convert_to_tensor([[
gradient_ops.gradients(y[0][0], x)[0],
gradient_ops.gradients(y[0][1], x)[0]
], [
gradient_ops.gradients(y[1][0], x)[0],
gradient_ops.gradients(y[1][1], x)[0]
]])
self.run_and_assert_equal(answer, jacobian_pfor)
self.run_and_assert_equal(answer, jacobian_while)
def test_jacobian_scan_shape(self):
# Shape x: [3, 4]
x = random_ops.random_uniform([3, 4])
elems = random_ops.random_uniform([6])
# Shape y: [6, 3, 4]
y = functional_ops.scan(lambda a, e: a + e, elems, initializer=x)
jacobian = gradients.jacobian(y, x)
expected_shape = [6, 3, 4, 3, 4]
self.assertAllEqual(expected_shape, jacobian.shape.as_list())
def test_jacobian_while_loop_shape(self):
# Shape x: [3, 4]
x = random_ops.random_uniform([3, 4])
_, y = while_loop.while_loop(lambda i, a: i > 5., lambda i, a:
(i + 1, a + i), (constant_op.constant(0.), x))
# Shape y: [2, 3]
y = y[:2, :3]
jacobian = gradients.jacobian(y, x)
expected_shape = [2, 3, 3, 4]
self.assertAllEqual(expected_shape, jacobian.shape.as_list())
def test_jacobian_unknown_shape(self):
with self.cached_session() as sess:
x = array_ops.placeholder(dtypes.float32, shape=[None, None])
y = math_ops.matmul(x, x, transpose_a=True)
jacobian_pfor = gradients.jacobian(y, x, use_pfor=True)
jacobian_while = gradients.jacobian(y, x, use_pfor=False)
answer = ops.convert_to_tensor([[
gradient_ops.gradients(y[0][0], x)[0],
gradient_ops.gradients(y[0][1], x)[0]
], [
gradient_ops.gradients(y[1][0], x)[0],
gradient_ops.gradients(y[1][1], x)[0]
]])
ans, pfor_value, while_value = sess.run(
[answer, jacobian_pfor, jacobian_while],
feed_dict={x: [[1, 2], [3, 4]]})
self.assertAllClose(ans, pfor_value)
self.assertAllClose(ans, while_value)
def test_jacobian_parallel_iterations(self):
x = constant_op.constant([[1., 2], [3, 4]])
y = math_ops.matmul(x, x)
self.assertAllClose(gradients.jacobian(y, x, parallel_iterations=2),
gradients.jacobian(y, x, parallel_iterations=3))
def test_batch_jacobian_bad_shapes(self):
x = random_ops.random_uniform([2, 2])
y = random_ops.random_uniform([3, 2])
with self.assertRaisesRegex(ValueError, "Need first dimension of `output`"):
gradients.batch_jacobian(y, x, use_pfor=True)
def test_batch_jacobian_bad_unknown_shapes(self):
with self.cached_session() as sess:
x = array_ops.placeholder(dtypes.float32)
y = array_ops.concat([x, x], axis=0)
jacobian = gradients.batch_jacobian(y, x)
with self.assertRaisesRegex(errors.InvalidArgumentError,
"assertion failed"):
sess.run(jacobian, feed_dict={x: [[1, 2], [3, 4]]})
def test_batch_jacobian_fixed_shape(self):
x = random_ops.random_uniform([2, 3, 5])
y = x * x
batch_jacobian_pfor = gradients.batch_jacobian(y, x, use_pfor=True)
batch_jacobian_while = gradients.batch_jacobian(y, x, use_pfor=False)
two_x = 2 * x
answer = array_ops_stack.stack(
[array_ops.diag(two_x[0]),
array_ops.diag(two_x[1])])
self.run_and_assert_equal(answer, batch_jacobian_pfor)
self.run_and_assert_equal(answer, batch_jacobian_while)
def test_batch_jacobian_unknown_shape(self):
with self.cached_session() as sess:
x = array_ops.placeholder(dtypes.float32)
y = x * x
batch_jacobian_pfor = gradients.batch_jacobian(y, x, use_pfor=True)
batch_jacobian_while = gradients.batch_jacobian(y, x, use_pfor=False)
two_x = 2 * x
answer = array_ops_stack.stack(
[array_ops.diag(two_x[0]),
array_ops.diag(two_x[1])])
ans, pfor_value, while_value = sess.run(
[answer, batch_jacobian_pfor, batch_jacobian_while],
feed_dict={x: [[1, 2], [3, 4]]})
self.assertAllClose(ans, pfor_value)
self.assertAllClose(ans, while_value)
def test_batch_jacobian_parallel_iterations(self):
x = constant_op.constant([[1., 2], [3, 4]])
w = constant_op.constant([[1., 2, 3, 4], [5, 6, 7, 8]])
y = math_ops.matmul(x, w)
self.assertAllClose(gradients.batch_jacobian(y, x, parallel_iterations=2),
gradients.batch_jacobian(y, x, parallel_iterations=3))
def test_fc_batch_jacobian(self):
pfor_jacobian, while_jacobian = create_fc_batch_jacobian(8, 4, 2)
self.run_and_assert_equal(pfor_jacobian, while_jacobian)
def test_lstm_batch_jacobian(self):
pfor_jacobian, while_jacobian = create_lstm_batch_jacobian(8, 4, 2,
inputs_size=128)
self.run_and_assert_equal(pfor_jacobian, while_jacobian)
@test_util.disable_xla("This test never passed for XLA")
def DISABLED_test_dynamic_lstm_batch_jacobian(self):
pfor_jacobian, while_gradients = create_dynamic_lstm_batch_jacobian(8, 4, 3)
with session.Session() as sess:
init = variables.global_variables_initializer()
self.evaluate(init)
pfor = self.evaluate(pfor_jacobian)
for i in range(4):
while_i = sess.run(while_gradients[i])
self.assertAllClose(while_i, pfor[:, i, ...])
def test_lstm_hessian(self):
pfor_hessian, while_hessian = create_lstm_hessian(2, 2, 2)
self.run_and_assert_equal(pfor_hessian, while_hessian)
def test_lstm_batch_hessian(self):
pfor_hessian, while_hessian = create_lstm_batch_hessian(2, 2, 2)
self.run_and_assert_equal(pfor_hessian, while_hessian)
def test_fc_per_eg_grad(self):
pfor_outputs, while_outputs = create_fc_per_eg_grad(8, 4, 2)
self.run_and_assert_equal(pfor_outputs, while_outputs)
def test_lstm_per_eg_grad(self):
pfor_outputs, while_outputs = create_lstm_per_eg_grad(8, 4, 2)
self.run_and_assert_equal(pfor_outputs, while_outputs)
def test_mnist_per_eg_grad(self):
# It looks like CUDNN_CONVOLUTION_BWD_DATA_ALGO_WINOGRAD_NONFUSED
# configuration of Winograd can cause low precision output resulting in
# tests failing. So we disable that here.
os.environ["TF_ENABLE_WINOGRAD_NONFUSED"] = "0"
data_format = ("channels_first"
if test.is_gpu_available() else "channels_last")
# Note that we are setting training=False here so that dropout produces
# the same result with pfor and with while_loop.
pfor_outputs, while_outputs = create_mnist_per_eg_grad(
4, data_format, training=False)
self.run_and_assert_equal(pfor_outputs, while_outputs, rtol=1e-3)
os.environ.pop("TF_ENABLE_WINOGRAD_NONFUSED", None)
def test_mnist_per_eg_jacobian(self):
# It looks like CUDNN_CONVOLUTION_BWD_DATA_ALGO_WINOGRAD_NONFUSED
# configuration of Winograd can cause low precision output resulting in
# tests failing. So we disable that here.
os.environ["TF_ENABLE_WINOGRAD_NONFUSED"] = "0"
data_format = ("channels_first"
if test.is_gpu_available() else "channels_last")
# Note that we are setting training=False here so that dropout produces
# the same result with pfor and with while_loop.
pfor_outputs, while_outputs = create_mnist_per_eg_jacobian(
2, data_format, training=False)
self.run_and_assert_equal(pfor_outputs, while_outputs, rtol=1e-3)
os.environ.pop("TF_ENABLE_WINOGRAD_NONFUSED", None)
def test_fc_jacobian(self):
jacobians, per_eg_jacobians_pfor, per_eg_jacobians_while = (
create_fc_per_eg_jacobians(batch_size=8,
activation_size=4,
num_layers=2))
self.run_and_assert_equal(jacobians, per_eg_jacobians_pfor,
rtol=2e-3, atol=1e-3)
self.run_and_assert_equal(jacobians, per_eg_jacobians_while,
rtol=2e-3, atol=1e-3)
def test_indexed_slice(self):
inp = random_ops.random_uniform([3, 2])
output = nn.embedding_lookup(inp, [0, 2])
pfor_jacobian = gradients.jacobian(output, inp, use_pfor=True)
while_jacobian = gradients.jacobian(output, inp, use_pfor=False)
self.run_and_assert_equal(while_jacobian, pfor_jacobian)
class GradientsBenchmarks(test.Benchmark):
def _run(self, targets, iters, name=None):
def _done(t):
# Note that we don't use tf.control_dependencies since that will not make
# sure that the computation on GPU has actually finished. So we fetch the
# first element of the output, and assume that this will not be called on
# empty tensors.
return array_ops.gather(array_ops.reshape(t, [-1]), 0)
targets = [_done(x) for x in nest.flatten(targets)]
sess = session.Session()
with sess:
init = variables.global_variables_initializer()
self.evaluate(init)
self.evaluate(targets)
begin = time.time()
for _ in range(iters):
self.evaluate(targets)
end = time.time()
avg_time_ms = (1000 * (end - begin)) / iters
self.report_benchmark(iters=iters, wall_time=avg_time_ms, name=name)
return avg_time_ms
def benchmark_fc_batch_jacobian(self):
with ops.Graph().as_default():
pfor_jacobian, while_jacobian = create_fc_batch_jacobian(100, 32, 20)
self._run(pfor_jacobian, 100, name="fc_batch_jacobian_pfor")
self._run(while_jacobian, 20, name="fc_batch_jacobian_while")
def benchmark_lstm_batch_jacobian(self):
with ops.Graph().as_default():
pfor_jacobian, while_jacobian = create_lstm_batch_jacobian(
100, 32, 8, inputs_size=128)
self._run(pfor_jacobian, 100, name="lstm_batch_jacobian_pfor")
self._run(while_jacobian, 20, name="lstm_batch_jacobian_while")
def benchmark_lstm_hessian(self):
with ops.Graph().as_default():
pfor_hessian, while_hessian = create_lstm_hessian(2, 2, 10)
self._run(pfor_hessian, 20, name="lstm_hessian_pfor")
self._run(while_hessian, 3, name="lstm_hessian_while_pfor")
def benchmark_lstm_batch_hessian(self):
with ops.Graph().as_default():
pfor_hessian, while_hessian = create_lstm_batch_hessian(4, 4, 10)
self._run(pfor_hessian, 100, name="lstm_batch_hessian_pfor")
self._run(while_hessian, 20, name="lstm_batch_hessian_while_pfor")
def benchmark_fc_per_eg_grad(self):
with ops.Graph().as_default():
pfor_outputs, while_outputs = create_fc_per_eg_grad(100, 32, 3)
self._run(pfor_outputs, 100, name="fc_per_eg_grad_pfor")
self._run(while_outputs, 20, name="fc_per_eg_grad_while")
def benchmark_lstm_per_eg_grad(self):
with ops.Graph().as_default():
pfor_outputs, while_outputs = create_lstm_per_eg_grad(100, 32, 8)
self._run(pfor_outputs, 100, name="lstm_per_eg_grad_pfor")
self._run(while_outputs, 20, name="lstm_per_eg_grad_while")
def benchmark_mnist_autobatch(self):
with ops.Graph().as_default():
data_format = ("channels_first"
if test.is_gpu_available() else "channels_last")
pfor_outputs, while_outputs, manual = create_mnist_autobatch(
100, data_format, training=False)
self._run(pfor_outputs, 100, name="mnist_pfor")
self._run(while_outputs, 20, name="mnist_while")
self._run(manual, 100, name="mnist_manual")
def benchmark_mnist_per_eg_grad(self):
with ops.Graph().as_default():
data_format = ("channels_first"
if test.is_gpu_available() else "channels_last")
pfor_outputs, while_outputs = create_mnist_per_eg_grad(
128, data_format, training=True)
self._run(pfor_outputs, 20, name="mnist_per_eg_grad_pfor")
self._run(while_outputs, 20, name="mnist_per_eg_grad_while")
def benchmark_mnist_per_eg_jacobian(self):
with ops.Graph().as_default():
if test.is_gpu_available():
data_format = "channels_first"
else:
data_format = "channels_last"
pfor_outputs, while_outputs = create_mnist_per_eg_jacobian(
16, data_format, training=True)
self._run(pfor_outputs, 20, name="mnist_per_eg_jacobian_pfor")
self._run(while_outputs, 20, name="mnist_per_eg_jacobian_while")
def benchmark_mnist_batch_jacobian(self):
with ops.Graph().as_default():
if test.is_gpu_available():
data_format = "channels_first"
else:
data_format = "channels_last"
pfor_outputs, while_outputs = create_mnist_batch_jacobian(
128, data_format, training=True)
self._run(pfor_outputs, 20, name="mnist_batch_jacobian_pfor")
self._run(while_outputs, 20, name="mnist_batch_jacobian_while")
def benchmark_fc_per_eg_jacobian(self):
with ops.Graph().as_default():
jacobians, per_eg_jacobians_pfor, per_eg_jacobians_while = (
create_fc_per_eg_jacobians(batch_size=128,
activation_size=32,
num_layers=3))
self._run(jacobians, 30, name="fc_jacobians_pfor")
self._run(per_eg_jacobians_pfor, 100,
name="fc_per_eg_jacobians_pfor")
self._run(per_eg_jacobians_while, 10,
name="fc_per_eg_jacobians_while")
if __name__ == "__main__":
test.main()
tensorflow/python/ops/parallel_for/index.md
+201
View File
@@ -0,0 +1,201 @@
# Vectorized map
*See also https://en.wikipedia.org/wiki/Automatic_vectorization*
TensorFlow provides in-graph looping constructs like `tf.while_loop` which are
similar to loops in other languages: they repeatedly run the loop body, not
keeping memory contiguous or taking advantage of hardware SIMD. When loop
iterations are independent, it is much more efficient to batch tensors together:
a matrix-matrix multiply instead of a loop over vector-matrix multiplies.
`tf.vectorized_map` provides a `tf.map_fn`-like API with the efficiency of
manual batching: while the API matches APIs which are implemented with a
`tf.while_loop` like `tf.map_fn`, `tf.vectorized_map` is implemented by using
batch dimensions of ops. This `tf.map_fn` style is often a more convenient way
to author models, as opposed to juggling a batch dimension explicitly:
```python
def f(args):
embeddings, index = args
# embeddings [vocab_size, embedding_dim]
# index []
# desired result: [embedding_dim]
return tf.gather(params=embeddings, indices=index)
@tf.function
def f_auto_vectorized(embeddings, indices):
# embeddings [num_heads, vocab_size, embedding_dim]
# indices [num_heads]
# desired result: [num_heads, embedding_dim]
return tf.vectorized_map(f, [embeddings, indices])
concrete_vectorized = f_auto_vectorized.get_concrete_function(
tf.TensorSpec(shape=[None, 100, 16], dtype=tf.float32),
tf.TensorSpec(shape=[None], dtype=tf.int32))
print(concrete_vectorized.graph.as_graph_def())
```
The vectorized graph contains many ops, but no loops. Instead,
`tf.vectorized_map` looks at the GatherV2 op and its attributes, and generates
the equivalent of `tf.gather(..., batch_dims=1)` without requiring the user to
know how to tell `tf.gather` and every other op they use which dimensions are
batch dimensions.
```python
gdef = concrete_vectorized.graph.as_graph_def()
print([n for n in gdef.node if n.op == "GatherV2"])
```
This prints a bunch of gathers related to pfor infrastructure, but at the time
of writing does include one with a `batch_dims` attribute of `1`.
## Vectorization as a post-trace graph transformation
`tf.vectorized_map` currently works as a graph-to-graph transformation
implemented in Python. This is mostly a historical artifact: it was conceived
before TensorFlow did op-by-op execution. It is similar to `tf.gradients`,
walking the connections in an existing graph and applying op-specific rules
(defined by `RegisterGradient` for `tf.gradients`, `RegisterPFor` for
`tf.vectorized_map`) in order to produce a new graph. While `tf.gradients` adds
a backward pass which references tensors in the forward pass (both execute),
`tf.vectorized_map` creates a transformed graph which executes in place of the
original graph.
For gradients, `tf.GradientTape` was introduced to provide an op-by-op version
of gradients, re-using the per-op `RegisterGradient` definitions. There is no
equivalent for vectorization. Instead, `tf.vectorized_map` wraps the function it
takes as an argument in `tf.function` in order to create a trace to
vectorize. This means that the user's function never executes eagerly, and if
`tf.vectorized_map` is called executing eagerly that the user's function is
re-traced and re-vectorized every call to `tf.vectorized_map`.
While `tf.vectorized_map` is the public-facing API, the implementation is
written in terms of [an integer-indexed for
loop](https://github.com/tensorflow/tensorflow/blob/8b000ce0d5395d399e08791ae9589b41358f651d/tensorflow/python/ops/parallel_for/control_flow_ops.py#L134). The
loop does not execute as a regular for loop, but this is a good mental model and
the implementation makes frequent references to `loop_len`, i.e. the number of
iterations for the hypothetical loop. The user-visible outputs should ideally
match the outputs of an equivalent real for loop, and this is how most of the
unit tests are written.
The virtual for loop setup includes a loop-variant integer loop index
tensor. "Loop-variant" just means a tensor with a different value on each
iteration; loop-variant tensors are represented with a leading extra dimension
corresponding to the loop iteration. `tf.vectorized_map`'s implementation
`tf.gather`s a slice of each input using the loop index and then runs the user's
function on those (loop-variant) values.
Anything with the loop index in its transitive input is loop-variant and must be
transformed. Ops like `tf.constant`, however, create loop-invariant values
(i.e. their values are the same on each loop iteration). Loop-invariant values
returned from vectorization may simply be tiled, but more frequently they feed
into ops with a mix of variant/invariant inputs to produce loop-variant
values. Converters for ops will sometimes have simpler special cases for
loop-invariant inputs, e.g. `tf.roll`'s converter is much simpler if the shift
amount is loop-invariant (a common case).
## Defining vectorizations
As with gradients, most ops have relatively straightforward definitions and
function call / control flow operations have complicated special cases. This
section covers the common cases.
As with `RegisterGradient`, converters are defined for op types (with a
corresponding `REGISTER_OP` macro in C++, named WithUppercase), not Python
endpoints. So if the user writes `tf.roll`, the corresponding
[`RegisterPFor("Roll")`
converter](https://github.com/tensorflow/tensorflow/blob/349172cf0ac29ba1346d244a40dc4761b4600f2e/tensorflow/python/ops/parallel_for/pfor.py#L2653)
is triggered since `tf.roll` is implemented with the "Roll" op.
Like gradients, the set of all TensorFlow ops would ideally be closed under
vectorization (i.e. vectorization would always produce ops which are themselves
vectorizable). In practice not all ops have pfor converters defined, and those
that do sometimes assume inputs are loop-invariant. A "fallback converter" runs
in these cases, adding a `tf.while_loop` in place of a real vectorization for
the op. This is safe but generally slower. `tf.while_loop` can run iterations in
parallel (non-strict execution), but other benefits of vectorization like
contiguous memory layouts and SIMD are not available.
For stateless ops which compute a deterministic value from their inputs, the
common case, pfor converters take loop-variant inputs with an extra dimension
("stacked") and emit an op or subgraph which treats this extra stacked dimension
(the tensor's zeroth dimension) as a batch dimension but otherwise computes the
same value as the original op. This may involve examining the original op's
attributes and forwarding them to newly emitted ops.
The general case for a converter has every input stacked and loop-variant, but
there are often more efficient special cases when some inputs are loop-invariant
and so may be handled "unstacked" (with no special zeroth dimension). Some
converters omit the general case entirely, simply requesting the unstacked input
and relying on the fallback converter triggering if that fails because the input
is loop-variant. Others have branches for various combinations of
stacked/unstacked inputs.
There are many examples of existing converters, all in
tensorflow/python/ops/parallel_for/pfor.py (the user-facing APIs are defined in
control_flow_ops.py in the same directory). Converters can be quite subtle, so
it is important to use the unit test macros to compare to a ground-truth for
loop and to ensure that all relevant combinations of loop variant/invariant
inputs are covered by these tests.
## Stateful ops
Stateful ops and ops with non-deterministic outputs are difficult to deal
with. One option is to use the fallback `tf.while_loop` converter for these ops,
so e.g. `tf.print` would print `loop_len` times with the different loop-variant
values. This makes sense from the "for loop as ground truth" mindset, but it's
less clear that this satisfies user expectations for `tf.vectorized_map` (which
doesn't explicitly mention a loop).
There isn't a great universal answer for this class of ops. Currently `tf.print`
[prints the full vectorized
tensors](https://github.com/tensorflow/tensorflow/blob/349172cf0ac29ba1346d244a40dc4761b4600f2e/tensorflow/python/ops/parallel_for/pfor.py#L3505-L3522) ([example](https://github.com/tensorflow/tensorflow/blob/8b202f08d52e8206af2bdb2112a62fafbc546ec7/tensorflow/python/ops/parallel_for/control_flow_ops_test.py#L956-L970))
rather than printing `loop_len` times. Stateful random ops are [vectorized by
adding an extra dimension to their output](https://github.com/tensorflow/tensorflow/blob/349172cf0ac29ba1346d244a40dc4761b4600f2e/tensorflow/python/ops/parallel_for/pfor.py#L3276-L3294) shape attributes, even though this
gives a different result (but follows the same distribution / independence
structure). Stateless random ops [use the `tf.while_loop` fallback converter](https://github.com/tensorflow/tensorflow/blob/349172cf0ac29ba1346d244a40dc4761b4600f2e/tensorflow/python/ops/parallel_for/pfor.py#L3360-L3377)
since users might care more about the exact values; this may want revisiting if
stateless random ops are used to implement popular APIs.
## Vectorization of control flow (while_loop, cond) and variants
Ops whose execution is defined by a serialized program (generally a FunctionDef
referenced by name in an attribute) need special handling, since the vectorized
op will reference a transformed serialized program.
For function call operations this is relatively straightforward: the converter
converts the function body and generates a new call operation referencing the
vectorized function body. (See `RegisterPfor("PartitionedCall")`; the code is
pretty readable.)
Cond (["If"/"StatelessIf" ops](https://github.com/tensorflow/tensorflow/blob/349172cf0ac29ba1346d244a40dc4761b4600f2e/tensorflow/python/ops/parallel_for/pfor.py#L4499); [example](https://github.com/tensorflow/tensorflow/blob/8b202f08d52e8206af2bdb2112a62fafbc546ec7/tensorflow/python/ops/parallel_for/control_flow_ops_test.py#L2041-L2053)) can be a bit more complicated if the Boolean
condition is loop variant, in which case inputs/outputs must be partitioned
between the branches and both run (although the ops in one branch could have
zero-sized inputs if the loop variant condition happened to not trigger that
branch for any iteration of the virtual for loop). If the condition Boolean is
loop invariant then cond is very similar to a function call operation, just with
two function bodies to transform.
[While loop vectorization](https://github.com/tensorflow/tensorflow/blob/349172cf0ac29ba1346d244a40dc4761b4600f2e/tensorflow/python/ops/parallel_for/pfor.py#L5001) is fairly complicated. This is unrelated to the
fallback converter for ops; it triggers when users define a graph with a
`tf.while_loop` and then request vectorization for it (although the fallback
converter can trigger this case when `tf.vectorized_map` is nested). At a high
level, while loop conversion is an iterative version of the
loop-variant-condition cond conversion. Only one while loop runs in the
vectorized graph, but it keeps track of which iterations of the virtual pfor
loop are done and only runs the while loop body for corresponding inputs. Once
all of the iterations of the virtual pfor loop would have finished their
`tf.while_loop`s the single vectorized loop terminates.
While loops accumulate values across iterations in TensorLists (aka
TensorArrays). These are variant-dtype tensors with a C++ vector of pointers to
other tensors. A straightforward conversion would simply stack variant tensors,
so rather than scalar variant-dtype tensors they would have shape
`[loop_len]`. However, this would make memory non-contiguous: the tensors across
each iteration of the virtual pfor loop would be separate Tensor objects in C++,
and concatenation / splitting would be necessary to push and pop
tensors. Instead, TensorLists are special-cased to use "internal vectorization":
the variant representing a vectorized/stacked TensorList remains a scalar, but
the shape of the tensors it contains has a special zeroth dimension. This makes
many common operations on vectorized TensorLists more efficient, but leads to
some complicated special cases when accessing the vectorization dimension.
tensorflow/python/ops/parallel_for/math_test.py
+805
View File
@@ -0,0 +1,805 @@
# 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 vectorization of math kernels."""
import itertools
from absl.testing import parameterized
from tensorflow.python.eager import backprop
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops as framework_ops
from tensorflow.python.framework import tensor
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import clip_ops
from tensorflow.python.ops import linalg_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn
from tensorflow.python.ops import random_ops
from tensorflow.python.ops import special_math_ops
from tensorflow.python.ops import tensor_array_grad # pylint: disable=unused-import
from tensorflow.python.ops.parallel_for import control_flow_ops as pfor_control_flow_ops
from tensorflow.python.ops.parallel_for.test_util import PForTestCase
from tensorflow.python.platform import test
@test_util.run_all_in_graph_and_eager_modes
class MathTest(PForTestCase, parameterized.TestCase):
def _test_unary_cwise_ops(self, ops, is_complex):
for op in ops:
with backprop.GradientTape(persistent=True) as g:
x = random_ops.random_uniform([3, 5])
g.watch(x)
if is_complex:
y = random_ops.random_uniform([3, 5])
g.watch(y)
x = math_ops.complex(x, y)
# pylint: disable=cell-var-from-loop
def loop_fn(i):
with g:
y = op(x)
x_i = array_ops.gather(x, i)
y_i = op(x_i)
outputs = [y_i]
# Build cross product of loop variant/invariant outputs and gradients.
for out in (y, y_i):
if out.dtype == dtypes.float32:
for output_gradients in (None, out * math_ops.cast(i, out.dtype)):
grad = g.gradient(out, x_i, output_gradients=output_gradients)
if grad is not None:
outputs.append(grad)
return outputs
# pylint: enable=cell-var-from-loop
self._test_loop_fn(loop_fn, 3)
def test_unary_cwise_complex_ops(self):
complex_ops = [
math_ops.angle,
math_ops.imag,
math_ops.complex_abs,
math_ops.real,
math_ops.conj,
]
self._test_unary_cwise_ops(complex_ops, True)
def test_unary_cwise_real_ops_1(self):
real_ops = [
lambda x: math_ops.acosh(1 + math_ops.square(x)),
math_ops.abs,
math_ops.acos,
math_ops.asin,
math_ops.asinh,
math_ops.atan,
math_ops.atanh,
math_ops.cos,
math_ops.cosh,
math_ops.digamma,
math_ops.erf,
math_ops.erfc,
math_ops.erfinv,
math_ops.exp,
math_ops.expm1,
math_ops.inv,
math_ops.is_finite,
math_ops.is_inf,
math_ops.lgamma,
math_ops.log,
math_ops.log1p,
math_ops.ndtri,
special_math_ops.bessel_i0e,
special_math_ops.bessel_i1e,
]
self._test_unary_cwise_ops(real_ops, False)
def test_unary_cwise_real_ops_2(self):
real_ops = [
math_ops.neg,
math_ops.negative,
math_ops.reciprocal,
math_ops.rint,
math_ops.round,
math_ops.rsqrt,
math_ops.sigmoid,
math_ops.sign,
math_ops.sin,
math_ops.sinh,
math_ops.sqrt,
math_ops.square,
math_ops.tan,
math_ops.tanh,
nn.elu,
nn.relu,
nn.relu6,
lambda t: nn.leaky_relu(t, alpha=0.1),
nn.selu,
nn.softplus,
nn.softsign,
]
self._test_unary_cwise_ops(real_ops, False)
def test_unary_cwise_no_grad(self):
for op in [math_ops.ceil, math_ops.floor, math_ops.logical_not]:
x = random_ops.random_uniform([3, 5])
if op == math_ops.logical_not:
x = x > 0
# pylint: disable=cell-var-from-loop
def loop_fn(i):
return op(array_ops.gather(x, i))
# pylint: enable=cell-var-from-loop
self._test_loop_fn(loop_fn, 3)
def test_binary_cwise_ops(self):
# Enable tensor equality to test `equal` and `not_equal` ops below.
default_equality = tensor.Tensor._USE_EQUALITY
tensor.enable_tensor_equality()
try:
logical_ops = [
math_ops.logical_and, math_ops.logical_or, math_ops.logical_xor
]
# Wrapper functions restricting the range of inputs of zeta and polygamma.
def safe_polygamma(x, y):
return math_ops.polygamma(
math_ops.round(clip_ops.clip_by_value(y, 1, 10)), x * x + 1)
def safe_zeta(x, y):
return math_ops.zeta(x * x + 1, y * y)
float_ops = [
math_ops.add,
math_ops.add_v2,
math_ops.atan2,
math_ops.complex,
math_ops.div,
math_ops.divide,
math_ops.div_no_nan,
math_ops.equal,
lambda x, y: framework_ops.convert_to_tensor(x == y),
lambda x, y: framework_ops.convert_to_tensor(x != y),
math_ops.floor_mod,
math_ops.greater,
math_ops.greater_equal,
math_ops.igamma,
math_ops.igammac,
math_ops.igamma_grad_a,
math_ops.less,
math_ops.less_equal,
math_ops.maximum,
math_ops.minimum,
math_ops.mod,
math_ops.multiply,
math_ops.not_equal,
math_ops.pow,
math_ops.squared_difference,
math_ops.subtract,
math_ops.truncate_mod,
safe_polygamma,
]
# FloorDiv fails on XLA due floor's discontinuities exacerbating small
# division differences.
if not test_util.is_xla_enabled():
float_ops += [math_ops.floor_div]
# TODO(b/168912036): Re-enable once GPU + XLA issues for Zeta are
# resolved.
if not test_util.is_gpu_available():
float_ops += [safe_zeta]
for op in logical_ops + float_ops:
x = random_ops.random_uniform([7, 3, 5])
y = random_ops.random_uniform([3, 5])
if op in logical_ops:
x = x > 0
y = y > 0
output_dtypes = []
# pylint: disable=cell-var-from-loop
def loop_fn(i):
x1 = array_ops.gather(x, i)
y1 = array_ops.gather(y, i)
outputs = [op(x, y), op(x1, y), op(x, y1), op(x1, y1), op(x1, x1)]
del output_dtypes[:]
output_dtypes.extend(t.dtype for t in outputs)
return outputs
# pylint: enable=cell-var-from-loop
self._test_loop_fn(loop_fn, 3)
finally:
if not default_equality:
tensor.disable_tensor_equality()
def test_approximate_equal(self):
x = random_ops.random_uniform([3, 5])
y = random_ops.random_uniform([3, 5])
def loop_fn(i):
x1 = array_ops.gather(x, i)
y1 = array_ops.gather(y, i)
return math_ops.approximate_equal(x1, y1)
self._test_loop_fn(loop_fn, 3)
def test_abs_complex(self):
r = pfor_control_flow_ops.vectorized_map(
math_ops.abs, math_ops.cast([0, -1], dtype=dtypes.complex128))
self.assertAllEqual(self.evaluate(r), [0, 1])
def test_colocate_with(self):
a = random_ops.random_uniform([3, 5])
def loop_fn(i):
x1 = array_ops.gather(a, i)
with framework_ops.colocate_with(x1):
return x1 * 2
self._test_loop_fn(loop_fn, 3)
def test_addn(self):
x = random_ops.random_uniform([2, 3, 5])
y = random_ops.random_uniform([3, 5])
z = random_ops.random_uniform([3, 5])
def loop_fn(i):
x1 = array_ops.gather(x, i)
return math_ops.add_n([x1, y, z])
self._test_loop_fn(loop_fn, 2)
def test_cross(self):
x = random_ops.random_uniform([4, 2, 3])
y = random_ops.random_uniform([4, 2, 3])
def loop_fn(i):
x_i = array_ops.gather(x, i)
y_i = array_ops.gather(y, i)
x_0 = array_ops.gather(x, 0)
return math_ops.cross(x_i, y_i), math_ops.cross(x_0, y_i)
self._test_loop_fn(loop_fn, 4)
@test_util.run_without_tensor_float_32(
"Calls matmul in parallel for-loop and compares result to calling matmul "
"in sequential for-loop")
def test_matmul(self):
for tr_a in (True, False):
for tr_b in (True, False):
for stack_a in (True, False):
for stack_b in (True, False):
shape_a = (5, 3) if tr_a else (3, 5)
if stack_a:
shape_a = (2,) + shape_a
shape_b = (7, 5) if tr_b else (5, 7)
if stack_b:
shape_b = (2,) + shape_b
x = random_ops.random_uniform(shape_a)
y = random_ops.random_uniform(shape_b)
# pylint: disable=cell-var-from-loop
def loop_fn(i):
a = array_ops.gather(x, i) if stack_a else x
b = array_ops.gather(y, i) if stack_b else y
return math_ops.matmul(a, b, transpose_a=tr_a, transpose_b=tr_b)
# pylint: enable=cell-var-from-loop
self._test_loop_fn(loop_fn, 2)
def test_batch_matmul(self):
for tr_a in (True, False):
for tr_b in (True, False):
for stack_a in (True, False):
for stack_b in (True, False):
shape_a = (4, 5, 3) if tr_a else (4, 3, 5)
if stack_a:
shape_a = (2,) + shape_a
shape_b = (4, 7, 5) if tr_b else (4, 5, 7)
if stack_b:
shape_b = (2,) + shape_b
x = random_ops.random_uniform(shape_a)
y = random_ops.random_uniform(shape_b)
# pylint: disable=cell-var-from-loop
def loop_fn(i):
a = array_ops.gather(x, i) if stack_a else x
b = array_ops.gather(y, i) if stack_b else y
return math_ops.matmul(a, b, transpose_a=tr_a, transpose_b=tr_b)
# pylint: enable=cell-var-from-loop
self._test_loop_fn(loop_fn, 2)
def test_batch_matmul_broadcast(self):
for broadcast_a in (True, False):
for broadcast_b in (True, False):
for stack_a in (True, False):
for stack_b in (True, False):
shape_a = (2, 3, 5) if broadcast_a else (4, 2, 3, 5)
shape_b = (2, 5, 7) if broadcast_b else (4, 2, 5, 7)
shape_a = (2,) + shape_a if stack_a else shape_a
shape_b = (2,) + shape_b if stack_b else shape_b
x = random_ops.random_uniform(shape_a)
y = random_ops.random_uniform(shape_b)
# pylint: disable=cell-var-from-loop
def loop_fn(i):
a = array_ops.gather(x, i) if stack_a else x
b = array_ops.gather(y, i) if stack_b else y
return math_ops.matmul(a, b)
# pylint: enable=cell-var-from-loop
self._test_loop_fn(loop_fn, 2)
def test_reduction(self):
x = random_ops.random_uniform([2, 3, 4, 5])
for op in [
math_ops.reduce_sum,
math_ops.reduce_prod,
math_ops.reduce_max,
math_ops.reduce_min,
math_ops.reduce_mean,
]:
for axis in ([1], None, [0, 2], constant_op.constant([1], dtypes.int64)):
for keepdims in (True, False):
# pylint: disable=cell-var-from-loop
def loop_fn(i):
a = array_ops.gather(x, i)
return op(a, axis=axis, keepdims=keepdims)
# pylint: enable=cell-var-from-loop
self._test_loop_fn(loop_fn, 2)
def test_boolean_reduction(self):
x = random_ops.random_uniform([2, 3, 4, 5]) > 0.5
for op in [math_ops.reduce_any, math_ops.reduce_all]:
for axis in ([1], None, [0, 2], constant_op.constant([1], dtypes.int64)):
for keepdims in (True, False):
# pylint: disable=cell-var-from-loop
def loop_fn(i):
a = array_ops.gather(x, i)
return op(a, axis=axis, keepdims=keepdims)
# pylint: enable=cell-var-from-loop
self._test_loop_fn(loop_fn, 2)
def test_argmin_argmax(self):
x = random_ops.random_uniform([2, 3, 4, 5])
for op in [math_ops.argmin, math_ops.argmax]:
for axis in (1, None, -1):
for output_dtype in (dtypes.int32, dtypes.int64, None):
# pylint: disable=cell-var-from-loop
def loop_fn(i):
a = array_ops.gather(x, i)
return op(a, axis=axis, output_type=output_dtype)
# pylint: enable=cell-var-from-loop
self._test_loop_fn(loop_fn, 2)
def test_bucketize(self):
x = random_ops.random_uniform([2, 3, 4])
def loop_fn(i):
a = array_ops.gather(x, i)
return math_ops.bucketize(a, [-1, 0.5, 1])
self._test_loop_fn(loop_fn, 2)
def test_clip_by_value(self):
x = random_ops.random_uniform([2, 3, 4])
def loop_fn(i):
a = array_ops.gather(x, i)
return clip_ops.clip_by_value(a, 0.5, 1.0)
self._test_loop_fn(loop_fn, 2)
def test_cum_sum(self):
x = random_ops.random_uniform([2, 3, 4, 5])
for axis in (1, -2, constant_op.constant(1, dtypes.int64)):
for exclusive in (True, False):
for reverse in (True, False):
# pylint: disable=cell-var-from-loop
def loop_fn(i):
a = array_ops.gather(x, i)
return math_ops.cumsum(
a, axis=axis, exclusive=exclusive, reverse=reverse)
# pylint: enable=cell-var-from-loop
self._test_loop_fn(loop_fn, 2)
def test_cum_prod(self):
x = random_ops.random_uniform([2, 3, 4, 5])
for axis in (1, -2, constant_op.constant(1, dtypes.int64)):
for exclusive in (True, False):
for reverse in (True, False):
# pylint: disable=cell-var-from-loop
def loop_fn(i):
a = array_ops.gather(x, i)
return math_ops.cumprod(
a, axis=axis, exclusive=exclusive, reverse=reverse)
# pylint: enable=cell-var-from-loop
self._test_loop_fn(loop_fn, 2)
def test_bias_add(self):
for data_format in ("NCHW", "NHWC"):
for stacked_value in (True, False):
x_shape = [3, 4, 5, 6]
if stacked_value:
x_shape = [2] + x_shape
x = random_ops.random_uniform(x_shape)
for stacked_bias in (True, False):
if not (stacked_value or stacked_bias):
continue
with backprop.GradientTape(persistent=True) as g:
bias_dim = -1
if data_format == "NCHW":
bias_dim = 2 if stacked_value else 1
bias_shape = [x_shape[bias_dim]]
if stacked_bias:
bias_shape = [2] + bias_shape
bias = random_ops.random_uniform(bias_shape)
g.watch(bias)
# pylint: disable=cell-var-from-loop
def loop_fn(i):
with g:
a = array_ops.gather(x, i) if stacked_value else x
b = array_ops.gather(bias, i) if stacked_bias else bias
y = nn.bias_add(a, b, data_format=data_format)
loss = math_ops.reduce_sum(y * y)
grad = g.gradient(loss, bias)
if stacked_bias:
# If we gather over bias in loop_fn, the gradient will be an
# instance of `IndexedSlices` with attrs `values` and `indices`.
return y, grad.values, grad.indices
else:
return y, grad
# pylint: enable=cell-var-from-loop
out_dtypes = [dtypes.float32, dtypes.float32]
if stacked_bias:
out_dtypes = out_dtypes + [dtypes.int32]
self._test_loop_fn(loop_fn, 2)
@parameterized.parameters(
*itertools.product(
(
math_ops.unsorted_segment_mean,
math_ops.unsorted_segment_sum,
math_ops.unsorted_segment_min,
math_ops.unsorted_segment_max,
math_ops.unsorted_segment_prod,
),
(
[[0, 0, 2], [0, 1, 2], [2, 2, 2]],
[[0, 0, 2, -1], [0, 1, 2, -1], [2, 2, 2, -1]],
),
)
)
def test_unsorted_segment_reduction(self, reduction_op, indices):
t = random_ops.random_uniform(constant_op.constant(indices).shape + [2])
for segment_ids_dtype in (dtypes.int32, dtypes.int64):
segment_ids = constant_op.constant(indices, dtype=segment_ids_dtype)
for num_segments_dtype in (dtypes.int32, dtypes.int64):
num_segments = constant_op.constant(3, dtype=num_segments_dtype)
# pylint: disable=cell-var-from-loop
def loop_fn(i):
data = array_ops.gather(t, i)
data_0 = array_ops.gather(t, 0)
seg_ids = array_ops.gather(segment_ids, i)
seg_ids_0 = array_ops.gather(segment_ids, 0)
return (reduction_op(data, seg_ids, num_segments),
reduction_op(data_0, seg_ids, num_segments),
reduction_op(data, seg_ids_0, num_segments))
# pylint: enable=cell-var-from-loop
self._test_loop_fn(loop_fn, 3)
@parameterized.parameters((math_ops.sparse_segment_sum_v2, True),
(math_ops.sparse_segment_mean_v2, True),
(math_ops.sparse_segment_sqrt_n_v2, True),
(math_ops.sparse_segment_sum_v2, False),
(math_ops.sparse_segment_mean_v2, False),
(math_ops.sparse_segment_sqrt_n_v2, False))
def test_sparse_segment(self, op_func, with_num_segments):
data = random_ops.random_uniform([3, 4, 2])
indices = constant_op.constant([[1, 2, 3], [0, 1, 2], [0, 2, 3]])
seg_ids = constant_op.constant([[0, 0, 2], [1, 1, 1], [0, 1, 1]])
if with_num_segments:
num_segments = 3
else:
num_segments = None
def loop_fn(i):
data_i = array_ops.gather(data, i)
data_0 = array_ops.gather(data, 0)
indices_i = array_ops.gather(indices, i)
indices_0 = array_ops.gather(indices, 0)
seg_ids_i = array_ops.gather(seg_ids, i)
seg_ids_0 = array_ops.gather(seg_ids, 0)
outputs = [
op_func(data_0, indices_i, seg_ids_0, num_segments=num_segments),
op_func(data_i, indices_i, seg_ids_0, num_segments=num_segments),
op_func(data_0, indices_0, seg_ids_0, num_segments=num_segments),
op_func(data_i, indices_0, seg_ids_0, num_segments=num_segments)
]
if with_num_segments:
# For this case, we support loop variant segment_ids as well.
outputs += [
op_func(data_0, indices_i, seg_ids_i, num_segments=num_segments),
op_func(data_i, indices_i, seg_ids_i, num_segments=num_segments),
op_func(data_0, indices_0, seg_ids_i, num_segments=num_segments),
op_func(data_i, indices_0, seg_ids_i, num_segments=num_segments)
]
return outputs
self._test_loop_fn(loop_fn, 3)
@parameterized.parameters(math_ops.sparse_segment_sum_grad,
math_ops.sparse_segment_mean_grad,
math_ops.sparse_segment_sqrt_n_grad)
def test_sparse_segment_grad(self, op_func):
grad = random_ops.random_uniform([3, 3, 2])
indices = constant_op.constant([1, 2, 3])
seg_ids = constant_op.constant([0, 0, 2])
dim0 = 4
def loop_fn(i):
grad_i = array_ops.gather(grad, i)
return op_func(grad_i, indices, seg_ids, dim0)
self._test_loop_fn(loop_fn, 3)
def test_cast(self):
x = constant_op.constant([[1], [2]])
y = constant_op.constant([[1.0], [2.0]])
def loop_fn(i):
return (math_ops.cast(array_ops.gather(x, i), dtypes.float32),
math_ops.cast(array_ops.gather(y, i), dtypes.int32))
self._test_loop_fn(loop_fn, 2)
def test_tanh_axpy(self):
a = constant_op.constant(3.)
x = random_ops.random_uniform([4, 5])
y = random_ops.random_uniform([6, 5])
n = x.shape[0]
def loop_fn(i):
return math_ops.tanh(a * array_ops.gather(x, i) + array_ops.gather(y, i))
self._test_loop_fn(loop_fn, n)
def test_select(self):
a = random_ops.random_uniform([2, 3, 5])
b = random_ops.random_uniform([2, 3, 5])
for cond_shape in [2], [2, 3], [2, 3, 5]:
cond = random_ops.random_uniform(cond_shape) > 0.5
# pylint: disable=cell-var-from-loop
def loop_fn(i):
a_i = array_ops.gather(a, i)
b_i = array_ops.gather(b, i)
cond_i = array_ops.gather(cond, i)
return array_ops.where(cond_i, a_i, b_i)
# pylint: enable=cell-var-from-loop
self._test_loop_fn(loop_fn, 2)
def test_selectv2_cond_needs_broadcast(self):
a = random_ops.random_uniform([2, 3, 5])
b = random_ops.random_uniform([2, 3, 5])
# wherev2 assumes all shapes are broadcastable with each other.
# This means that we can only specify conditions that are
# broadcastable with [3, 5].
for cond_shape in [2], [2, 1], [2, 5], [2, 3, 1], [2, 3, 5]:
cond = random_ops.random_uniform(cond_shape) > 0.5
# pylint: disable=cell-var-from-loop
def loop_fn(i):
a_i = array_ops.gather(a, i)
b_i = array_ops.gather(b, i)
cond_i = array_ops.gather(cond, i)
return array_ops.where_v2(cond_i, a_i, b_i)
# pylint: enable=cell-var-from-loop
self._test_loop_fn(loop_fn, 2)
def test_selectv2_args_need_broadcast(self):
a = random_ops.random_uniform([2, 5])
b = random_ops.random_uniform([2, 3, 5])
# wherev2 assumes all shapes are broadcastable with each other.
# This means that we can only specify conditions that are
# broadcastable with [3, 5].
for cond_shape in [2], [2, 1], [2, 5], [2, 3, 1], [2, 3, 5]:
cond = random_ops.random_uniform(cond_shape) > 0.5
# pylint: disable=cell-var-from-loop
def loop_fn(i):
a_i = array_ops.gather(a, i)
b_i = array_ops.gather(b, i)
cond_i = array_ops.gather(cond, i)
return array_ops.where_v2(cond_i, a_i, b_i)
# pylint: enable=cell-var-from-loop
self._test_loop_fn(loop_fn, 2)
def test_selectv2_cond_fixed(self):
cond = random_ops.random_uniform([3, 5]) > 0.5
b = random_ops.random_uniform([2, 3, 5])
# wherev2 assumes all shapes are broadcastable with each other.
# This means that we can only specify conditions that are
# broadcastable with [3, 5].
for a_shape in [2], [2, 1], [2, 5], [2, 3, 1], [2, 3, 5]:
a = random_ops.random_uniform(a_shape)
# pylint: disable=cell-var-from-loop
def loop_fn(i):
a_i = array_ops.gather(a, i)
b_i = array_ops.gather(b, i)
return array_ops.where_v2(cond, a_i, b_i)
# pylint: enable=cell-var-from-loop
self._test_loop_fn(loop_fn, 2)
@test_util.run_all_in_graph_and_eager_modes
class LinalgTest(PForTestCase):
def test_cholesky(self):
z = random_ops.random_normal([2, 3, 3])
x = (
math_ops.matmul(z, array_ops.matrix_transpose(z)) # Ensure pos. def.
+ linalg_ops.eye(3)) # Ensure well-conditioned.
def loop_fn(i):
return linalg_ops.cholesky(array_ops.gather(x, i))
self._test_loop_fn(loop_fn, 2)
def test_log_matrix_determinant(self):
for x_shape in ([3, 4, 2, 2], [3, 2, 2]):
x = random_ops.random_normal(x_shape)
# pylint: disable=cell-var-from-loop
def loop_fn(i):
return linalg_ops.log_matrix_determinant(array_ops.gather(x, i))
# pylint: enable=cell-var-from-loop
self._test_loop_fn(loop_fn, 3)
def test_matrix_inverse(self):
x = (random_ops.random_uniform([3, 4, 2, 2]) + 10 * linalg_ops.eye(2)
) # Ensure well-conditioned.
for adjoint in (True, False):
# pylint: disable=cell-var-from-loop
def loop_fn(i):
return linalg_ops.matrix_inverse(
array_ops.gather(x, i), adjoint=adjoint)
# pylint: enable=cell-var-from-loop
self._test_loop_fn(loop_fn, 2)
def test_matrix_solve(self):
for adjoint in (True, False):
for stack_a in (True, False):
for stack_b in (True, False):
shape_a = (2, 4, 3, 3) if stack_a else (4, 3, 3)
shape_b = (2, 4, 3, 5) if stack_b else (4, 3, 5)
x = (random_ops.random_uniform(shape_a) + 10 * linalg_ops.eye(3)
) # Ensure well-conditioned.
y = random_ops.random_uniform(shape_b)
# pylint: disable=cell-var-from-loop
def loop_fn(i):
a = array_ops.gather(x, i) if stack_a else x
b = array_ops.gather(y, i) if stack_b else y
return linalg_ops.matrix_solve(a, b, adjoint=adjoint)
# pylint: enable=cell-var-from-loop
self._test_loop_fn(loop_fn, 2)
def test_matrix_triangular_solve(self):
for lower in (True, False):
for adjoint in (True, False):
for stack_a in (True, False):
for stack_b in (True, False):
shape_a = (2, 4, 3, 3) if stack_a else (4, 3, 3)
shape_b = (2, 4, 3, 5) if stack_b else (4, 3, 5)
x = array_ops.matrix_band_part(
random_ops.random_uniform(shape_a) +
linalg_ops.eye(3), # Ensure well-conditioned.
*((-1, 0) if lower else (0, -1))) # Ensure triangular.
y = random_ops.random_uniform(shape_b)
# pylint: disable=cell-var-from-loop
def loop_fn(i):
a = array_ops.gather(x, i) if stack_a else x
b = array_ops.gather(y, i) if stack_b else y
return linalg_ops.matrix_triangular_solve(
a, b, lower=lower, adjoint=adjoint)
# pylint: enable=cell-var-from-loop
self._test_loop_fn(loop_fn, 2)
def test_self_adjoint_eig(self):
z = random_ops.random_normal([2, 3, 3])
x = z + array_ops.matrix_transpose(z) # Ensure self-adjoint.
def loop_fn(i):
return (linalg_ops.self_adjoint_eig(array_ops.gather(x, i)),
linalg_ops.self_adjoint_eigvals(array_ops.gather(x, i)))
self._test_loop_fn(loop_fn, 2)
@test_util.run_without_tensor_float_32(
"Calls einsum in parallel for-loop and compares result to calling einsum "
"in sequential for-loop")
def test_einsum(self):
b = 10
x_series = random_ops.random_uniform([b, 9, 9])
y_series = random_ops.random_uniform([b, 9, 1])
def loop_fn(i):
x = array_ops.gather(x_series, 0) # invariant.
y = array_ops.gather(y_series, 0) # invariant.
x_i = array_ops.gather(x_series, i)
y_i = array_ops.gather(y_series, i)
z0 = special_math_ops.einsum("ab->b", x_i)
z1 = special_math_ops.einsum("ab,bc->ac", x_i, y)
z2 = special_math_ops.einsum("ab,bc->ac", x, y_i)
z3 = special_math_ops.einsum("ab,bc->ac", x, y)
z4 = special_math_ops.einsum("ab,bc->ac", x_i, y_i)
z5 = special_math_ops.einsum("cd,ce->de", y_i, x_i) # Includes transpose.
outputs = [z0, z1, z2, z3, z4, z5]
return outputs
self._test_loop_fn(loop_fn, b)
if __name__ == "__main__":
test.main()
tensorflow/python/ops/parallel_for/pfor.py
+5229
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/ops/parallel_for/pfor_test.py
+77
View File
@@ -0,0 +1,77 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import random_ops
from tensorflow.python.ops.parallel_for import pfor
from tensorflow.python.platform import test
class PForTest(test.TestCase):
def test_rank_known(self):
with ops.Graph().as_default():
x = array_ops.placeholder(dtypes.float32, [None, None])
rank = pfor._rank(x)
self.assertIsInstance(rank, int)
self.assertEqual(rank, 2)
def test_rank_unknown(self):
with ops.Graph().as_default():
x = array_ops.placeholder(dtypes.float32)
rank = pfor._rank(x)
self.assertIsInstance(rank, tensor.Tensor)
def test_size_known(self):
with ops.Graph().as_default():
x = array_ops.placeholder(dtypes.float32, [3, 5])
size = pfor._size(x)
self.assertIsInstance(size, int)
self.assertEqual(size, 3 * 5)
def test_size_unknown(self):
with ops.Graph().as_default():
x = array_ops.placeholder(dtypes.float32, [3, None])
size = pfor._size(x, dtypes.int32)
self.assertIsInstance(size, tensor.Tensor)
self.assertEqual(size.dtype, dtypes.int32)
size = pfor._size(x, dtypes.int64)
self.assertIsInstance(size, tensor.Tensor)
self.assertEqual(size.dtype, dtypes.int64)
def test_expand_dims_static(self):
x = random_ops.random_uniform([3, 5])
axis = 1
num_axes = 2
expected = array_ops.reshape(x, [3, 1, 1, 5])
actual = pfor._expand_dims(x, axis, num_axes)
self.assertAllEqual(expected, actual)
def test_expand_dims_dynamic(self):
x = random_ops.random_uniform([3, 5])
axis = 1
num_axes = constant_op.constant([2])
expected = array_ops.reshape(x, [3, 1, 1, 5])
actual = pfor._expand_dims(x, axis, num_axes)
self.assertAllEqual(expected, actual)
if __name__ == '__main__':
test.main()
tensorflow/python/ops/parallel_for/test_util.py
+76
View File
@@ -0,0 +1,76 @@
# 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 utility."""
import numpy as np
from tensorflow.python.ops import variables
from tensorflow.python.ops.parallel_for import control_flow_ops as pfor_control_flow_ops
from tensorflow.python.platform import test
from tensorflow.python.util import nest
class PForTestCase(test.TestCase):
"""Base class for test cases."""
def _run_targets(self, targets1, targets2=None, run_init=True):
targets1 = nest.flatten(targets1)
targets2 = ([] if targets2 is None else nest.flatten(targets2))
assert len(targets1) == len(targets2) or not targets2
if run_init:
init = variables.global_variables_initializer()
self.evaluate(init)
return self.evaluate(targets1 + targets2)
# TODO(agarwal): Allow tests to pass down tolerances.
def run_and_assert_equal(self, targets1, targets2, rtol=1e-4, atol=1e-5):
outputs = self._run_targets(targets1, targets2)
outputs = nest.flatten(outputs) # flatten SparseTensorValues
n = len(outputs) // 2
for i in range(n):
if outputs[i + n].dtype != np.object_:
self.assertAllClose(outputs[i + n], outputs[i], rtol=rtol, atol=atol)
else:
self.assertAllEqual(outputs[i + n], outputs[i])
def _test_loop_fn(self,
loop_fn,
iters,
parallel_iterations=None,
fallback_to_while_loop=False,
rtol=1e-4,
atol=1e-5):
t1 = pfor_control_flow_ops.pfor(
loop_fn,
iters=iters,
fallback_to_while_loop=fallback_to_while_loop,
parallel_iterations=parallel_iterations)
loop_fn_dtypes = nest.map_structure(lambda x: x.dtype, t1)
t2 = pfor_control_flow_ops.for_loop(loop_fn, loop_fn_dtypes, iters=iters,
parallel_iterations=parallel_iterations)
def _check_shape(a, b):
msg = (
"Inferred static shapes are different between two loops:"
f" {a.shape} vs {b.shape}."
)
# TODO(b/268146947): should assert bool(a.shape) == bool(b.shape),
# since both should be either defined or undefined. But it does not work.
if b.shape:
self.assertEqual(a.shape.as_list()[0], b.shape.as_list()[0], msg)
# TODO(b/268146947): self.assertShapeEqual(a, b, msg) does not work.
nest.map_structure(_check_shape, t1, t2)
self.run_and_assert_equal(t1, t2, rtol=rtol, atol=atol)
tensorflow/python/ops/parallel_for/xla_control_flow_ops_test.py
+244
View File
@@ -0,0 +1,244 @@
# 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.
# ==============================================================================
"""XLA tests for pfor."""
# pylint: disable=g-direct-tensorflow-import
from tensorflow.compiler.tf2xla.python import xla as xla_ops
from tensorflow.python.compiler.xla import jit
from tensorflow.python.compiler.xla import xla
from tensorflow.python.eager import context
from tensorflow.python.eager import def_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_v2_toggles
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import random_ops
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import while_loop
from tensorflow.python.ops.parallel_for import control_flow_ops as pfor_control_flow_ops
from tensorflow.python.ops.parallel_for.test_util import PForTestCase
from tensorflow.python.platform import test
@test_util.run_all_in_graph_and_eager_modes
class PForTest(PForTestCase):
def __init__(self, method_name="runTest"):
super(PForTest, self).__init__(method_name)
context.context().enable_xla_devices()
def test_xla_einsum(self):
num_loop = 10
x_series = random_ops.random_uniform([num_loop, 9, 9])
y_series = random_ops.random_uniform([num_loop, 9, 1])
def loop_fn(i):
x = array_ops.gather(x_series, 0) # invariant.
y = array_ops.gather(y_series, 0) # invariant.
x_i = array_ops.gather(x_series, i)
y_i = array_ops.gather(y_series, i)
z1 = xla_ops.einsum(x_i, y, "ab,bc->ac")
z2 = xla_ops.einsum(x, y_i, "ab,bc->ac")
z3 = xla_ops.einsum(x, y, "ab,bc->ac")
z4 = xla_ops.einsum(x_i, y_i, "ab,bc->ac")
z5 = xla_ops.einsum(y_i, x_i, "cd,ce->de") # Includes transpose.
outputs = [z1, z2, z3, z4, z5]
return outputs
self._test_loop_fn(loop_fn, num_loop)
def test_xla(self):
def compute(x):
return math_ops.reduce_mean(x, axis=0, keepdims=True)
def vectorized_compute(x):
return pfor_control_flow_ops.vectorized_map(compute, x)
result = xla.compile(
vectorized_compute, inputs=[array_ops.ones((10, 5, 3))])
self.run_and_assert_equal(result, array_ops.ones((10, 1, 3)))
def test_function_jit_compile(self):
def compute(x):
return math_ops.reduce_mean(x, axis=0, keepdims=True)
@def_function.function(jit_compile=True)
def vectorized_compute(x):
return pfor_control_flow_ops.vectorized_map(compute, x)
result = vectorized_compute(array_ops.ones((10, 5, 3)))
self.run_and_assert_equal(result, array_ops.ones((10, 1, 3)))
def test_xla_while_loop(self):
def compute(x):
return math_ops.reduce_mean(x, axis=0, keepdims=True)
def vectorized_compute(x, i):
inp = array_ops.gather(x, i)
output = pfor_control_flow_ops.vectorized_map(compute, inp)
output.set_shape([5, 1])
return output
def while_compute(x):
return while_loop.while_loop_v2(
lambda i, _: i < 10,
lambda i, y: (i + 1, y + vectorized_compute(x, i)),
(0, array_ops.zeros([5, 1])))[1]
result = xla.compile(while_compute, inputs=[array_ops.ones((10, 5, 3))])
expected = array_ops.ones([5, 1]) * 10
self.run_and_assert_equal(expected, result)
def test_reduce_mean(self):
x = random_ops.random_uniform([8, 3])
@def_function.function(jit_compile=True)
def f():
def loop_fn(i, pfor_config):
x_i = array_ops.gather(x, i)
return x_i - pfor_config.reduce_mean(x_i)
return pfor_control_flow_ops.pfor(loop_fn, 8)
output = f()
ans = x - math_ops.reduce_mean(x, axis=0)
output_val, ans_val = self.evaluate([output, ans])
self.assertAllClose(ans_val, output_val)
def _make_unstacked(cond, body, pfor_config):
def _cond(*args):
return math_ops.reduce_any(pfor_config.reduce_concat(args[0]))
def _body(*args):
not_done = args[0]
args = args[1:]
not_done = math_ops.logical_and(not_done, cond(*args))
outputs = body(*args)
return (not_done,) + tuple(
array_ops.where_v2(not_done, x, y) for x, y in zip(outputs, args))
return _cond, _body
@test_util.run_all_in_graph_and_eager_modes
class WhileV2Test(PForTestCase):
def setUp(self):
self._enabled = control_flow_v2_toggles.control_flow_v2_enabled()
control_flow_v2_toggles.enable_control_flow_v2()
super(WhileV2Test, self).setUp()
def tearDown(self):
if not self._enabled:
control_flow_v2_toggles.disable_control_flow_v2()
super(WhileV2Test, self).tearDown()
def _test_loop_fn(self, loop_fn, iters, force_xla=False):
def f():
return pfor_control_flow_ops.pfor(loop_fn, iters)
@def_function.function
def jit_f():
with jit.experimental_jit_scope():
return f()
out = f()
jit_out = jit_f()
self.run_and_assert_equal(out, jit_out)
# TODO(agarwal): The following may complain about uncompilable nodes. Hence
# these are currently not enabled for all tests.
if force_xla:
out_exp_compile_f = def_function.function(jit_compile=True)(f)()
self.run_and_assert_equal(out, out_exp_compile_f)
out_xla_compile_f = xla.compile(f, inputs=[])
self.run_and_assert_equal(out, out_xla_compile_f)
def test_stateless_while(self):
x = random_ops.random_uniform([3, 5])
lengths = constant_op.constant([4, 0, 2])
def loop_fn(i):
x_i = array_ops.gather(x, i)
lengths_i = array_ops.gather(lengths, i)
return while_loop.while_loop(
lambda j, _: j < lengths_i,
lambda j, t: (j + 1, t + array_ops.gather(x_i, j)),
[0, 0.])
self._test_loop_fn(loop_fn, 3)
def test_while_with_variable(self):
if not context.executing_eagerly():
self.skipTest("Flaky with tf.Session")
v = resource_variable_ops.ResourceVariable(5.)
def loop_fn(_):
_, output = while_loop.while_loop(
lambda j, x: j < 4,
lambda j, x: (j + 1, x + v),
[0, 0.])
return output
self._test_loop_fn(loop_fn, 3)
def test_while_unstacked_condition(self):
def loop_fn(i):
return while_loop.while_loop(
lambda j, x: j < 4,
lambda j, x: (j + 1, x + i), [0, 0])
self._test_loop_fn(loop_fn, 3, force_xla=True)
def test_while_force_unstacked_condition(self):
# The while_loop in this setup is similar to the one in test_stateless_while
# whose condition is loop variant. However here we wrap the cond and body of
# the loop in a way that makes the while_loop condition pfor loop invariant.
# This allows xla compilation to work since the vectorized code no longer
# needs to perform dynamic partitioning of the inputs.
x = random_ops.random_uniform([3, 5])
lengths = constant_op.constant([4, 0, 2])
def loop_fn(i, pfor_config):
x_i = array_ops.gather(x, i)
lengths_i = array_ops.gather(lengths, i)
def _cond(j, _):
return j < lengths_i
def _body(j, t):
return (j + 1, t + array_ops.gather(x_i, j))
cond, body = _make_unstacked(_cond, _body, pfor_config)
return while_loop.while_loop(
cond,
body,
[True, 0, 0.])
self._test_loop_fn(loop_fn, 3, force_xla=True)
if __name__ == "__main__":
test.main()