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/eager/polymorphic_function/BUILD
+638
View File
@@ -0,0 +1,638 @@
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test", "tf_py_strict_test")
load("//tensorflow/compiler/tests:build_defs.bzl", "tf_xla_py_strict_test")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
py_library(
name = "attributes",
srcs = ["attributes.py"],
strict_deps = True,
visibility = ["//tensorflow/python:__subpackages__"],
deps = [
"//tensorflow/core:protos_all_py",
"//tensorflow/python/util:compat",
],
)
py_library(
name = "autograph_util",
srcs = ["autograph_util.py"],
strict_deps = True,
deps = [
"//tensorflow/python/autograph/core:converter",
"//tensorflow/python/autograph/impl:api",
"//tensorflow/python/util:tf_decorator_py",
],
)
py_library(
name = "transform",
srcs = ["transform.py"],
strict_deps = True,
visibility = ["//tensorflow/python:__subpackages__"],
deps = [],
)
pytype_strict_library(
name = "atomic_function",
srcs = ["atomic_function.py"],
visibility = ["//tensorflow/python:__subpackages__"],
deps = [
":attributes",
":function_type_utils",
"//tensorflow/core:protos_all_py",
"//tensorflow/core/function/polymorphism:function_type",
"//tensorflow/python/client:pywrap_tf_session",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:record",
"//tensorflow/python/framework:auto_control_deps_utils",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:error_interpolation",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:func_graph",
"//tensorflow/python/framework:function_def_to_graph",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:handle_data_util",
"//tensorflow/python/types:core",
"//tensorflow/python/util:compat",
"//tensorflow/python/util:function_utils",
"//tensorflow/python/util:tf_stack",
],
)
cuda_py_strict_test(
name = "atomic_function_test",
size = "medium",
srcs = ["atomic_function_test.py"],
deps = [
":atomic_function",
":polymorphic_function",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/platform:client_testlib",
],
)
py_library(
name = "concrete_function",
srcs = ["concrete_function.py"],
strict_deps = True,
visibility = ["//tensorflow/python:__subpackages__"],
deps = [
":atomic_function",
":attributes",
":function_type_utils",
":saved_model_exported_concrete",
"//tensorflow/core:protos_all_py",
"//tensorflow/core/function/polymorphism:function_type",
"//tensorflow/python:pywrap_tfe",
"//tensorflow/python/eager:backprop_util",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:forwardprop_util",
"//tensorflow/python/eager:graph_only_ops",
"//tensorflow/python/eager:record",
"//tensorflow/python/framework:composite_tensor",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:func_graph",
"//tensorflow/python/framework:indexed_slices",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/framework:type_spec",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:default_gradient",
"//tensorflow/python/ops:gradients_util",
"//tensorflow/python/ops:handle_data_util",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/profiler:trace",
"//tensorflow/python/trackable:base",
"//tensorflow/python/types:core",
"//tensorflow/python/util:compat",
"//tensorflow/python/util:nest",
"//tensorflow/python/util:object_identity",
],
)
cuda_py_strict_test(
name = "concrete_function_test",
size = "small",
srcs = ["concrete_function_test.py"],
deps = [
":atomic_function",
":concrete_function",
":polymorphic_function",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:func_graph",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/ops/ragged:ragged_tensor",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/util:compat",
"@absl_py//absl/testing:parameterized",
],
)
py_library(
name = "tracing_compilation",
srcs = ["tracing_compilation.py"],
strict_deps = True,
visibility = ["//tensorflow/python:__subpackages__"],
deps = [
":attributes",
":concrete_function",
":function_context",
":function_type_utils",
":tf_method_target",
":transform",
"//tensorflow/core/function/capture:capture_container",
"//tensorflow/core/function/polymorphism:function_cache",
"//tensorflow/core/function/polymorphism:function_type",
"//tensorflow/core/function/trace_type",
"//tensorflow/python/autograph/core:ag_ctx",
"//tensorflow/python/eager:monitoring",
"//tensorflow/python/framework:func_graph",
"//tensorflow/python/framework:ops",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/profiler:trace",
"//tensorflow/python/util:compat",
],
)
py_library(
name = "tf_method_target",
srcs = ["tf_method_target.py"],
strict_deps = True,
visibility = [
"//tensorflow/python/autograph/impl:__pkg__",
"//tensorflow/python/eager:__pkg__",
],
deps = [
"//tensorflow/python/util:tf_inspect",
],
)
py_library(
name = "polymorphic_function",
srcs = ["polymorphic_function.py"],
strict_deps = True,
visibility = ["//tensorflow:internal"],
deps = [
":attributes",
":autograph_util",
":compiler_ir",
":eager_function_run",
":function_type_utils",
":tf_method_target",
":tracing_compilation",
"//tensorflow/core:protos_all_py",
"//tensorflow/core/function/capture:capture_container",
"//tensorflow/core/function/polymorphism:function_cache",
"//tensorflow/core/function/trace_type",
"//tensorflow/python/distribute/parallel_device",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:lift_to_graph",
"//tensorflow/python/eager:monitoring",
"//tensorflow/python/framework:composite_tensor",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:func_graph",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:cond",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/ops:control_flow_util",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/profiler:trace",
"//tensorflow/python/trackable:base",
"//tensorflow/python/types:core",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:nest",
"//tensorflow/python/util:object_identity",
"//tensorflow/python/util:tf_decorator_py",
"//tensorflow/python/util:tf_export",
"//tensorflow/python/util:traceback_utils",
],
)
cuda_py_strict_test(
name = "polymorphic_function_test",
size = "medium",
srcs = ["polymorphic_function_test.py"],
shard_count = 15,
tags = [
"nomac", # b/157056289
],
deps = [
":attributes",
":polymorphic_function",
"//tensorflow/core/function/capture:capture_container",
"//tensorflow/core/function/trace_type",
"//tensorflow/python/autograph/core:ag_ctx",
"//tensorflow/python/autograph/core:converter",
"//tensorflow/python/autograph/lang:directives",
"//tensorflow/python/checkpoint",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/ops:iterator_ops",
"//tensorflow/python/eager:backprop",
"//tensorflow/python/eager:cancellation",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:lift_to_graph",
"//tensorflow/python/framework:composite_tensor",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:extension_type",
"//tensorflow/python/framework:function",
"//tensorflow/python/framework:indexed_slices",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:random_seed",
"//tensorflow/python/framework:sparse_tensor",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/framework:test_ops",
"//tensorflow/python/framework:type_spec",
"//tensorflow/python/module",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:check_ops",
"//tensorflow/python/ops:clip_ops",
"//tensorflow/python/ops:cond_v2",
"//tensorflow/python/ops:control_flow_assert",
"//tensorflow/python/ops:data_flow_ops",
"//tensorflow/python/ops:functional_ops",
"//tensorflow/python/ops:gradients_impl",
"//tensorflow/python/ops:list_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/ops:random_ops_gen",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:script_ops",
"//tensorflow/python/ops:sendrecv_ops_gen",
"//tensorflow/python/ops:string_ops",
"//tensorflow/python/ops:training_ops_gen",
"//tensorflow/python/ops:variable_scope",
"//tensorflow/python/ops:variables",
"//tensorflow/python/ops/ragged:ragged_factory_ops",
"//tensorflow/python/ops/ragged:ragged_tensor",
"//tensorflow/python/ops/structured:structured_tensor",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/saved_model:load",
"//tensorflow/python/saved_model:save",
"//tensorflow/python/saved_model:save_context",
"//tensorflow/python/saved_model:save_options",
"//tensorflow/python/util:compat",
"//tensorflow/python/util:nest",
"//tensorflow/python/util:tf_decorator_py",
"@absl_py//absl/testing:parameterized",
],
)
tf_py_strict_test(
name = "polymorphic_function_test_cpu_only",
srcs = ["polymorphic_function_test_cpu_only.py"],
# --config=cuda implicitly links in XLA.
tags = [
"no_cuda_on_cpu_tap",
"no_gpu",
"no_oss", # No way to force no XLA linkage in OSS build from here.
"no_pip",
],
deps = [
":polymorphic_function",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/platform:client_testlib",
"@absl_py//absl/testing:parameterized",
],
)
tf_xla_py_strict_test(
name = "polymorphic_function_xla_jit_test",
srcs = ["polymorphic_function_xla_jit_test.py"],
enabled_backends = [
"cpu",
],
tags = [
"no_mac",
"no_oss", # TODO(b/295654746)
"no_pip",
"no_windows",
],
use_xla_device = False,
deps = [
":polymorphic_function",
"//tensorflow/compiler/tests:xla_test",
"//tensorflow/python/eager:backprop",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:collective_ops",
"//tensorflow/python/ops:cond",
"//tensorflow/python/ops:control_flow_assert",
"//tensorflow/python/ops:control_flow_util",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:string_ops",
"//tensorflow/python/ops:summary_ops_v2",
"//tensorflow/python/ops:tensor_array_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/ops:while_loop",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/util:nest",
],
)
tf_xla_py_strict_test(
name = "polymorphic_function_xla_test",
srcs = ["polymorphic_function_xla_test.py"],
tags = [
"no_pip",
"no_windows",
"nomac",
],
deps = [
":polymorphic_function",
"//tensorflow/compiler/tests:xla_test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "gradients_test",
size = "medium",
srcs = ["gradients_test.py"],
shard_count = 5,
deps = [
":polymorphic_function",
"//tensorflow/python/eager:backprop",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:indexed_slices",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:cond",
"//tensorflow/python/ops:gradients_impl",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nn_grad",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:variable_scope",
"//tensorflow/python/ops:variables",
"//tensorflow/python/ops:while_loop",
"//tensorflow/python/platform:client_testlib",
"@absl_py//absl/testing:parameterized",
],
)
py_library(
name = "composite_tensor_utils",
srcs = ["composite_tensor_utils.py"],
strict_deps = True,
visibility = ["//tensorflow:internal"],
deps = [
"//tensorflow/python/framework:composite_tensor",
"//tensorflow/python/util:_pywrap_utils",
"//tensorflow/python/util:nest",
],
)
tf_py_strict_test(
name = "tracing_compilation_test",
size = "medium",
srcs = ["tracing_compilation_test.py"],
deps = [
":function_type_utils",
":tracing_compilation",
"//tensorflow/core:protos_all_py",
"//tensorflow/core/function/capture:capture_container",
"//tensorflow/core/function/polymorphism:function_cache",
"//tensorflow/python/eager:backprop",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/framework:test_ops",
"//tensorflow/python/module",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:functional_ops_gen",
"//tensorflow/python/ops:gradients_impl",
"//tensorflow/python/ops:init_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:resource_variable_ops_gen",
"//tensorflow/python/ops:variable_scope",
"//tensorflow/python/ops:variables",
"//tensorflow/python/ops:while_loop",
"//tensorflow/python/ops/ragged:ragged_factory_ops",
"//tensorflow/python/ops/ragged:ragged_tensor",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/saved_model:load",
"//tensorflow/python/saved_model:save",
"//tensorflow/python/util:compat",
"//tensorflow/python/util:nest",
"@absl_py//absl/testing:parameterized",
],
)
cuda_py_strict_test(
name = "argument_naming_test",
size = "medium",
srcs = ["argument_naming_test.py"],
deps = [
":polymorphic_function",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"@absl_py//absl/testing:parameterized",
],
)
cuda_py_strict_test(
name = "collection_test",
size = "medium",
srcs = ["collection_test.py"],
deps = [
":polymorphic_function",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
],
)
py_library(
name = "eager_function_run",
srcs = ["eager_function_run.py"],
strict_deps = True,
visibility = ["//tensorflow:internal"],
deps = [
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "function_context",
srcs = ["function_context.py"],
strict_deps = True,
visibility = ["//visibility:private"],
deps = [
"//tensorflow/core/function/polymorphism:function_cache",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:device",
"//tensorflow/python/framework:func_graph",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/saved_model:save_context",
],
)
py_library(
name = "saved_model_utils",
srcs = ["saved_model_utils.py"],
strict_deps = True,
visibility = ["//tensorflow:internal"],
deps = [
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_util",
"//tensorflow/python/saved_model/registration",
"//tensorflow/python/trackable:base",
"//third_party/py/numpy",
],
)
py_library(
name = "saved_model_exported_concrete",
srcs = ["saved_model_exported_concrete.py"],
strict_deps = True,
visibility = ["//tensorflow:internal"],
deps = [
":function_type_utils",
"//tensorflow/python/trackable:base",
],
)
py_library(
name = "function_type_utils",
srcs = ["function_type_utils.py"],
strict_deps = True,
visibility = ["//tensorflow:internal"],
deps = [
":composite_tensor_utils",
"//tensorflow/core/function/polymorphism:function_type",
"//tensorflow/core/function/trace_type",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:type_spec",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/util:nest",
"//third_party/py/numpy",
"@six_archive//:six",
],
)
tf_py_strict_test(
name = "function_spec_test",
size = "medium",
srcs = ["function_spec_test.py"],
deps = [
":function_type_utils",
"//tensorflow/core/function/polymorphism:function_type",
"//tensorflow/core/function/trace_type",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/util:tf_decorator_py",
"@absl_py//absl/testing:parameterized",
],
)
py_library(
name = "compiler_ir",
srcs = ["compiler_ir.py"],
strict_deps = True,
visibility = ["//visibility:private"],
deps = [
"//tensorflow/core/function/trace_type",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/util:nest",
],
)
tf_xla_py_strict_test(
name = "compiler_ir_test",
srcs = ["compiler_ir_test.py"],
disabled_backends = [
"cpu_ondemand",
"gpu_a100",
"gpu_h100",
],
tags = [
"no_mac",
"no_pip",
"no_windows",
],
use_xla_device = False,
deps = [
":compiler_ir",
":polymorphic_function",
"//tensorflow/compiler/tests:xla_test",
"//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:array_ops_gen",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/util:nest",
],
)
tensorflow/python/eager/polymorphic_function/argument_naming_test.py
+242
View File
@@ -0,0 +1,242 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
from absl.testing import parameterized
from tensorflow.core.protobuf import config_pb2
from tensorflow.python.eager.polymorphic_function import polymorphic_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_spec
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
class ArgumentNamingTests(test.TestCase, parameterized.TestCase):
"""Tests for recognizable export signatures from concrete functions."""
def testBasic(self):
@polymorphic_function.function
def fn(a, b):
return a + b, a * b
# Call the function to make def_function happy
fn(array_ops.ones([]), array_ops.ones([]))
fn_op = fn.get_concrete_function(
tensor_spec.TensorSpec(shape=(None,), dtype=dtypes.float32),
tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32))
self.assertEqual(
['a', 'b'],
[inp.op.name for inp in fn_op.inputs])
self.assertEqual(
[b'a', b'b'],
[inp.op.get_attr('_user_specified_name') for inp in fn_op.inputs])
self.assertLen(fn_op.graph.structured_outputs, 2)
self.assertAllClose(
[3., 2.],
fn_op(constant_op.constant(1.), constant_op.constant(2.)))
self.assertAllClose(
[3., 2.],
fn_op(a=constant_op.constant(1.), b=constant_op.constant(2.)))
def testVariable(self):
@polymorphic_function.function
def fn(a, b):
return a + b, a * b
# Call the function to make def_function happy
fn(array_ops.ones([]), array_ops.ones([]))
fn_op = fn.get_concrete_function(
tensor_spec.TensorSpec(shape=(None,), dtype=dtypes.float32),
variables.Variable(1.))
self.assertEqual(
['a', 'b'],
[inp.op.name for inp in fn_op.inputs])
self.assertEqual(
[b'a', b'b'],
[inp.op.get_attr('_user_specified_name') for inp in fn_op.inputs])
self.assertLen(fn_op.graph.structured_outputs, 2)
def testDictReturned(self):
@polymorphic_function.function
def fn(x, z=(1., 2.), y=3.):
z1, z2 = z
return {'alpha': x + y + z1, 'beta': x * y + z2}
# Call the function to make def_function happy
fn(array_ops.ones([]))
fn_op = fn.get_concrete_function(
x=tensor_spec.TensorSpec(shape=(None,), dtype=dtypes.float32),
y=tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32))
self.assertEqual(
['x', 'y'],
[inp.op.name for inp in fn_op.inputs])
self.assertEqual(
[b'x', b'y'],
[inp.op.get_attr('_user_specified_name') for inp in fn_op.inputs])
self.assertEqual({'alpha', 'beta'},
set(fn_op.graph.structured_outputs.keys()))
fn_op2 = fn.get_concrete_function(
z=(tensor_spec.TensorSpec(shape=(None,), dtype=dtypes.float32,
name='z_first'),
tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32,
name='z_second')),
y=tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32, name='custom'),
x=4.)
self.assertEqual(
['z_first', 'z_second', 'custom'],
[inp.op.name for inp in fn_op2.inputs])
self.assertEqual(
[b'z_first', b'z_second', b'custom'],
[inp.op.get_attr('_user_specified_name') for inp in fn_op2.inputs])
fn_op3 = fn.get_concrete_function(
tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32, name='custom'),
z=(tensor_spec.TensorSpec(shape=(None,), dtype=dtypes.float32,
name='z1'),
tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32, name='z2')),
y=tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32))
self.assertEqual(
['custom', 'z1', 'z2', 'y'],
[inp.op.name for inp in fn_op3.inputs])
self.assertEqual(
[b'custom', b'z1', b'z2', b'y'],
[inp.op.get_attr('_user_specified_name') for inp in fn_op3.inputs])
def testMethod(self):
class HasMethod(object):
@polymorphic_function.function
def method(self, x):
return x
has_method = HasMethod()
# Call the function to make def_function happy
HasMethod.method(has_method, array_ops.ones([]))
class_op = HasMethod.method.get_concrete_function(
has_method, tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32))
self.assertEqual(
['x'],
[inp.op.name for inp in class_op.inputs])
self.assertEqual(
[b'x'],
[inp.op.get_attr('_user_specified_name') for inp in class_op.inputs])
# Call the function to make def_function happy
has_method.method(array_ops.ones([]))
method_op = has_method.method.get_concrete_function(
tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32))
self.assertEqual(
['x'],
[inp.op.name for inp in method_op.inputs])
self.assertEqual(
[b'x'],
[inp.op.get_attr('_user_specified_name') for inp in method_op.inputs])
# TODO(allenl): It should be possible to override names when exporting. Do
# TensorSpec names need to go in cache keys? Or maybe get_concrete_function
# should always retrace?
self.skipTest('Not working')
method_op = has_method.method.get_concrete_function(
tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32, name='y'))
self.assertEqual(
['y'],
[inp.op.name for inp in method_op.inputs])
self.assertEqual(
[b'y'],
[inp.op.get_attr('_user_specified_name') for inp in method_op.inputs])
def testMethodSignature(self):
class HasMethod(object):
@polymorphic_function.function(
input_signature=(tensor_spec.TensorSpec(
shape=None, dtype=dtypes.float64, name='y'),))
def method(self, x):
hash(self) # No weak proxies passed as `self`
return x
has_method = HasMethod()
# Call the function to make def_function happy
has_method.method(array_ops.ones([], dtype=dtypes.float64))
method_op = has_method.method.get_concrete_function()
self.assertEqual(
['y'],
[inp.op.name for inp in method_op.inputs])
self.assertEqual(
[b'y'],
[inp.op.get_attr('_user_specified_name') for inp in method_op.inputs])
method_op2 = has_method.method.get_concrete_function()
self.assertEqual(
['y'],
[inp.op.name for inp in method_op2.inputs])
self.assertEqual(
[b'y'],
[inp.op.get_attr('_user_specified_name') for inp in method_op2.inputs])
def testVariadic(self):
@polymorphic_function.function
def variadic_fn(x, *args, **kwargs):
return x + math_ops.add_n(list(args) + list(kwargs.values()))
# Call the function to make def_function happy
variadic_fn(array_ops.ones([]), array_ops.ones([]))
variadic_op = variadic_fn.get_concrete_function(
tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32),
tensor_spec.TensorSpec(shape=None, dtype=dtypes.float32, name='y'),
tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32),
tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32,
name='second_variadic'),
z=tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32),
zz=tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32, name='cust'))
self.assertEqual(
['x', 'y', 'args_1', 'second_variadic', 'z', 'cust'],
[inp.op.name for inp in variadic_op.inputs])
self.assertEqual(
[b'x', b'y', b'args_1', b'second_variadic', b'z', b'cust'],
[inp.op.get_attr('_user_specified_name') for inp in variadic_op.inputs])
def testVariadicInputSignature(self):
@polymorphic_function.function(
input_signature=(
tensor_spec.TensorSpec(shape=None, dtype=dtypes.float32),
tensor_spec.TensorSpec(shape=None, dtype=dtypes.float32, name='y'),
tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32),
tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32, name='z'),
))
def variadic_fn(x, *args):
return x + math_ops.add_n(list(args))
# Call the function to make def_function happy
variadic_fn(array_ops.ones([]), array_ops.ones([]),
array_ops.ones([]), array_ops.ones([]))
variadic_op = variadic_fn.get_concrete_function()
self.assertIn(b'variadic_fn', variadic_op.name)
self.assertEqual(
['x', 'y', 'args_1', 'z'],
[inp.op.name for inp in variadic_op.inputs])
self.assertEqual(
[b'x', b'y', b'args_1', b'z'],
[inp.op.get_attr('_user_specified_name')
for inp in variadic_op.inputs])
if __name__ == '__main__':
ops.enable_eager_execution(
config=config_pb2.ConfigProto(device_count={'CPU': 4}))
test.main()
tensorflow/python/eager/polymorphic_function/atomic_function.py
+696
View File
@@ -0,0 +1,696 @@
# 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.
# ==============================================================================
"""Implementation for AtomicFunction."""
import dataclasses
import traceback
import typing
from typing import Any, Dict, List, Optional, Sequence, Union
from tensorflow.core.framework import attr_value_pb2
from tensorflow.core.framework import function_pb2
from tensorflow.core.framework import graph_debug_info_pb2
from tensorflow.core.function.polymorphism import function_type as function_type_lib
from tensorflow.python.client import pywrap_tf_session
from tensorflow.python.eager import context
from tensorflow.python.eager import record
from tensorflow.python.eager.polymorphic_function import attributes as attributes_lib
from tensorflow.python.eager.polymorphic_function import function_type_utils
from tensorflow.python.framework import auto_control_deps_utils as acd
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import error_interpolation
from tensorflow.python.framework import errors
from tensorflow.python.framework import func_graph as func_graph_module
from tensorflow.python.framework import function_def_to_graph
from tensorflow.python.framework import ops
from tensorflow.python.ops import handle_data_util
from tensorflow.python.types import core
from tensorflow.python.util import compat
from tensorflow.python.util import function_utils
from tensorflow.python.util import tf_stack
# TODO(fmuham): Should be lowered to FunctionDef/FunctionRecord.
@dataclasses.dataclass(frozen=True)
class CallOptions:
"""Specifies additional configuration for an AtomicFunction call."""
# Used by ACD to identify the CollectiveManager this function is scoped in.
collective_manager_ids_used: List[int] = dataclasses.field(
default_factory=list
)
# Used by ACD to list Ops/Tensors/Callables that must be called in advance.
control_captures: List[Any] = dataclasses.field(default_factory=list)
# Determines what kind of partitioned call is used for this function.
is_stateful: bool = False
# Maps the (scope_id, name) in runtime to associated AtomicFunctions.
RUNTIME_FUNCTION_REFS = {}
class AtomicFunction(core.AtomicFunction):
"""A Python callable for functions in the TF Runtime.
Provides core functionality for tf.function including:
- automatic lifecycle management of runtime functions
- structured inputs (including captures) and structured outputs
- calls from both eager and graph mode
- dependency tracking of children functions
- runtime error interpolation to identify user code stack traces
- control dependencies (including automatic)
"""
__slots__ = [
"_name",
"_bound_context",
"_function_type",
"_children",
"_call_options",
"_cached_definition",
"_cached_graph",
"_generated_graph",
]
def __init__(
self,
name: Union[str, bytes],
bound_context: context.Context,
function_type: function_type_lib.FunctionType,
children: Optional[List["AtomicFunction"]] = None,
call_options: CallOptions = CallOptions(),
cached_graph: Optional[func_graph_module.FuncGraph] = None,
):
"""Construct a new AtomicFunction.
Args:
name: str/bytes name of the runtime function in the bound context.
bound_context: interface to the runtime for the AtomicFunction.
function_type: input/output contract for the AtomicFunction
children: list of AtomicFunctions that are needed to call this one.
call_options: extra configuration options for the call.
cached_graph: FuncGraph that this AtomicFunction was generated from (if
known). Otherwise it will lazily construct a new corresponding FuncGraph
if ever needed.
"""
self._name = compat.as_bytes(name)
self._bound_context = bound_context
self._function_type = function_type
self._children = children if children else []
self._call_options = call_options
self._cached_definition = None
self._cached_graph = cached_graph
self._generated_graph = None
ref_key = (self._bound_context.function_scope_id, self.name)
if ref_key not in RUNTIME_FUNCTION_REFS:
RUNTIME_FUNCTION_REFS[ref_key] = 1
else:
RUNTIME_FUNCTION_REFS[ref_key] += 1
@property
def name(self) -> bytes:
"""Name represented in UTF-8 encoded bytes."""
return self._name
@property
def function_type(self) -> function_type_lib.FunctionType:
"""Represents the input/output contract of this function."""
return self._function_type
@property
def children(self) -> List["AtomicFunction"]:
"""AtomicFunctions needed as dependencies for this one."""
return self._children
@property
def definition(self) -> function_pb2.FunctionDef:
"""Current FunctionDef in the Runtime."""
return self._bound_context.get_function_def(self.name)
@property
def attributes(self) -> Any:
"""Returns FunctionDef attributes in the Runtime."""
attrs = self.definition.attr
# Remove construction context since it is specific to runtime and this fn.
attrs.pop(attributes_lib.EAGER_RUNTIME_CONSTRUCTION_CONTEXT, None)
return attrs
@property
def graph_debug_info(self) -> graph_debug_info_pb2.GraphDebugInfo:
"""A GraphDebugInfo proto mapping nodes to corresponding stack traces."""
return self._bound_context.get_graph_debug_info(self.name)
@property
def call_options(self) -> CallOptions:
"""Call options declared for this AtomicFunction."""
return self._call_options
@property
def graph_call_attrs(self) -> Dict[str, Any]:
"""Returns a dictionary of attributes needed to add a call in graph."""
attrs = {
"is_stateful": self.call_options.is_stateful,
"tout": [
o.dtype.as_datatype_enum for o in self.function_type.flat_outputs
],
"xla_compile_attr": self.cached_definition.attr.get(
attributes_lib.XLA_COMPILE, None
),
}
attrs.update(self._bound_context.function_call_options.as_attrs())
return attrs
@property
def _c_func(self) -> Any:
"""Returns a scoped pybind object containing FunctionRecord in runtime."""
return self._bound_context.get_c_function(self.name)
# TODO(fmuham): Move caching to dependent code and remove method.
@property
def cached_definition(self) -> function_pb2.FunctionDef:
"""Cached FunctionDef (not guaranteed to be fresh)."""
if self._cached_definition is None:
self._cached_definition = self.definition
return self._cached_definition
@property
def graph(self) -> func_graph_module.FuncGraph:
"""Returns a FuncGraph corresponding to the AtomicFunction."""
if self._cached_graph:
return self._cached_graph
# Lazily generate the graph if one is not specified.
if not self._generated_graph:
self._generated_graph = to_func_graph(self)
return self._generated_graph
def call_with_captures(
self, args: Sequence[Any], kwargs: Dict[str, Any], captures: Sequence[Any]
) -> Any:
"""Calls with args, kwargs, captures and returns structured output."""
bound_parameters = self.function_type.bind(*args, **kwargs)
tensor_inputs = self.function_type.unpack_inputs(bound_parameters)
capture_inputs = self.function_type.unpack_captures(captures)
return self.call_preflattened(tensor_inputs + capture_inputs)
def call_preflattened(self, args: Sequence[core.Tensor]) -> Any:
"""Calls with flattened tensor inputs and returns the structured output."""
flat_outputs = self.call_flat(*args)
return self.function_type.pack_output(flat_outputs)
def call_flat(self, *args: core.Tensor) -> Sequence[core.Tensor]:
"""Calls with flat tensor inputs and returns flat tensor outputs.
Args:
*args: arguments to call this function with.
Returns:
The outputs of the function call.
Raises:
ValueError: if the number of arguments is incorrect.
FunctionAlreadyGarbageCollectedError: if the function is no longer
available to be called because it has been garbage collected.
"""
expected_len = len(self.cached_definition.signature.input_arg)
if len(args) != expected_len:
raise ValueError(
f"Signature specifies {expected_len} arguments, got: {len(args)}."
f" Expected inputs: {self.cached_definition.signature.input_arg}."
f" Received inputs: {args}."
f" Function Type: {self.function_type!r}"
)
with InterpolateRuntimeError(self):
with ops.control_dependencies(self._call_options.control_captures):
# The caller must use record_operation to record this operation in the
# eager case, so we enforce the same requirement for the non-eager
# case by explicitly pausing recording. We don't have a gradient
# registered for PartitionedCall, so recording this operation confuses
# forwardprop code (GradientTape manages to ignore it).
with record.stop_recording():
if self._bound_context.executing_eagerly():
outputs = self._bound_context.call_function(
self.name,
list(args),
len(self.function_type.flat_outputs),
)
else:
outputs = make_call_op_in_graph(
self,
list(args),
self._bound_context.function_call_options.as_attrs(),
)
for i, output_type in enumerate(self.function_type.flat_outputs):
handle_data = output_type.dtype._handle_data # pylint: disable=protected-access
if handle_data:
handle_data_util.set_handle_data(
outputs[i], handle_data.shape_inference
)
# TODO(fmuham): Use FunctionType cast here for all cases.
if not self._bound_context.executing_eagerly():
for i, output_type in enumerate(self.function_type.flat_outputs):
outputs[i].set_shape(output_type.shape)
return outputs
def __call__(self, *args, **kwargs) -> Any:
if self.function_type.captures:
raise ValueError(
"The FunctionType defines captured inputs. Use call_with_captures"
" instead."
)
return self.call_with_captures(args, kwargs, [])
def __del__(self):
if self._generated_graph:
func_graph_module.dismantle_func_graph(self._generated_graph)
if RUNTIME_FUNCTION_REFS is None:
return
key = (self._bound_context.function_scope_id, self.name)
RUNTIME_FUNCTION_REFS[key] -= 1
if RUNTIME_FUNCTION_REFS[key] < 0:
raise RuntimeError(
f"AtomicFunction Refcounting for {self.name} is invalid."
)
if RUNTIME_FUNCTION_REFS[key] == 0:
try:
self._bound_context.remove_function(self.name)
RUNTIME_FUNCTION_REFS.pop(key)
except TypeError:
# Suppress some exceptions, mainly for the case when we're running on
# module deletion. Things that can go wrong include the context module
# already being unloaded, self._handle._handle_data no longer being
# valid, and so on. Printing warnings in these cases is silly
# (exceptions raised from __del__ are printed as warnings to stderr).
pass # 'NoneType' object is not callable when the handle has been
# partially unloaded.
except AttributeError:
pass # 'NoneType' object has no attribute 'eager_mode' when context has
# been unloaded. Will catch other module unloads as well.
def __str__(self):
return f"<AtomicFunction> {compat.as_str(self.name)}{self.function_type}"
def __repr__(self):
return (
f"AtomicFunction(name={self.name},\n"
f"bound_context={self._bound_context},\n"
f"function_type={self.function_type!r},\n"
f"children={self._children!s},\n"
f"call_options={self._call_options},\n"
f"cached_graph={self._cached_graph})"
)
def _set_read_only_resource_inputs_attr(
op: ops.Operation, func_graph: func_graph_module.FuncGraph
):
"""Sets the list of resource inputs which are read-only.
This is used by AutomaticControlDependencies.
Args:
op: PartitionedCall Operation.
func_graph: FuncGraph.
"""
read_only_indices = acd.get_read_only_resource_input_indices_graph(func_graph)
ops.set_int_list_attr(
op, acd.READ_ONLY_RESOURCE_INPUTS_ATTR, read_only_indices
)
def partitioned_call_op(
name: str,
args: Sequence[core.Tensor],
is_stateful: bool,
tout: Sequence[Any],
config: Any = None,
executor_type: Optional[str] = None,
xla_compile_attr: Any = None,
) -> ops.Operation:
"""Generates a function call op respecting device annotations.
Args:
name: Name of the function to call.
args: The arguments of the function, including captured inputs.
is_stateful: If the function is stateful.
tout: a list containing the output dtypes enums
config: (Optional) A `tensorflow::ConfigProto` proto, serialized. If `None`,
all optimizations are disabled. Currently only handled for eager defined
functions.
executor_type: (Optional) A string for the name of the executor to be used
in the function call. If not set, or set to an empty string, the default
tensorflow executor will be used.
xla_compile_attr: (Optional) value of the XLA compilation attribute.
Returns:
Returns the operation.
"""
if config is None:
config = function_utils.get_disabled_rewriter_config()
if executor_type is None:
executor_type = ""
# The generated binding returns an empty list for functions that don't
# return any Tensors, hence the need to use `create_op` directly.
args = [ops.convert_to_tensor(x) for x in args]
tin_attr = attr_value_pb2.AttrValue(
list=attr_value_pb2.AttrValue.ListValue(
type=[x.dtype.as_datatype_enum for x in args]
)
)
tout_attr = attr_value_pb2.AttrValue(
list=attr_value_pb2.AttrValue.ListValue(type=tout)
)
func_attr = attr_value_pb2.AttrValue(
func=attr_value_pb2.NameAttrList(name=name)
)
executor_type_attr = attr_value_pb2.AttrValue(
s=compat.as_bytes(executor_type)
)
# When running in graph mode, the graph and function graphs are optimized
# (i.e. run through grappler) per the session options, so we can disable any
# eager-specific rewriting.
config_proto = attr_value_pb2.AttrValue(s=config)
op_name = "StatefulPartitionedCall" if is_stateful else "PartitionedCall"
# Propagate the attribute indicating the need to compile from function to the
# call itself.
op_attrs = {
"Tin": tin_attr,
"Tout": tout_attr,
"f": func_attr,
"config_proto": config_proto,
"executor_type": executor_type_attr,
}
if xla_compile_attr is not None:
op_attrs[attributes_lib.XLA_COMPILE] = xla_compile_attr
op = ops.get_default_graph().create_op(
op_name, args, tout, name=op_name, attrs=op_attrs
)
return op
def make_call_op_in_graph(
atomic: AtomicFunction,
tensor_inputs: Sequence[core.Tensor],
context_call_attrs: Dict[str, Any],
):
"""Adds an AtomicFunction to graph."""
graph = ops.get_default_graph()
graph._add_function_recursive(atomic) # pylint: disable=protected-access
op = partitioned_call_op( # pytype: disable=wrong-arg-types # always-use-property-annotation
name=atomic.name,
args=tensor_inputs,
is_stateful=atomic.call_options.is_stateful,
tout=[
o.dtype.as_datatype_enum for o in atomic.function_type.flat_outputs
],
config=context_call_attrs["config_proto"],
executor_type=context_call_attrs["executor_type"],
xla_compile_attr=atomic.cached_definition.attr.get(
attributes_lib.XLA_COMPILE, None
),
)
_set_read_only_resource_inputs_attr(op, atomic.graph)
ops.set_int_list_attr(
op,
acd.COLLECTIVE_MANAGER_IDS,
atomic._call_options.collective_manager_ids_used, # pylint: disable=protected-access
)
return op.outputs
def from_function_def(
function_def: function_pb2.FunctionDef,
function_type: function_type_lib.FunctionType,
) -> AtomicFunction:
"""Create a new AtomicFunction from FunctionDef + FunctionType."""
bound_context = context.context()
if bound_context.has_function(compat.as_bytes(function_def.signature.name)):
raise ValueError("Function already registered in context.")
bound_context.add_function_def(function_def)
return AtomicFunction(
function_def.signature.name, bound_context, function_type
)
def from_func_graph(
name: Union[str, bytes],
graph: func_graph_module.FuncGraph,
attrs: Dict[str, attr_value_pb2.AttrValue],
function_type: Optional[function_type_lib.FunctionType] = None,
overwrite: bool = False,
) -> AtomicFunction:
"""Initializes an AtomicFunction from FuncGraph.
Args:
name: str, the name for the created function.
graph: Graph, the graph containing the operations in the function
attrs: dict mapping names of attributes to their AttrValue values
function_type: known FunctionType to use, otherwise one is derived.
overwrite: overwrites function definition in the current context if needed
Returns:
An AtomicFunction instance.
"""
if attrs and attributes_lib.IMPLEMENTS in attrs:
# The alternative is to silently drop "implements" tag
# but it seems likely it would lead to hard to catch bugs.
# Another alternative is to make func_body to preserve the order
# of arguments if variables are present. Yet another option
# is to automatically replace variables as arguments to functions
# to v.read_value() whenever "implements" tag is present
# Anytime we annotate existing function we probably want to wrap
# it with safe read_value for backward compatibility.
has_resource_vars = any(
inp.dtype == dtypes.resource for inp in graph.inputs
)
captured_inputs = graph.external_captures + graph.deferred_external_captures
assert not any(
(has_resource_vars, captured_inputs)
), (
'Function {name} has "{attr}={value}" attribute and thus can not '
"depend on any tensors outside of its signature or modify variables. "
"\n\nNote: variables are always captured and cause function "
"re-tracing for every variable called.\n"
" inputs: {inputs}\n captures: {captured}\n\n"
"To pass a variable to such function use "
"use variable.read_value().".format(
name=graph.name,
attr=attributes_lib.IMPLEMENTS,
value=attrs[attributes_lib.IMPLEMENTS],
inputs=graph.inputs,
captured=captured_inputs,
)
)
input_ops = set(arg.op for arg in graph.inputs)
operations = [op for op in graph.get_operations() if op not in input_ops]
graph_output_names = graph._output_names # pylint: disable=protected-access
if graph_output_names is not None and all(
ops.tensor_id(t) in graph_output_names for t in graph.outputs
):
output_names = [
compat.as_bytes(graph_output_names[ops.tensor_id(t)])
for t in graph.outputs
]
if len(set(output_names)) != len(output_names):
# There are duplicate names for some reason, probably an invalid
# signature. Revert to auto-naming.
output_names = []
else:
output_names = []
with graph._c_graph.get() as c_graph: # pylint: disable=protected-access
fn = pywrap_tf_session.TF_GraphToFunction_wrapper(
c_graph,
compat.as_str(name),
False,
[o._c_op for o in operations], # pylint: disable=protected-access
[t._as_tf_output() for t in graph.inputs], # pylint: disable=protected-access
[t._as_tf_output() for t in graph.outputs], # pylint: disable=protected-access
output_names,
[o._c_op for o in graph.control_outputs], # pylint: disable=protected-access
[], # control_output_names
None,
compat.as_str(""),
)
attrs = attributes_lib.parse_func_attrs(attrs or {})
for attr_name, attr_value in attrs.items():
serialized = attr_value.SerializeToString()
pywrap_tf_session.TF_FunctionSetAttrValueProto(
fn, compat.as_str(attr_name), serialized
)
name = compat.as_bytes(name)
bound_context = context.context()
if overwrite and bound_context.has_function(name):
bound_context.remove_function(name)
bound_context.add_c_function(fn)
pywrap_tf_session.TF_DeleteFunction(fn)
call_options = CallOptions(
collective_manager_ids_used=getattr(
graph, "collective_manager_ids_used", []
),
control_captures=graph.function_captures.control,
is_stateful=any(op._is_stateful for op in operations), # pylint: disable=protected-access
)
if not function_type:
function_type = function_type_utils.derive_from_graph(graph)
return AtomicFunction(
name,
bound_context,
function_type,
list(graph._functions.values()), # pylint: disable=protected-access,
call_options,
cached_graph=graph,
)
def to_func_graph(atomic: AtomicFunction) -> func_graph_module.FuncGraph:
"""Generate a FuncGraph from an AtomicFunction."""
# pylint: disable=protected-access
input_signature, output_signature = function_type_lib.to_structured_signature(
atomic.function_type
)
with ops.Graph().as_default():
# Insert dependencies in the default graph so the new graph can pull them.
for f in atomic.children:
ops.get_default_graph()._add_function(f)
result = function_def_to_graph.function_def_to_graph(
atomic.definition,
structured_input_signature=input_signature,
structured_outputs=output_signature,
propagate_device_spec=True,
include_library_functions=False,
)
for f in atomic.children:
result._add_function(f)
# Set input shapes and handle data
for i, input_type in enumerate(atomic.function_type.flat_inputs):
handle_data = input_type.dtype._handle_data
if handle_data:
handle_data_util.set_handle_data(
result.inputs[i], handle_data.shape_inference
)
result.inputs[i].set_shape(input_type.shape)
# Set output shapes and handle data
for i, output_type in enumerate(atomic.function_type.flat_outputs):
handle_data = output_type.dtype._handle_data
if handle_data:
handle_data_util.set_handle_data(
result.outputs[i], handle_data.shape_inference
)
result.outputs[i].set_shape(output_type.shape)
result.collective_manager_ids_used = (
atomic.call_options.collective_manager_ids_used,
)
# pylint: enable=protected-access
return result
class InterpolateRuntimeError(object):
"""Context Manager that interpolates exceptions received by AtomicFunction."""
DENY_LIST_PHRASES = ["<embedded"]
def __init__(self, top_level_func):
self._func = top_level_func
def interpolate(self, message, node_names, graph_debug_info):
"""Uses the GraphDebugInfo to generate an error message."""
error_message = ["Graph execution error:", ""]
traces = tf_stack.LoadTracesFromDebugInfo(graph_debug_info)
for node_name in node_names:
error_message.append(
f"Detected at node {node_name} defined at (most recent call last):"
)
if node_name in traces:
stack_trace = traces[node_name]
for formatted_frame in traceback.format_list(stack_trace):
if not any(p in formatted_frame for p in self.DENY_LIST_PHRASES):
error_message.append(formatted_frame)
else:
error_message.append("<stack traces unavailable>")
error_message.append(message.strip())
return "\n".join(error_message)
def __enter__(self):
pass
def __exit__(self, typ, exc, tb):
if not exc or not isinstance(exc, errors.OpError):
return False
exc = typing.cast(errors.OpError, exc)
message = compat.as_text(exc.message)
parsed_message, func_tags, node_tags = error_interpolation.parse_message(
message
)
deepest_func = None
for func_tag in func_tags:
if func_tag.name == compat.as_str(self._func.name):
deepest_func = self._func
elif deepest_func:
next_func = None
for child_func in deepest_func.children:
if func_tag.name == compat.as_str(child_func.name):
next_func = child_func
break
if next_func is not None and isinstance(next_func, AtomicFunction):
deepest_func = next_func
if deepest_func:
exc._message = self.interpolate(
parsed_message,
[t.name for t in node_tags],
deepest_func.graph_debug_info,
)
return False
tensorflow/python/eager/polymorphic_function/atomic_function_test.py
+180
View File
@@ -0,0 +1,180 @@
# 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.eager.polymorphic_function import atomic_function
from tensorflow.python.eager.polymorphic_function import polymorphic_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.platform import test
def get_function_def_and_type(foo, inputs):
"""Traces `foo` generate the FunctionDef and FunctionType."""
concrete = polymorphic_function.function(foo).get_concrete_function(*inputs)
atomic = concrete._inference_function
return atomic.definition, atomic.function_type
class AtomicFunctionTest(test.TestCase):
def test_call_eager(self):
definition, func_type = get_function_def_and_type(
lambda x, y: x + y, (constant_op.constant(1), constant_op.constant(2))
)
atomic = atomic_function.from_function_def(definition, func_type)
self.assertRegex(
str(atomic),
r"<AtomicFunction> .*(x: TensorSpec.*, y: TensorSpec.*) ->"
r" TensorSpec.*",
)
self.assertRegex(
repr(atomic).replace("\n", " "),
r"AtomicFunction.*name.*bound_context.*function_type.*"
r"children.*call_options.*cached_graph.*",
)
self.assertEqual(
atomic.call_flat(constant_op.constant(3), constant_op.constant(4))[
0
].numpy(),
7,
)
def test_call_graph(self):
definition, func_type = get_function_def_and_type(
lambda x, y: x + y, (constant_op.constant(1), constant_op.constant(2))
)
atomic = atomic_function.from_function_def(definition, func_type)
@polymorphic_function.function
def foo(a, b):
return atomic.call_flat(a, b)[0]
self.assertEqual(
foo(constant_op.constant(3), constant_op.constant(4)).numpy(),
7,
)
def test_variable_input_eager(self):
definition, func_type = get_function_def_and_type(
lambda x, y: x + y,
(resource_variable_ops.ResourceVariable(1), constant_op.constant(2)),
)
atomic = atomic_function.from_function_def(definition, func_type)
self.assertEqual(
atomic.call_flat(
resource_variable_ops.ResourceVariable(3)._handle,
constant_op.constant(4),
)[0].numpy(),
7,
)
def test_variable_input_graph(self):
definition, func_type = get_function_def_and_type(
lambda x, y: x + y,
(resource_variable_ops.ResourceVariable(1), constant_op.constant(2)),
)
atomic = atomic_function.from_function_def(definition, func_type)
@polymorphic_function.function
def foo(a, b):
return atomic.call_flat(a, b)[0]
self.assertEqual(
foo(
resource_variable_ops.ResourceVariable(3)._handle,
constant_op.constant(4),
).numpy(),
7,
)
def test_call_with_captures(self):
my_capture = constant_op.constant(2)
@polymorphic_function.function
def foo(x):
my_dict = {}
my_dict["my_tensor"] = x["my_tensor"]
my_dict["my_resource"] = x["my_variable"].handle
my_dict["my_capture"] = my_capture
my_dict["my_ints"] = x["my_ints"]
return my_dict
structured_inputs = {
"my_tensor": constant_op.constant(1),
"my_variable": resource_variable_ops.ResourceVariable(1),
"my_ints": [1, 2, 3],
}
function_def, function_type = get_function_def_and_type(
foo, (structured_inputs,)
)
atomic = atomic_function.from_function_def(function_def, function_type)
with self.assertRaisesRegex(ValueError, "Use call_with_captures instead."):
atomic(structured_inputs)
result = atomic.call_with_captures((structured_inputs,), {}, [my_capture])
self.assertEqual(
result["my_tensor"].numpy(), structured_inputs["my_tensor"].numpy()
)
self.assertEqual(result["my_resource"].dtype, dtypes.resource)
self.assertEqual(result["my_capture"].numpy(), my_capture.numpy())
self.assertEqual(result["my_ints"][0].numpy(), 1)
self.assertEqual(result["my_ints"][1].numpy(), 2)
self.assertEqual(result["my_ints"][2].numpy(), 3)
def test_call(self):
@polymorphic_function.function
def foo(x):
my_dict = {}
my_dict["my_tensor"] = x["my_tensor"]
my_dict["my_resource"] = x["my_variable"].handle
my_dict["my_ints"] = x["my_ints"]
return my_dict
structured_inputs = {
"my_tensor": constant_op.constant(1),
"my_variable": resource_variable_ops.ResourceVariable(1),
"my_ints": [1, 2, 3],
}
function_def, function_type = get_function_def_and_type(
foo, (structured_inputs,)
)
atomic = atomic_function.from_function_def(function_def, function_type)
result = atomic(structured_inputs)
self.assertEqual(
result["my_tensor"].numpy(), structured_inputs["my_tensor"].numpy()
)
self.assertEqual(result["my_resource"].dtype, dtypes.resource)
self.assertEqual(result["my_ints"][0].numpy(), 1)
self.assertEqual(result["my_ints"][1].numpy(), 2)
self.assertEqual(result["my_ints"][2].numpy(), 3)
if __name__ == "__main__":
ops.enable_eager_execution()
test.main()
tensorflow/python/eager/polymorphic_function/attributes.py
+184
View File
@@ -0,0 +1,184 @@
# 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.
# ==============================================================================
"""This file lists FunctionDef attributes and corresponding allowlists."""
from tensorflow.core.framework import attr_value_pb2
from tensorflow.python.util import compat
# IMPORTANT: The usage of all the attributes below should be considered tech
# debt and new additions to this list are discouraged.
#
# Historically, attributes have been used as means to pipe extra information
# down to runtime that is not related to the actual function definition itself.
#
# This information is better layered independently and future work is encouraged
# to pursue that direction instead.
API_IMPLEMENTS = "api_implements"
API_PREFERRED_DEVICE = "api_preferred_device"
BACKWARD_FUNCTION = "backward_function_name"
DISABLE_ACD = "_disable_acd"
DISABLE_CALL_SHAPE_INFERENCE = "_disable_call_shape_inference"
DISABLE_SUMMARIES_AT_RUNTIME = "disable_summaries_at_runtime"
EAGER_RUNTIME_CONSTRUCTION_CONTEXT = "_construction_context"
FORWARD_FUNCTION = "forward_function_name"
GO_BACKWARDS = "go_backwards"
IMPLEMENTS = "_implements"
INPUT_SHAPES = "_input_shapes"
INTS_ON_DEVICE = "experimental_ints_on_device"
NO_INLINE = "_noinline"
ORIGINAL_FUNCTION_NAME = "_original_func_name"
OUTPUTS_ON_OP_DEVICE = "_OutputsOnOpDevice"
QUANTIZED_COMPOSITE_FUNCTION = "tf_quant.composite_function"
QUANTIZED_OPS = "tf_quant.quantized_ops"
RUNS_AT_MOST_ONCE = "function_runs_at_most_once"
RUNTIME_CONSTANT_OPTIMIZATION = "runtime_constant_optimization"
SHARED_RENDEZVOUS = "shared_rendezvous"
TF_DATA_FUNCTION = "_tf_data_function"
TFTRT_ALLOW_BUILD_AT_RUNTIME = "_tftrt_allow_build_at_runtime"
TFTRT_CONVERT_FUNCTION = "_tftrt_convert_function"
TFTRT_IS_DYN_OP = "_tftrt_is_dyn_op"
TFTRT_LOGGER = "_tftrt_trt_logger_name"
TFTRT_MAX_BATCH_SIZE = "_tftrt_max_batch_size"
TFTRT_MAX_CACHED_ENGINES = "_tftrt_max_cached_engines"
TFTRT_MAX_WORKSPACE_SIZE = "_tftrt_max_workspace_size_bytes"
TFTRT_MIN_SEGMENT_SIZE = "_tftrt_minimum_segment_size"
TFTRT_PRECISION_MODE = "_tftrt_precision_mode"
TFTRT_PROFILE_STRATEGY = "_tftrt_profile_strategy"
TFTRT_USE_CALIBRATION = "_tftrt_use_calibration"
TFTRT_USE_IMPLICIT_BATCH = "_tftrt_use_implicit_batch"
TIME_MAJOR = "time_major"
XLA_COMPILE = "_XlaMustCompile"
XLA_COMPILE_OPTIONAL = "_XlaCompile"
XLA_SCOPE = "_XlaScope"
XLA_SEPARATE_COMPILED_GRADIENTS = "_XlaSeparateCompiledGradients"
POLYMORPHIC_FUNCTION_ALLOWLIST = frozenset({
API_IMPLEMENTS,
API_PREFERRED_DEVICE,
DISABLE_ACD,
DISABLE_SUMMARIES_AT_RUNTIME,
GO_BACKWARDS,
IMPLEMENTS,
INTS_ON_DEVICE,
NO_INLINE,
RUNS_AT_MOST_ONCE,
RUNTIME_CONSTANT_OPTIMIZATION,
TF_DATA_FUNCTION,
TIME_MAJOR,
OUTPUTS_ON_OP_DEVICE,
})
TRACING_COMPILATION_ALLOWLIST = frozenset().union(
POLYMORPHIC_FUNCTION_ALLOWLIST,
{
SHARED_RENDEZVOUS,
XLA_COMPILE,
},
)
MONOMORPHIC_FUNCTION_ALLOWLIST = frozenset().union(
TRACING_COMPILATION_ALLOWLIST,
{
BACKWARD_FUNCTION,
DISABLE_CALL_SHAPE_INFERENCE,
EAGER_RUNTIME_CONSTRUCTION_CONTEXT,
FORWARD_FUNCTION,
INPUT_SHAPES,
ORIGINAL_FUNCTION_NAME,
QUANTIZED_COMPOSITE_FUNCTION,
QUANTIZED_OPS,
TFTRT_ALLOW_BUILD_AT_RUNTIME,
TFTRT_CONVERT_FUNCTION,
TFTRT_IS_DYN_OP,
TFTRT_LOGGER,
TFTRT_MAX_BATCH_SIZE,
TFTRT_MAX_CACHED_ENGINES,
TFTRT_MAX_WORKSPACE_SIZE,
TFTRT_MIN_SEGMENT_SIZE,
TFTRT_PRECISION_MODE,
TFTRT_PROFILE_STRATEGY,
TFTRT_USE_CALIBRATION,
TFTRT_USE_IMPLICIT_BATCH,
XLA_COMPILE_OPTIONAL,
XLA_SCOPE,
XLA_SEPARATE_COMPILED_GRADIENTS,
},
)
def _parse_func_attr_value(key, value):
"""Converts a python object to an attr_value_pb2.AttrValue object."""
if isinstance(value, attr_value_pb2.AttrValue):
return value
# bool type check has to happen before int since bool is a subclass of int.
elif isinstance(value, bool):
return attr_value_pb2.AttrValue(b=value)
elif isinstance(value, int):
return attr_value_pb2.AttrValue(i=value)
elif isinstance(value, float):
return attr_value_pb2.AttrValue(f=value)
elif isinstance(value, (str, bytes)):
return attr_value_pb2.AttrValue(s=compat.as_bytes(value))
elif isinstance(value, list):
list_value = attr_value_pb2.AttrValue.ListValue()
for v in value:
if isinstance(v, bool):
list_value.b.append(v)
elif isinstance(v, int):
list_value.i.append(v)
elif isinstance(v, float):
list_value.f.append(v)
elif isinstance(v, (str, bytes)):
list_value.s.append(compat.as_bytes(v))
else:
raise ValueError(
f"Attributes for {key} must be bool, int, float, or string. "
f"Got {type(v)}."
)
return attr_value_pb2.AttrValue(list=list_value)
else:
raise ValueError(
f"Attribute {key} must be bool, int, float, string, list, or "
f"AttrValue. Got {type(value)}."
)
def parse_func_attrs(attributes, allowlist=None):
"""Convert the keyword arguments into function_def attributes.
Currently only support primitive types: bool, int, float and string.
Args:
attributes: the dictionary of attributes.
allowlist: set of attribute names allowed.
Returns:
A dict of attributes where the key is the name of attribute and the value
is the AttrValue proto.
Raises:
ValueError: If the kwargs contains unallowlisted name or unsupported value
types.
"""
if not allowlist:
allowlist = MONOMORPHIC_FUNCTION_ALLOWLIST
attrs = {}
for key, value in attributes.items():
if key not in allowlist:
raise ValueError(
f"Allowlist does not support `{key}` as an attribute.")
attrs[key] = _parse_func_attr_value(key, value)
return attrs
tensorflow/python/eager/polymorphic_function/autograph_util.py
+59
View File
@@ -0,0 +1,59 @@
# 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.
# ==============================================================================
# pylint: disable=unidiomatic-typecheck
"""Autograph utility functions for polymorphic_function."""
from tensorflow.python.autograph.core import converter
from tensorflow.python.autograph.impl import api
from tensorflow.python.util import tf_decorator
def py_func_from_autograph(
python_func,
autograph_options=None,
):
"""Compile a python function using autograph, for use with FuncGraph.
Args:
python_func: the Python function to compile.
autograph_options: additional knobs to control when `autograph=True`.
See https://www.tensorflow.org/guide/autograph for more information.
Returns:
python_func, converted using autograph.
"""
_, original_func = tf_decorator.unwrap(python_func)
def autograph_handler(*args, **kwargs):
"""Calls a converted version of original_func."""
try:
return api.converted_call(
original_func,
args,
kwargs,
options=converter.ConversionOptions(
recursive=True,
optional_features=autograph_options,
user_requested=True,
))
except Exception as e: # pylint:disable=broad-except
if hasattr(e, "ag_error_metadata"):
raise e.ag_error_metadata.to_exception(e)
else:
raise
# Wrapping around a decorator allows checks like tf_inspect.getargspec
# to be accurate.
converted_func = tf_decorator.make_decorator(original_func, autograph_handler)
return tf_decorator.rewrap(python_func, original_func, converted_func)
tensorflow/python/eager/polymorphic_function/collection_test.py
+69
View File
@@ -0,0 +1,69 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
from tensorflow.core.protobuf import config_pb2
from tensorflow.python.eager.polymorphic_function import polymorphic_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
class FunctionCollectionTest(test.TestCase):
def testCollectionValueAccess(self):
"""Read values from graph collections inside of defun."""
with ops.Graph().as_default() as g:
with self.session(graph=g):
x = 2
y = 5
ops.add_to_collection('x', x)
ops.add_to_collection('y', y)
@polymorphic_function.function
def fn():
x_const = constant_op.constant(ops.get_collection('x')[0])
y_const = constant_op.constant(ops.get_collection('y')[0])
z = math_ops.add(x_const, y_const)
ops.add_to_collection('z', 7)
return z
self.assertEqual(7, int(self.evaluate(fn())))
self.assertEqual(ops.get_collection('x'), [2])
self.assertEqual(ops.get_collection('y'), [5])
self.assertEqual(ops.get_collection('z'), [])
def testCollectionVariableValueAccess(self):
"""Read variable value from graph collections inside of defun."""
with ops.Graph().as_default() as g:
with self.session(graph=g):
v = resource_variable_ops.ResourceVariable(1.0)
@polymorphic_function.function
def f():
return v.read_value()
self.evaluate(variables.global_variables_initializer())
self.assertEqual(1.0, float(self.evaluate(f())))
self.assertLen(ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES), 1)
if __name__ == '__main__':
ops.enable_eager_execution(
config=config_pb2.ConfigProto(device_count={'CPU': 4}))
test.main()
tensorflow/python/eager/polymorphic_function/compiler_ir.py
+130
View File
@@ -0,0 +1,130 @@
# 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.
# ==============================================================================
"""Implementation for defining get_compiler_ir."""
from typing import List, Optional
import warnings
from tensorflow.core.function import trace_type
from tensorflow.python.eager import context
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import tensor_spec
from tensorflow.python.ops import random_ops
from tensorflow.python.util import nest
def maybe_get_device_name(device_name):
# TODO(cheshire): This is a hack to get the current "preferred" device,
# there is no current API to get it otherwise.
if device_name is None:
device_name = random_ops.random_normal([]).device
return device_name
# TODO(fmuham): Use trace_type._flatten here instead when available
def make_handledata_tensor_specs(resource_vars):
"""Convert tf.Variable list to its corresponding TensorSpec list."""
if not all(x.dtype is dtypes.resource for x in resource_vars):
raise RuntimeError("Resource_vars must be tf.resource list.")
inner_context = trace_type.InternalTracingContext()
trace_type_inputs = trace_type.from_value(
tuple(resource_vars), inner_context
).components
def to_resource_spec(traced_input):
try:
handle_data = traced_input.dtype._handle_data.shape_inference # pylint: disable=protected-access
shape_and_type = handle_data.shape_and_type[0]
spec = tensor_spec.TensorSpec(
shape=shape_and_type.shape, dtype=shape_and_type.dtype
)
return spec
except Exception as e:
raise ValueError(
"Fail to convert tf.Variable list to TensorSpec list. The error"
" is: %s" % e
) from e
return [to_resource_spec(trace_type) for trace_type in trace_type_inputs]
def from_concrete_function(
concrete_fn,
specialized_flat_specs: Optional[List[tensor_spec.TensorSpec]] = None,
):
"""Generate the Compiler Ir from tf concrete function with TensorSpec.
Args:
concrete_fn: returned by using get_concrete_function.
specialized_flat_specs: specialized flat tf.TensorSpecs for function args.
Returns:
Function callable that generate the HLO text.
Raises:
ValueError: if concrete_fn is not "compilable" without concrete
inputs.
"""
context.ensure_initialized()
fn_name = concrete_fn.name
filtered_flat_specs = specialized_flat_specs or list(
nest.flatten(concrete_fn.structured_input_signature)
)
if not all(s.shape.is_fully_defined() for s in filtered_flat_specs):
raise ValueError(
f"Only support static input shape but got inputs = {concrete_fn.inputs}"
)
def compiler_ir_generator(stage="hlo", device_name=None, platform_name=None):
"""Gets the compiler IR bytes.
Args:
stage: The exported stage for the given function.
device_name: The name of the device with the form as
"/job:localhost/replica:0/task:0/device:CPU:0", "/device:TPU:0" etc.
When this is used, actual device is needed for getting the compiler IR.
platform_name: The name of the platform, e.g. "TPU". See the comment in
`get_compiler_ir` in `context.py`.
Returns:
The compiler IR bytes.
"""
if device_name is not None:
if platform_name is not None:
raise ValueError(
"device_name and platform_name cannot be provided at the same time."
)
warnings.warn("device_name is being deprecated. Use platform_name.")
device_name = maybe_get_device_name(device_name)
res_bytes = context.context().get_compiler_ir(
device_name=device_name,
platform_name=platform_name,
function_name=fn_name,
flat_args=filtered_flat_specs,
captured_inputs=concrete_fn.captured_inputs,
stage=stage,
)
if stage in (
# Ordered by IrExportStage enum order
"stablehlo_serialized",
"hlo_serialized",
"optimized_hlo_serialized",
"optimized_hlo_proto_serialized",
):
return res_bytes
else:
return res_bytes.decode("utf-8")
return compiler_ir_generator
tensorflow/python/eager/polymorphic_function/compiler_ir_test.py
+241
View File
@@ -0,0 +1,241 @@
# 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.compiler.tests import xla_test
from tensorflow.python.eager.polymorphic_function import compiler_ir
from tensorflow.python.eager.polymorphic_function import polymorphic_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
from tensorflow.python.util import nest
class CompilerIrTest(xla_test.XLATestCase):
def _compareTwoMethodsCompilerIROutput(self, f, args, kwargs):
flat_args = list(args) + list(kwargs.values())
if not all([isinstance(x, tensor.Tensor) for x in flat_args]):
self.skipTest('It only support args and kwargs are all tf.Tensor types.')
args_spec = nest.map_structure(tensor.TensorSpec.from_tensor, args)
kwargs_spec = nest.map_structure(tensor.TensorSpec.from_tensor, kwargs)
hlo_1 = f.experimental_get_compiler_ir(*args, **kwargs)()
hlo_2 = f.experimental_get_compiler_ir(*args_spec, **kwargs_spec)()
if hlo_1 != hlo_2:
self.fail(
'The tensor_spec way experimental_get_compiler_ir give diff result'
' to normal experimental_get_compiler_ir.'
f' \nhlo(concrete_input):\n{hlo_1}\nhlo(tensor_spec):\n{hlo_2}\n'
)
# Check that StableHLO conversion succeeds
hlo_3 = f.experimental_get_compiler_ir(*args, **kwargs)(stage='stablehlo')
self.assertIn('stablehlo', hlo_3)
# Check that StableHLO bytecode conversion succeeds.
# MLIR bytecode files all begin with magic `MLiR` byte, check for byte.
hlo_4 = f.experimental_get_compiler_ir(*args, **kwargs)(
stage='stablehlo_serialized'
)
self.assertIn(b'ML\xefR', hlo_4)
def test_zero_input(self):
with ops.device('device:{}:0'.format(self.device)):
@polymorphic_function.function(jit_compile=True, autograph=False)
def fun_tf():
return array_ops.zeros((10), dtype=dtypes.int32)
self._compareTwoMethodsCompilerIROutput(fun_tf, [], {})
def test_constant_slice(self):
with ops.device('device:{}:0'.format(self.device)):
# Constant slice. This is the common case.
x = array_ops.zeros((10,), dtype=dtypes.int32)
@polymorphic_function.function(jit_compile=True, autograph=False)
def fun_tf(x):
begin = 0
return x[begin:5]
self._compareTwoMethodsCompilerIROutput(fun_tf, [x], {})
def test_compile_time_constant(self):
with ops.device('device:{}:0'.format(self.device)):
# Non-constant slice, but compile-time constant depending only on shapes.
x = array_ops.zeros((10,), dtype=dtypes.int32)
@polymorphic_function.function(jit_compile=True, autograph=False)
def fun_tf(x):
# begin is a compile-time constant, even if x is not
begin = array_ops.shape_v2(x)[0] - 2
return x[begin:]
self._compareTwoMethodsCompilerIROutput(fun_tf, [x], {})
def test_capture_constant(self):
with ops.device('device:{}:0'.format(self.device)):
# Capture a constant
outer_ct = [3.0]
x = ops.convert_to_tensor([2.0, 3.0, 4.0], dtype=dtypes.float32)
@polymorphic_function.function(jit_compile=True, autograph=False)
def fun_tf(x):
return x * gen_array_ops.broadcast_to(outer_ct, x.shape) + 1.0
self._compareTwoMethodsCompilerIROutput(fun_tf, [x], {})
def test_unsupported_dynamic_input(self):
with ops.device('device:{}:0'.format(self.device)):
@polymorphic_function.function(jit_compile=True)
def f(x):
return x
with self.assertRaisesRegex(
ValueError, 'Only support static input shape but got'
):
args_spec = [tensor.TensorSpec((None), dtype=dtypes.float32)]
concrete_fn = f.get_concrete_function(*args_spec)
_ = compiler_ir.from_concrete_function(concrete_fn)(stage='hlo')
def test_unsupported_shape_depend_input(self):
with ops.device('device:{}:0'.format(self.device)):
# Those cases output shapes are dynamic.
@polymorphic_function.function(jit_compile=True)
def f2(x):
return x[x[0] : 0]
args = [ops.convert_to_tensor([1, 2, 3, 4])]
args_spec = nest.map_structure(tensor.TensorSpec.from_tensor, args)
concrete_fn = f2.get_concrete_function(*args_spec)
_ = compiler_ir.from_concrete_function(concrete_fn)(stage='hlo')
def test_make_handledata_tensor_specs(self):
with ops.device('device:{}:0'.format(self.device)):
v1 = variables.Variable([0.1, 0.1])
v3 = variables.Variable([1], dtype=dtypes.int32)
@polymorphic_function.function(jit_compile=True)
def f4(a, b):
return (a + b) * v1 - math_ops.cast(v3, dtypes.float32)
a = constant_op.constant([1.1, 1.1])
b = constant_op.constant([2.2, 2.2])
kwargs = {'b': a, 'a': b}
kwargs_spec = nest.map_structure(
tensor.TensorSpec.from_tensor, kwargs
)
concrete_fn = f4.get_concrete_function(**kwargs_spec)
captured_inputs = concrete_fn.captured_inputs
captured_spec = compiler_ir.make_handledata_tensor_specs(captured_inputs)
self.assertEqual(len(captured_spec), 2)
self.assertEqual(
captured_spec[0], tensor.TensorSpec((2), dtype=dtypes.float32)
)
self.assertEqual(
captured_spec[1], tensor.TensorSpec((1), dtype=dtypes.int32)
)
def test_capture_variable_1(self):
if 'gpu' in self.device.lower():
self.skipTest('Skip test on GPU')
with ops.device('device:{}:0'.format(self.device)):
v1 = variables.Variable([0.1, 0.1])
v3 = variables.Variable([1], dtype=dtypes.int32)
@polymorphic_function.function(jit_compile=True)
def f4(a, b):
return (a + b) * v1 - math_ops.cast(v3, dtypes.float32)
a = constant_op.constant([1.1, 1.1])
b = constant_op.constant([2.2, 2.2])
kwargs = {'b': a, 'a': b}
self._compareTwoMethodsCompilerIROutput(f4, [], kwargs)
def test_capture_variable_2(self):
if 'gpu' in self.device.lower():
self.skipTest('Skip test on GPU')
with ops.device('device:{}:0'.format(self.device)):
v2 = variables.Variable(2.0, dtype=dtypes.float32)
v3 = variables.Variable(3.0, dtype=dtypes.float32)
@polymorphic_function.function(jit_compile=True)
def fun_tf(x):
# Defining tf.constants inside func_body is okay.
t4 = constant_op.constant(4.0, dtype=dtypes.float32)
t5 = constant_op.constant(5.0, dtype=dtypes.float32)
return (x * v3 + t4 + v2) * v3 + t5
x = constant_op.constant(2.0, dtype=dtypes.float32)
self._compareTwoMethodsCompilerIROutput(fun_tf, [x], {})
def test_capture_constants(self):
if 'gpu' in self.device.lower():
self.skipTest('Skip test on GPU')
with ops.device('device:{}:0'.format(self.device)):
v2 = variables.Variable(2.0, dtype=dtypes.float32)
v3 = variables.Variable(3.0, dtype=dtypes.float32)
t4 = constant_op.constant([4.0, 5.0], dtype=dtypes.float32)
t5 = constant_op.constant([5.0, 6.0], dtype=dtypes.float32)
@polymorphic_function.function(jit_compile=True)
def fun_tf(x):
return (x * v3 + t4 + v2) * v3 + t5
x = constant_op.constant([2.0, 3.0], dtype=dtypes.float32)
self._compareTwoMethodsCompilerIROutput(fun_tf, [x], {})
def test_from_concrete_function_with_args(self):
with ops.device('device:{}:0'.format(self.device)):
v2 = variables.Variable(2.0, dtype=dtypes.float32)
v3 = variables.Variable(3.0, dtype=dtypes.float32)
# Capturing tf.constants outside func_body is not okay.
t4 = constant_op.constant(4.0, dtype=dtypes.float32)
t5 = constant_op.constant(5.0, dtype=dtypes.float32)
@polymorphic_function.function(jit_compile=True)
def fun_tf(x):
return (x * v3 + t4 + v2) * v3 + t5
concrete_fn = fun_tf.get_concrete_function(
tensor.TensorSpec((None,), dtype=dtypes.float32)
)
x = tensor.TensorSpec((10,), dtype=dtypes.float32)
hlo_1 = compiler_ir.from_concrete_function(concrete_fn, [x])(stage='hlo')
self.assertIn('f32[10]', hlo_1)
x = tensor.TensorSpec((20,), dtype=dtypes.float32)
hlo_2 = compiler_ir.from_concrete_function(concrete_fn, [x])(stage='hlo')
self.assertIn('f32[20]', hlo_2)
if __name__ == '__main__':
ops.enable_eager_execution()
test.main()
tensorflow/python/eager/polymorphic_function/composite_tensor_utils.py
+39
View File
@@ -0,0 +1,39 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Utility to manipulate CompositeTensors in tf.function."""
from tensorflow.python.framework import composite_tensor
from tensorflow.python.util import _pywrap_utils
from tensorflow.python.util import nest
# TODO(b/240337581, b/240337099): Remove this function when we de-alias
# dt_resource tensors or tf.nest support is_leaf.
def flatten_with_variables(inputs):
"""Flattens `inputs` but don't expand `ResourceVariable`s."""
# We assume that any CompositeTensors have already converted their components
# from numpy arrays to Tensors, so we don't need to expand composites here for
# the numpy array conversion. Instead, we do so because the flattened inputs
# are eventually passed to ConcreteFunction()._call_flat, which requires
# expanded composites.
flat_inputs = []
for value in nest.flatten(inputs):
if (isinstance(value, composite_tensor.CompositeTensor) and
not _pywrap_utils.IsResourceVariable(value)):
components = value._type_spec._to_components(value) # pylint: disable=protected-access
flat_inputs.extend(flatten_with_variables(components))
else:
flat_inputs.append(value)
return flat_inputs
tensorflow/python/eager/polymorphic_function/concrete_function.py
+1771
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/eager/polymorphic_function/concrete_function_test.py
+157
View File
@@ -0,0 +1,157 @@
# 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 absl.testing import parameterized
from tensorflow.core.framework import attr_value_pb2
from tensorflow.python.eager.polymorphic_function import atomic_function
from tensorflow.python.eager.polymorphic_function import concrete_function as cf
from tensorflow.python.eager.polymorphic_function import polymorphic_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import func_graph as func_graph_module
from tensorflow.python.framework import ops
from tensorflow.python.ops import variables
from tensorflow.python.ops.ragged import ragged_tensor
from tensorflow.python.platform import test
from tensorflow.python.util import compat
class ConcreteFunctionTest(test.TestCase, parameterized.TestCase):
def concrete_function_with_attrs(self, attrs):
func_graph = func_graph_module.FuncGraph("f")
return cf.ConcreteFunction.from_func_graph(func_graph, None, attrs=attrs)
@parameterized.parameters(
({"api_implements": True}, attr_value_pb2.AttrValue(b=True)),
({"api_implements": 1}, attr_value_pb2.AttrValue(i=1)),
({"api_implements": 1.0}, attr_value_pb2.AttrValue(f=1.0)),
(
{"api_implements": "test"},
attr_value_pb2.AttrValue(s=compat.as_bytes("test")),
),
)
def test_parses_func_attr_scalar_values(self, attrs, expected):
self.assertEqual(
self.concrete_function_with_attrs(attrs=attrs).function_def.attr[
"api_implements"
],
expected,
)
def test_parses_func_attr_list_values(self):
self.assertProtoEquals(
r"""
list {
s: 'test'
b: True
i: 1
f: 1.0
}
""",
self.concrete_function_with_attrs(
attrs={"api_implements": ["test", True, 1, 1.0]}
).function_def.attr["api_implements"],
)
def test_raises_value_error_for_invalid_attr(self):
with self.assertRaisesRegex(ValueError, "Attribute api_implements must be"):
self.concrete_function_with_attrs(attrs={"api_implements": None})
def test_generate_from_atomic(self):
@polymorphic_function.function
def add_dicts(dict_a, dict_b):
result = {}
for key in dict_a.keys():
result[key] = dict_a[key] + dict_b[key]
return result
dict_a = {
"tensor": constant_op.constant(1),
"variable": variables.Variable(2),
"ragged_tensor": ragged_tensor.RaggedTensor.from_row_splits(
values=[3, 1, 4, 1, 5, 9, 2, 6], row_splits=[0, 4, 4, 7, 8, 8]
),
"python_int": 4,
}
dict_b = {
"tensor": constant_op.constant(2),
"variable": variables.Variable(5),
"ragged_tensor": ragged_tensor.RaggedTensor.from_row_splits(
values=[4, 2, 4, 1, 6, 9, 3, 6], row_splits=[0, 4, 4, 7, 8, 8]
),
"python_int": 5,
}
original_concrete_fn = add_dicts.get_concrete_function(dict_a, dict_b)
# Get the atomic function and delete everything else.
atomic_fn = original_concrete_fn._inference_function
del add_dicts
del original_concrete_fn
# Regenerate the ConcreteFunction.
concrete_fn = cf.ConcreteFunction(atomic_fn)
result = concrete_fn(dict_a, dict_b)
# Call and check results.
self.assertEqual(result["tensor"].numpy(), 3)
self.assertEqual(result["variable"].numpy(), 7)
self.assertEqual(
result["ragged_tensor"].flat_values.numpy().tolist(),
[7, 3, 8, 2, 11, 18, 5, 12],
)
self.assertEqual(result["python_int"].numpy(), 9)
def test_generate_from_def(self):
@polymorphic_function.function
def add_dicts(dict_a, dict_b):
result = {}
for key in dict_a.keys():
result[key] = dict_a[key] + dict_b[key]
return result
dict_a = {
"tensor": constant_op.constant(1),
"variable": variables.Variable(2),
"python_int": 4,
}
dict_b = {
"tensor": constant_op.constant(2),
"variable": variables.Variable(5),
"python_int": 5,
}
original_concrete_fn = add_dicts.get_concrete_function(dict_a, dict_b)
# Get FunctionDef + FunctionType and delete everything else.
function_def = original_concrete_fn.function_def
function_type = original_concrete_fn.function_type
del add_dicts
del original_concrete_fn
# Regenerate the ConcreteFunction.
atomic_fn = atomic_function.from_function_def(function_def, function_type)
concrete_fn = cf.ConcreteFunction(atomic_fn)
result = concrete_fn(dict_a, dict_b)
# Call and check results.
self.assertEqual(result["tensor"].numpy(), 3)
self.assertEqual(result["variable"].numpy(), 7)
self.assertEqual(result["python_int"].numpy(), 9)
if __name__ == "__main__":
ops.enable_eager_execution()
test.main()
tensorflow/python/eager/polymorphic_function/eager_function_run.py
+111
View File
@@ -0,0 +1,111 @@
# 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.
# ==============================================================================
# pylint: disable=unidiomatic-typecheck
"""Eager semantics for polymorphic function."""
from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
RUN_FUNCTIONS_EAGERLY = False
@tf_export("config.functions_run_eagerly")
def functions_run_eagerly():
"""Returns the value of the `run_functions_eagerly` setting."""
return RUN_FUNCTIONS_EAGERLY
@tf_export("config.run_functions_eagerly")
def run_functions_eagerly(run_eagerly):
"""Enables / disables eager execution of `tf.function`s.
Calling `tf.config.run_functions_eagerly(True)` will make all
invocations of `tf.function` run eagerly instead of running as a traced graph
function. This can be useful for debugging. As the code now runs line-by-line,
you can add arbitrary `print` messages or pdb breakpoints to monitor the
inputs/outputs of each Tensorflow operation. However, you should avoid using
this for actual production because it significantly slows down execution.
>>> def my_func(a):
... print(f'a: {a}')
... return a + a
>>> a_fn = tf.function(my_func)
>>> # A side effect the first time the function is traced
>>> # In tracing time, `a` is printed with shape and dtype only
>>> a_fn(tf.constant(1))
a: Tensor("a:0", shape=(), dtype=int32)
<tf.Tensor: shape=(), dtype=int32, numpy=2>
>>> # `print` is a python side effect, it won't execute as the traced function
>>> # is called
>>> a_fn(tf.constant(2))
<tf.Tensor: shape=(), dtype=int32, numpy=4>
>>> # Now, switch to eager running
>>> tf.config.run_functions_eagerly(True)
>>> # The code now runs eagerly and the actual value of `a` is printed
>>> a_fn(tf.constant(2))
a: 2
<tf.Tensor: shape=(), dtype=int32, numpy=4>
>>> # Turn this back off
>>> tf.config.run_functions_eagerly(False)
Note: This flag has no effect on functions passed into tf.data transformations
as arguments. tf.data functions are never executed eagerly and are always
executed as a compiled Tensorflow Graph.
Args:
run_eagerly: Boolean. Whether to run functions eagerly.
"""
global RUN_FUNCTIONS_EAGERLY
RUN_FUNCTIONS_EAGERLY = bool(run_eagerly)
@deprecation.deprecated(
None, "Use `tf.config.run_functions_eagerly` instead of the experimental "
"version.")
@tf_export("config.experimental_run_functions_eagerly")
def experimental_run_functions_eagerly(run_eagerly):
"""Enables / disables eager execution of `tf.function`s.
Calling `tf.config.experimental_run_functions_eagerly(True)` will make all
invocations of `tf.function` run eagerly instead of running as a traced graph
function.
See `tf.config.run_functions_eagerly` for an example.
Note: This flag has no effect on functions passed into tf.data transformations
as arguments. tf.data functions are never executed eagerly and are always
executed as a compiled Tensorflow Graph.
Args:
run_eagerly: Boolean. Whether to run functions eagerly.
Returns:
None
"""
return run_functions_eagerly(run_eagerly)
@deprecation.deprecated(
None,
"Use tf.config.functions_run_eagerly instead of the experimental version.")
@tf_export("config.experimental_functions_run_eagerly")
def experimental_functions_run_eagerly():
"""Returns the value of the `experimental_run_functions_eagerly` setting."""
return functions_run_eagerly()
tensorflow/python/eager/polymorphic_function/function_context.py
+127
View File
@@ -0,0 +1,127 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Context information for a tf.function."""
from typing import NamedTuple, Any
from tensorflow.core.function.polymorphism import function_cache
from tensorflow.python.eager import context
from tensorflow.python.framework import device as pydev
from tensorflow.python.framework import func_graph as func_graph_module
from tensorflow.python.framework import ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.saved_model import save_context
# EagerContext is used by tf.function to identify cases where tracing
# needs to occur due to a change in conditions other than the arguments.
class EagerContext(NamedTuple):
parent_graph: Any
device_functions: Any
colocation_stack: Any
in_cross_replica_context: Any
variable_policy: Any
xla_context_id: Any
def make_function_context(scope_type=None) -> function_cache.FunctionContext:
"""Generates a FunctionContext based on current contextual info."""
ctx = context.context()
# Don't need to open an init_scope if the tf.function call is in eager mode
# already.
executing_eagerly = ctx.executing_eagerly()
parent_graph = None
xla_context_id = 0
if not executing_eagerly:
# We want to force function retracing for each different
# XLAControlFlowContext, so add `xla_context_id` to the context.
xla_context = _enclosing_xla_context()
if xla_context is not None and xla_context.RequiresUniqueFunctionRetracing(
):
xla_context_id = id(xla_context)
with ops.init_scope():
# The graph, or whether we're executing eagerly, should be a part of the
# cache key so we don't improperly capture tensors such as variables.
executing_eagerly = ctx.executing_eagerly()
parent_graph = None if executing_eagerly else ops.get_default_graph()
# pylint: disable=protected-access
default_graph = ops.get_default_graph()
# TODO(b/117617952): The current distribution strategy will affect graph
# building (e.g. accessing different variables from different devices) and
# so requires retracing for each device.
strategy_stack = default_graph._distribution_strategy_stack
uses_distribution_strategy = (
strategy_stack and
strategy_stack[-1].strategy.extended._retrace_functions_for_each_device)
if executing_eagerly:
colocation_stack = ()
if uses_distribution_strategy:
device_functions = (pydev.merge_device(ctx.device_name),)
else:
device_functions = ()
else:
colocation_stack = tuple(default_graph._colocation_stack.peek_objs())
if (uses_distribution_strategy or
func_graph_module.device_stack_has_callable(
default_graph._device_function_stack)):
# Putting the device in the cache key ensures that call-site device
# annotations are respected.
device_functions = tuple(default_graph._device_functions_outer_to_inner)
else:
device_functions = ()
in_cross_replica_context = False
try:
in_cross_replica_context = (strategy_stack[-1].replica_context is None) # pylint: disable=protected-access
except (AttributeError, IndexError):
pass
if save_context.in_save_context():
variable_policy = (
save_context.get_save_options().experimental_variable_policy)
else:
variable_policy = None
return function_cache.FunctionContext(
EagerContext(
parent_graph,
device_functions,
colocation_stack,
in_cross_replica_context,
variable_policy,
xla_context_id,
),
scope_type,
)
def _enclosing_xla_context():
"""Returns the XLAControlFlowContext, which exists inside a tpu.rewrite()."""
graph = ops.get_default_graph()
while graph is not None:
# pylint: disable=protected-access
context_ = graph._get_control_flow_context()
# pylint: enable=protected-access
while context_ is not None:
if isinstance(context_, control_flow_ops.XLAControlFlowContext):
return context_
context_ = context_.outer_context
# This may be a FuncGraph due to defuns or v2 control flow. We need to
# find the original graph with the XLAControlFlowContext.
graph = getattr(graph, "outer_graph", None)
return None
tensorflow/python/eager/polymorphic_function/function_spec_test.py
+567
View File
@@ -0,0 +1,567 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT 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 function_type_utils."""
from absl.testing import parameterized
from tensorflow.core.function import trace_type
from tensorflow.core.function.polymorphism import function_type as function_type_lib
from tensorflow.python.eager.polymorphic_function import function_type_utils
from tensorflow.python.framework import tensor_spec
from tensorflow.python.platform import test
from tensorflow.python.util import tf_decorator
def dummy_tf_decorator(func):
def wrapper(*args, **kwargs):
return func(*args, **kwargs)
return tf_decorator.make_decorator(func, wrapper)
def transparent_decorator(func):
return func
class FunctionSpecTest(test.TestCase, parameterized.TestCase):
@parameterized.product(
({
'input_signature': None,
'type_constraint': (None, None, None)
}, {
'input_signature': (tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)),
'type_constraint': (tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None))
}, {
'input_signature': ([
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)
], tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)),
'type_constraint': (trace_type.from_value([
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)
], trace_type.InternalTracingContext(is_legacy_signature=True)),
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None))
}),
decorator=(dummy_tf_decorator, transparent_decorator),
)
def test_required_only(self, input_signature, type_constraint, decorator):
@decorator
def foo(x, y, z): # pylint: disable=unused-argument
pass
spec = function_type_utils.FunctionSpec.from_function_and_signature(
foo, input_signature)
self.assertEqual(
tuple(spec.fullargspec),
(['x', 'y', 'z'], None, None, None, [], None, {}))
self.assertEqual(spec.input_signature, input_signature)
self.assertEqual(spec.default_values, {})
self.assertEqual(
spec.function_type,
function_type_lib.FunctionType([
function_type_lib.Parameter(
'x', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD, False,
type_constraint[0]),
function_type_lib.Parameter(
'y', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD, False,
type_constraint[1]),
function_type_lib.Parameter(
'z', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD, False,
type_constraint[2])
]))
@parameterized.product(
({
'input_signature': None,
'type_constraint': (None, None, None)
}, {
'input_signature': (tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)),
'type_constraint': (tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None))
}, {
'input_signature': (tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)),
'type_constraint':
(tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None), trace_type.from_value(3))
}, {
'input_signature': (tensor_spec.TensorSpec(shape=None),),
'type_constraint':
(tensor_spec.TensorSpec(shape=None), trace_type.from_value(2),
trace_type.from_value(3))
}, {
'input_signature': ([
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)
], tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)),
'type_constraint': (trace_type.from_value([
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)
], trace_type.InternalTracingContext(is_legacy_signature=True)),
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None))
}),
decorator=(dummy_tf_decorator, transparent_decorator),
)
def test_optional_only(self, input_signature, type_constraint, decorator):
@decorator
def foo(x=1, y=2, z=3): # pylint: disable=unused-argument
pass
spec = function_type_utils.FunctionSpec.from_function_and_signature(
foo, input_signature)
self.assertEqual(
tuple(spec.fullargspec),
(['x', 'y', 'z'], None, None, (1, 2, 3), [], None, {}))
self.assertEqual(spec.input_signature, input_signature)
self.assertEqual(spec.default_values, {'x': 1, 'y': 2, 'z': 3})
self.assertEqual(
spec.function_type,
function_type_lib.FunctionType([
function_type_lib.Parameter(
'x', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD, True,
type_constraint[0]),
function_type_lib.Parameter(
'y', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD, True,
type_constraint[1]),
function_type_lib.Parameter(
'z', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD, True,
type_constraint[2])
]))
@parameterized.product(
({
'input_signature': None,
'type_constraint': (None, None, None)
}, {
'input_signature': (tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)),
'type_constraint': (tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None))
}, {
'input_signature': (tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)),
'type_constraint':
(tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None), trace_type.from_value(3))
}, {
'input_signature': ([
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)
], tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)),
'type_constraint': (trace_type.from_value([
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)
], trace_type.InternalTracingContext(is_legacy_signature=True)),
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None))
}),
decorator=(dummy_tf_decorator, transparent_decorator),
)
def test_required_and_optional(self, input_signature, type_constraint,
decorator):
@decorator
def foo(x, y, z=3): # pylint: disable=unused-argument
pass
spec = function_type_utils.FunctionSpec.from_function_and_signature(
foo, input_signature)
self.assertEqual(
tuple(spec.fullargspec),
(['x', 'y', 'z'], None, None, (3,), [], None, {}))
self.assertEqual(spec.input_signature, input_signature)
self.assertEqual(spec.default_values, {'z': 3})
self.assertEqual(
spec.function_type,
function_type_lib.FunctionType([
function_type_lib.Parameter(
'x', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD, False,
type_constraint[0]),
function_type_lib.Parameter(
'y', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD, False,
type_constraint[1]),
function_type_lib.Parameter(
'z', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD, True,
type_constraint[2])
]))
@parameterized.product(
({
'input_signature': (tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)),
'type_constraint': (tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None))
}, {
'input_signature': ([
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)
], tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)),
'type_constraint': (trace_type.from_value([
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)
], trace_type.InternalTracingContext(is_legacy_signature=True)),
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None))
}),
decorator=(dummy_tf_decorator, transparent_decorator),
)
def test_varargs(self, input_signature, type_constraint, decorator):
@decorator
def foo(*my_var_args): # pylint: disable=unused-argument
pass
spec = function_type_utils.FunctionSpec.from_function_and_signature(
foo, input_signature)
self.assertEqual(
tuple(spec.fullargspec),
(['my_var_args_0', 'my_var_args_1', 'my_var_args_2'
], None, None, None, [], None, {}))
self.assertEqual(spec.input_signature, input_signature)
self.assertEqual(spec.default_values, {})
self.assertEqual(
spec.function_type,
function_type_lib.FunctionType([
function_type_lib.Parameter(
'my_var_args_0', function_type_lib.Parameter.POSITIONAL_ONLY,
False, type_constraint[0]),
function_type_lib.Parameter(
'my_var_args_1', function_type_lib.Parameter.POSITIONAL_ONLY,
False, type_constraint[1]),
function_type_lib.Parameter(
'my_var_args_2', function_type_lib.Parameter.POSITIONAL_ONLY,
False, type_constraint[2])
]))
@parameterized.product(
({
'input_signature': None,
'type_constraint': (None, None, None)
}, {
'input_signature': (tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)),
'type_constraint': (tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None))
}, {
'input_signature': (tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)),
'type_constraint':
(tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None), trace_type.from_value(3))
}, {
'input_signature': ([
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)
], tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)),
'type_constraint': (trace_type.from_value([
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)
], trace_type.InternalTracingContext(is_legacy_signature=True)),
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None))
}),
decorator=(dummy_tf_decorator, transparent_decorator),
)
def test_kwonly(self, input_signature, type_constraint, decorator):
@decorator
def foo(x, y, *, z=3): # pylint: disable=unused-argument
pass
spec = function_type_utils.FunctionSpec.from_function_and_signature(
foo, input_signature)
self.assertEqual(
tuple(spec.fullargspec), (['x', 'y'], None, None, None, ['z'], {
'z': 3
}, {}))
self.assertEqual(spec.input_signature, input_signature)
self.assertEqual(spec.default_values, {'z': 3})
self.assertEqual(
spec.function_type,
function_type_lib.FunctionType([
function_type_lib.Parameter(
'x', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD, False,
type_constraint[0]),
function_type_lib.Parameter(
'y', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD, False,
type_constraint[1]),
function_type_lib.Parameter(
'z', function_type_lib.Parameter.KEYWORD_ONLY, True,
type_constraint[2])
]))
@parameterized.product(
({
'input_signature': None,
'type_constraint': (None, None, None)
}, {
'input_signature': (tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)),
'type_constraint': (None, tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None))
}, {
'input_signature': (tensor_spec.TensorSpec(shape=None),),
'type_constraint': (None, tensor_spec.TensorSpec(shape=None),
trace_type.from_value(1))
}, {
'input_signature': ([
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)
], tensor_spec.TensorSpec(shape=None)),
'type_constraint':
(None,
trace_type.from_value([
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)
], trace_type.InternalTracingContext(is_legacy_signature=True)),
tensor_spec.TensorSpec(shape=None))
}),
decorator=(dummy_tf_decorator, transparent_decorator),
)
def test_method_bound_internal(
self, input_signature, type_constraint, decorator
):
def testing_decorator(func):
spec = function_type_utils.FunctionSpec.from_function_and_signature(
func, input_signature
)
self.assertEqual(
tuple(spec.fullargspec),
(['self', 'x', 'y'], None, None, (1,), [], None, {}),
)
self.assertEqual(spec.default_values, {'y': 1})
self.assertEqual(
spec.function_type,
function_type_lib.FunctionType([
function_type_lib.Parameter(
'self',
function_type_lib.Parameter.POSITIONAL_OR_KEYWORD,
False,
type_constraint[0],
),
function_type_lib.Parameter(
'x',
function_type_lib.Parameter.POSITIONAL_OR_KEYWORD,
False,
type_constraint[1],
),
function_type_lib.Parameter(
'y',
function_type_lib.Parameter.POSITIONAL_OR_KEYWORD,
True,
type_constraint[2],
),
]),
)
return func
class MyClass:
@testing_decorator
def foo(self, x, y=1):
pass
MyClass().foo(1)
@parameterized.product(
({
'input_signature': None,
'type_constraint': (None, None, None)
}, {
'input_signature': (tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)),
'type_constraint': (tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None))
}, {
'input_signature': (tensor_spec.TensorSpec(shape=None),),
'type_constraint': (tensor_spec.TensorSpec(shape=None),
trace_type.from_value(1))
}, {
'input_signature': ([
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)
], tensor_spec.TensorSpec(shape=None)),
'type_constraint': (trace_type.from_value([
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)
], trace_type.InternalTracingContext(is_legacy_signature=True)),
tensor_spec.TensorSpec(shape=None))
}),
decorator=(dummy_tf_decorator, transparent_decorator),
)
def test_method_bound_external(
self, input_signature, type_constraint, decorator
):
class MyClass:
@decorator
def foo(self, x, y=1):
pass
spec = function_type_utils.FunctionSpec.from_function_and_signature(
MyClass().foo, input_signature)
self.assertEqual(
tuple(spec.fullargspec),
(['x', 'y'], None, None, (1,), [], None, {}),
)
self.assertEqual(spec.default_values, {'y': 1})
self.assertEqual(
spec.function_type,
function_type_lib.FunctionType([
function_type_lib.Parameter(
'x', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD, False,
type_constraint[0]),
function_type_lib.Parameter(
'y', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD, True,
type_constraint[1])
]))
@parameterized.product(
({
'input_signature': None,
'type_constraint': (None, None, None)
}, {
'input_signature': (tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)),
'type_constraint': (None, tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None))
}, {
'input_signature': (tensor_spec.TensorSpec(shape=None),),
'type_constraint': (None, tensor_spec.TensorSpec(shape=None),
trace_type.from_value(1))
}, {
'input_signature': ([
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)
], tensor_spec.TensorSpec(shape=None)),
'type_constraint':
(None,
trace_type.from_value([
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None)
], trace_type.InternalTracingContext(is_legacy_signature=True)),
tensor_spec.TensorSpec(shape=None))
}),
decorator=(dummy_tf_decorator, transparent_decorator),
)
def test_method_unbound(self, input_signature, type_constraint, decorator):
class MyClass:
@decorator
def foo(self, x, y=1):
pass
spec = function_type_utils.FunctionSpec.from_function_and_signature(
MyClass.foo, input_signature)
self.assertEqual(
tuple(spec.fullargspec),
(['self', 'x', 'y'], None, None, (1,), [], None, {}))
self.assertEqual(spec.input_signature, input_signature)
self.assertEqual(spec.default_values, {'y': 1})
self.assertEqual(
spec.function_type,
function_type_lib.FunctionType([
function_type_lib.Parameter(
'self', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD,
False, type_constraint[0]),
function_type_lib.Parameter(
'x', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD, False,
type_constraint[1]),
function_type_lib.Parameter(
'y', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD, True,
type_constraint[2])
]))
def test_spec_summary(self):
input_signature = (
tensor_spec.TensorSpec(shape=None),
tensor_spec.TensorSpec(shape=None),
)
@dummy_tf_decorator
def foo(x=2, y=3): # pylint: disable=unused-argument
pass
spec = function_type_utils.FunctionSpec.from_function_and_signature(
foo, input_signature
)
self.assertEqual(
spec.signature_summary(True),
'Input Parameters:\n'
+ ' x (POSITIONAL_OR_KEYWORD):'
' TensorSpec(shape=<unknown>, dtype=tf.float32, name=None)\n'
+ ' y'
' (POSITIONAL_OR_KEYWORD): TensorSpec(shape=<unknown>,'
' dtype=tf.float32, name=None)\n'
+ 'Output Type:\n'
+ ' None\n'
+ 'Captures:\n'
+ ' None\n'
+ 'Defaults:\n'
+ ' x: 2\n'
+ ' y: 3',
)
# TODO(fmuham): Remove when is_same_structure is removed.
class SameStructureTest(test.TestCase):
def test_same_structure(self):
self.assertTrue(
function_type_utils.is_same_structure([1, 2, 3], [1, 2, 3], True)
)
self.assertTrue(
function_type_utils.is_same_structure([1, 2, 3], [1, 2, 4], False)
)
self.assertFalse(
function_type_utils.is_same_structure([1, 2, 3], [1, 2, 4], True)
)
self.assertFalse(
function_type_utils.is_same_structure([1, 2, 3], [1, 2, 3, 4], False)
)
if __name__ == '__main__':
test.main()
tensorflow/python/eager/polymorphic_function/function_type_utils.py
+548
View File
@@ -0,0 +1,548 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Utilities for using FunctionType with tf.function."""
import functools
import inspect
from typing import Any, Dict, Tuple
import six
from tensorflow.core.function import trace_type
from tensorflow.core.function.polymorphism import function_type as function_type_lib
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor
from tensorflow.python.framework import type_spec
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.util import nest
def to_fullargspec(
function_type: function_type_lib.FunctionType,
default_values: Dict[str, Any],
) -> inspect.FullArgSpec:
"""Generates backwards compatible FullArgSpec from FunctionType."""
args = []
varargs = None
varkw = None
defaults = []
kwonlyargs = []
kwonlydefaults = {}
for parameter in function_type.parameters.values():
if parameter.kind in [
inspect.Parameter.POSITIONAL_ONLY,
inspect.Parameter.POSITIONAL_OR_KEYWORD,
]:
args.append(parameter.name)
if parameter.default is not inspect.Parameter.empty:
defaults.append(default_values[parameter.name])
elif parameter.kind is inspect.Parameter.KEYWORD_ONLY:
kwonlyargs.append(parameter.name)
if parameter.default is not inspect.Parameter.empty:
kwonlydefaults[parameter.name] = default_values[parameter.name]
elif parameter.kind is inspect.Parameter.VAR_POSITIONAL:
varargs = parameter.name
elif parameter.kind is inspect.Parameter.VAR_KEYWORD:
varkw = parameter.name
return inspect.FullArgSpec(
args,
varargs,
varkw,
tuple(defaults) if defaults else None,
kwonlyargs,
kwonlydefaults if kwonlydefaults else None,
annotations={},
)
def _to_default_values(fullargspec):
"""Returns default values from the function's inspected fullargspec."""
if fullargspec.defaults is not None:
defaults = {
name: value
for name, value in zip(
fullargspec.args[-len(fullargspec.defaults) :], fullargspec.defaults
)
}
else:
defaults = {}
if fullargspec.kwonlydefaults is not None:
defaults.update(fullargspec.kwonlydefaults)
defaults = {
function_type_lib.sanitize_arg_name(name): value
for name, value in defaults.items()
}
return defaults
def to_function_type(fullargspec):
"""Generates FunctionType and default values from fullargspec."""
default_values = _to_default_values(fullargspec)
parameters = []
for arg in fullargspec.args:
arg_name = function_type_lib.sanitize_arg_name(arg)
parameters.append(
function_type_lib.Parameter(
arg_name,
function_type_lib.Parameter.POSITIONAL_OR_KEYWORD,
arg_name in default_values,
None,
)
)
if fullargspec.varargs is not None:
parameters.append(
function_type_lib.Parameter(
fullargspec.varargs,
function_type_lib.Parameter.VAR_POSITIONAL,
False,
None,
)
)
for kwarg in fullargspec.kwonlyargs:
parameters.append(
function_type_lib.Parameter(
function_type_lib.sanitize_arg_name(kwarg),
function_type_lib.Parameter.KEYWORD_ONLY,
kwarg in default_values,
None,
)
)
if fullargspec.varkw is not None:
parameters.append(
function_type_lib.Parameter(
fullargspec.varkw,
function_type_lib.Parameter.VAR_KEYWORD,
False,
None,
)
)
return function_type_lib.FunctionType(parameters), default_values
def to_input_signature(function_type):
"""Extracts an input_signature from function_type instance."""
constrained_parameters = list(function_type.parameters.keys())
# self does not have a constraint in input_signature
if "self" in constrained_parameters:
constrained_parameters.pop(0)
# There are no parameters to constrain.
if not constrained_parameters:
return tuple()
constraints = []
is_auto_constrained = False
for parameter_name in constrained_parameters:
parameter = function_type.parameters[parameter_name]
constraint = None
if parameter.type_constraint:
# Generate legacy constraint representation.
constraint = parameter.type_constraint.placeholder_value(
trace_type.InternalPlaceholderContext(unnest_only=True)
)
if any(
not isinstance(arg, tensor.TensorSpec)
for arg in nest.flatten([constraint], expand_composites=True)
):
# input_signature only supports contiguous TensorSpec composites
is_auto_constrained = True
break
else:
constraints.append(constraint)
# All constraints were generated by FunctionType
if is_auto_constrained and not constraints:
return tuple()
# If the list is empty then there was no input_signature specified.
return tuple(constraints) if constraints else None
def to_arg_names(function_type):
"""Generates a list of arg names from a FunctionType."""
arg_names = []
for p in function_type.parameters.values():
if p.kind in {
function_type_lib.Parameter.POSITIONAL_ONLY,
function_type_lib.Parameter.POSITIONAL_OR_KEYWORD,
}:
arg_names.append(p.name)
return arg_names
# TODO(b/214462107): Minimize API surface for FunctionSpec.
class FunctionSpec(object):
"""Specification of how to bind arguments to a function.
Deprecated. Please use FunctionType instead.
"""
@classmethod
def from_function_and_signature(
cls, python_function, input_signature, is_pure=False, jit_compile=None
):
"""Creates a FunctionSpec instance given a python function and signature.
Args:
python_function: a function to inspect
input_signature: a signature of the function (None, if variable)
is_pure: if True all input arguments (including variables and constants)
will be converted to tensors and no variable changes allowed.
jit_compile: see `tf.function`
Returns:
instance of FunctionSpec
"""
function_type, default_values = make_function_type(
python_function, input_signature)
# Get the function's name. Remove functools.partial wrappers if necessary.
while isinstance(python_function, functools.partial):
python_function = python_function.func
name = getattr(python_function, "__name__", "f")
return FunctionSpec(
function_type,
default_values,
is_pure=is_pure,
jit_compile=jit_compile,
name=name,
)
@classmethod
def from_fullargspec_and_signature(
cls,
fullargspec,
input_signature,
is_pure=False,
name=None,
jit_compile=None,
):
"""Construct FunctionSpec from legacy FullArgSpec format."""
function_type, default_values = to_function_type(fullargspec)
if input_signature:
input_signature = tuple(input_signature)
_validate_signature(input_signature)
function_type = function_type_lib.add_type_constraints(
function_type, input_signature, default_values
)
return FunctionSpec(
function_type, default_values, is_pure, name, jit_compile
)
def __init__(
self,
function_type,
default_values,
is_pure=False,
name=None,
jit_compile=None,
):
"""Constructs a FunctionSpec describing a python function.
Args:
function_type: A FunctionType describing the python function signature.
default_values: Dictionary mapping parameter names to default values.
is_pure: if True all input arguments (including variables and constants)
will be converted to tensors and no variable changes allowed.
name: Name of the function
jit_compile: see `tf.function`.
"""
self._function_type = function_type
self._default_values = default_values
self._fullargspec = to_fullargspec(function_type, default_values)
self._is_pure = is_pure
self._jit_compile = jit_compile
# TODO(edloper): Include name when serializing for SavedModel?
self._name = name or "f"
self._input_signature = to_input_signature(function_type)
@property
def default_values(self):
"""Returns dict mapping parameter names to default values."""
return self._default_values
@property
def function_type(self):
"""Returns a FunctionType representing the Python function signature."""
return self._function_type
@property
def fullargspec(self):
return self._fullargspec
# TODO(fmuham): Replace usages with FunctionType and remove.
@property
def input_signature(self):
return self._input_signature
# TODO(fmuham): Replace usages with FunctionType and remove.
@property
def flat_input_signature(self):
return tuple(nest.flatten(self.input_signature, expand_composites=True))
@property
def is_pure(self):
return self._is_pure
@property
def jit_compile(self):
return self._jit_compile
# TODO(fmuham): Replace usages and remove.
@property
def arg_names(self):
return to_arg_names(self.function_type)
def signature_summary(self, default_values=False):
"""Returns a string summarizing this function's signature.
Args:
default_values: If true, then include default values in the signature.
Returns:
A `string`.
"""
summary = f"{self._function_type!r}"
if default_values:
summary += "\nDefaults:"
if self.default_values:
for name, value in self.default_values.items():
summary += f"\n {name}: {value!r}"
else:
summary += "\n None"
return summary
def make_function_type(python_function, input_signature):
"""Generates a FunctionType for python_function."""
_validate_signature(input_signature)
function_type = function_type_lib.FunctionType.from_callable(
python_function
)
default_values = function_type_lib.FunctionType.get_default_values(
python_function
)
if input_signature is not None:
input_signature = tuple(input_signature)
function_type = function_type_lib.add_type_constraints(
function_type, input_signature, default_values
)
return function_type, default_values
def make_canonicalized_monomorphic_type(
args: Any,
kwargs: Any,
capture_types: Any,
polymorphic_type,
) -> Tuple[function_type_lib.FunctionType, trace_type.InternalTracingContext]:
"""Generates function type given the function arguments."""
kwargs = {
function_type_lib.sanitize_arg_name(name): value
for name, value in kwargs.items()
}
function_type, type_context = (
function_type_lib.canonicalize_to_monomorphic(
args, kwargs, {}, capture_types, polymorphic_type
)
)
return function_type, type_context
def canonicalize_function_inputs(
args, kwargs, function_type, default_values=None, is_pure=False
):
"""Canonicalizes `args` and `kwargs`.
Canonicalize the inputs to the Python function using FunctionType.
In particular, we parse the varargs and kwargs that the
original function was called with into a tuple corresponding to the
Python function's positional (named) arguments and a dictionary
corresponding to its kwargs. Missing default arguments are added.
If the FunctionType has an type constraints, then they are used to convert
arguments to tensors; otherwise, any inputs containing numpy arrays are
converted to tensors.
Args:
args: The varargs this object was called with.
kwargs: The keyword args this function was called with.
function_type: FunctionType to canonicalize against.
default_values: Default values to use.
is_pure: Force variable inputs to Tensors.
Returns:
A canonicalized ordering of the inputs, as well as full and filtered
(Tensors and Variables only) versions of their concatenated flattened
representations, represented by a tuple in the form (args, kwargs,
flat_args, filtered_flat_args). Here: `args` is a full list of bound
arguments, and `kwargs` contains only true keyword arguments, as opposed
to named arguments called in a keyword-like fashion.
Raises:
ValueError: If a keyword in `kwargs` cannot be matched with a positional
argument when an input signature is specified, or when the inputs
do not conform to the input signature.
"""
default_values = {} if not default_values else default_values
if is_pure:
args, kwargs = _convert_variables_to_tensors(args, kwargs)
bound_arguments = bind_function_inputs(
args, kwargs, function_type, default_values
)
return bound_arguments
def bind_function_inputs(args, kwargs, function_type, default_values):
"""Bind `args` and `kwargs` into a canonicalized signature args, kwargs."""
sanitized_kwargs = {
function_type_lib.sanitize_arg_name(k): v for k, v in kwargs.items()
}
if len(kwargs) != len(sanitized_kwargs):
raise ValueError(
"Name collision after sanitization. Please rename "
"tf.function input parameters. Original: "
f"{sorted(kwargs.keys())}, Sanitized: "
f"{sorted(sanitized_kwargs.keys())}"
)
try:
bound_arguments = function_type.bind_with_defaults(
args, sanitized_kwargs, default_values
)
except Exception as e:
raise TypeError(
f"Binding inputs to tf.function failed due to `{e}`. "
f"Received args: {args} and kwargs: {sanitized_kwargs} for signature:"
f" {function_type}."
) from e
return bound_arguments
def _validate_signature(signature):
"""Checks the input_signature to be valid."""
if signature is None:
return
if not isinstance(signature, (tuple, list)):
raise TypeError(
"input_signature must be either a tuple or a list, got "
f"{type(signature)}."
)
# TODO(xjun): Allow VariableSpec once we figure out API for de-aliasing.
variable_specs = _get_variable_specs(signature)
if variable_specs:
raise TypeError(
f"input_signature doesn't support VariableSpec, got {variable_specs}"
)
if any(
not isinstance(arg, tensor.TensorSpec)
for arg in nest.flatten(signature, expand_composites=True)
):
bad_args = [
arg
for arg in nest.flatten(signature, expand_composites=True)
if not isinstance(arg, tensor.TensorSpec)
]
raise TypeError(
"input_signature must be a possibly nested sequence of "
f"TensorSpec objects, got invalid args {bad_args} with "
f"types {list(six.moves.map(type, bad_args))}."
)
def _to_tensor_or_tensor_spec(x):
return (
x
if isinstance(x, (tensor.Tensor, tensor.TensorSpec))
else ops.convert_to_tensor(x)
)
def _convert_variables_to_tensors(args, kwargs):
args = [_to_tensor_or_tensor_spec(x) for x in args]
kwargs = {kw: _to_tensor_or_tensor_spec(x) for kw, x in kwargs.items()}
return tuple(args), kwargs
def _get_variable_specs(args):
"""Returns `VariableSpecs` from `args`."""
variable_specs = []
for arg in nest.flatten(args):
if not isinstance(arg, type_spec.TypeSpec):
continue
if isinstance(arg, resource_variable_ops.VariableSpec):
variable_specs.append(arg)
elif not isinstance(arg, tensor.TensorSpec):
# arg is a CompositeTensor spec.
variable_specs.extend(_get_variable_specs(arg._component_specs)) # pylint: disable=protected-access
return variable_specs
def derive_from_graph(func_graph):
"""Derives a FunctionType from FuncGraph."""
# TODO(fmuham): Include structure info from structured_inputs
input_signature = (
tuple(trace_type.from_value(i) for i in func_graph.inputs),
{},
)
# TODO(fmuham): Include output structure info from structured_outputs
output_signature = tuple(trace_type.from_value(o) for o in func_graph.outputs)
return function_type_lib.from_structured_signature(
input_signature,
output_signature,
func_graph.function_captures.capture_types,
)
# TODO(fmuham): Replace usages with TraceType and remove.
def is_same_structure(structure1, structure2, check_values=False):
"""Check two structures for equality, optionally of types and of values."""
try:
nest.assert_same_structure(structure1, structure2, expand_composites=True)
except (ValueError, TypeError):
return False
if check_values:
flattened1 = nest.flatten(structure1, expand_composites=True)
flattened2 = nest.flatten(structure2, expand_composites=True)
# First check the types to avoid AttributeErrors.
if any(type(f1) is not type(f2) for f1, f2 in zip(flattened1, flattened2)):
return False
return flattened1 == flattened2
return True
tensorflow/python/eager/polymorphic_function/gradients_test.py
+927
View File
@@ -0,0 +1,927 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
from absl.testing import parameterized
from tensorflow.python.eager import backprop
from tensorflow.python.eager import context
from tensorflow.python.eager.polymorphic_function import polymorphic_function
from tensorflow.python.framework import config
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import indexed_slices
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import cond
from tensorflow.python.ops import gradients_impl
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn_grad
from tensorflow.python.ops import nn_ops
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import variables
from tensorflow.python.ops import while_loop
from tensorflow.python.platform import test
_COS_DERIVATIVES = [math_ops.cos,
lambda x: -math_ops.sin(x),
lambda x: -math_ops.cos(x),
math_ops.sin,
math_ops.cos]
class FunctionGradientsTest(test.TestCase, parameterized.TestCase):
def setUp(self):
super(FunctionGradientsTest, self).setUp()
cpus = config.list_physical_devices('CPU')
# Set 4 virtual CPUs
config.set_logical_device_configuration(cpus[0], [
context.LogicalDeviceConfiguration(),
context.LogicalDeviceConfiguration(),
context.LogicalDeviceConfiguration(),
context.LogicalDeviceConfiguration()
])
def testGraphModeWithGradients(self):
v = resource_variable_ops.ResourceVariable(1.0, name='v')
@polymorphic_function.function
def step():
def inner():
return v * v
return backprop.implicit_grad(inner)()[0][0]
self.assertAllEqual(step(), 2.0)
def testGraphGradientVariable(self):
with ops.Graph().as_default(), self.cached_session():
v = variables.Variable(1.0)
@polymorphic_function.function
def f():
return 2.0 * v
node = f()
grads, = gradients_impl.gradients(node, v)
v.initializer.run()
self.assertAllEqual(grads, 2.0)
self.assertEqual(grads.shape, v.shape)
def testSymbolicHigherOrder(self):
@polymorphic_function.function
def f(x, order):
y = polymorphic_function.function(lambda: math_ops.cos(x))()
for _ in range(order):
y, = gradients_impl.gradients(y, [x])
return y
for order, expected in enumerate(_COS_DERIVATIVES):
self.assertAllClose(
expected(constant_op.constant(1.)),
f(constant_op.constant(1.), order))
@parameterized.parameters([dict(persistent=True),
dict(persistent=False)])
def testSymbolicHigherOrderUnderTape(self, persistent):
@polymorphic_function.function
def f(x, order):
with backprop.GradientTape(persistent=persistent) as tape:
tape.watch(x)
# Note that having a tape active, even if we don't use it, forces us
# down a different function call path. Symbolic gradients should work
# here too; correctness of tape gradients are tested elsewhere.
y = polymorphic_function.function(lambda: math_ops.cos(x))()
tape_dy = tape.gradient(y, x)
for _ in range(order):
y, = gradients_impl.gradients(y, [x])
if order > 0:
y1 = tape_dy
for _ in range(order - 1):
y1, = gradients_impl.gradients(y1, [x])
else:
y1 = y
return y, y1
for order, expected_f in enumerate(_COS_DERIVATIVES):
expected = self.evaluate(expected_f(constant_op.constant(1.)))
self.assertAllClose(
(expected, expected),
f(constant_op.constant(1.), order))
def testIteratedGradientsNested(self):
def _grad(f):
def _grad_function(primal):
with backprop.GradientTape() as tape:
tape.watch(primal)
primal_out = f(primal)
return tape.gradient(primal_out, primal)
return _grad_function
@polymorphic_function.function
def _forward(x):
return math_ops.cos(x)
f = _forward
traced_f = polymorphic_function.function(f)
one = constant_op.constant(1.)
for expected in _COS_DERIVATIVES:
self.assertAllClose(expected(one), f(one))
self.assertAllClose(expected(one), traced_f(one))
self.assertAllClose(expected(one), polymorphic_function.function(f)(one))
f = _grad(f)
traced_f = polymorphic_function.function(_grad(traced_f))
def testIteratedGradientsNestedWithVariable(self):
def _grad(f):
def _grad_function():
with backprop.GradientTape() as tape:
primal_out = f()
g, = tape.gradient(primal_out, tape.watched_variables())
return g
return _grad_function
v = variables.Variable(2.)
@polymorphic_function.function
def _forward():
return math_ops.cos(v)
f = _forward
two = constant_op.constant(2.)
for expected in _COS_DERIVATIVES:
self.assertAllClose(expected(two), f())
self.assertAllClose(expected(two), polymorphic_function.function(f)())
f = _grad(f)
def testIteratedGradientsPersistent(self):
@polymorphic_function.function
def _forward(z):
return math_ops.cos(z)
f = _forward
with backprop.GradientTape(persistent=True) as tape:
start = constant_op.constant(1.)
tape.watch(start)
x = f(start)
for expected in _COS_DERIVATIVES:
self.assertAllClose(expected(start), x)
x = tape.gradient(x, start)
def testHigherOrderWithVariable(self):
v = variables.Variable(1.)
@polymorphic_function.function
def _forward():
return math_ops.cos(v)
f = _forward
with backprop.GradientTape(persistent=True) as tape:
x = f()
for expected in _COS_DERIVATIVES:
self.assertAllClose(expected(constant_op.constant(1.)), x)
x, = tape.gradient(x, tape.watched_variables())
def testGradientsChained(self):
@polymorphic_function.function
def _forward(z):
return math_ops.cos(z)
f = _forward
x = constant_op.constant(1.)
with backprop.GradientTape() as t:
t.watch(x)
y = f(x)
with backprop.GradientTape() as tt:
doutputs = constant_op.constant(2.)
tt.watch(doutputs)
g = t.gradient(y, x, doutputs)
self.assertAllClose(-2. * math_ops.sin(x), g)
gg = tt.gradient(g, doutputs)
# We're taking gradients with respect to doutputs, which is just a linear
# function of the gradient.
self.assertAllClose(-math_ops.sin(x), gg)
def testSymGradGatherNd(self):
with ops.Graph().as_default(), self.cached_session():
@polymorphic_function.function
def f(x):
return array_ops.gather_nd(x, [[0]])
c = constant_op.constant([[2.]])
f_c = f(c)
g, = gradients_impl.gradients(f_c, c)
self.assertAllEqual(self.evaluate(g).values, [[1.0]])
def testNoSymGradNestedDefun(self):
@polymorphic_function.function
def outer():
@polymorphic_function.function
def f(x):
return array_ops.gather_nd(x, [[0]])
c = constant_op.constant([[2.]])
f_c = f(c)
g, = gradients_impl.gradients(f_c, c)
self.assertIsInstance(g, indexed_slices.IndexedSlices)
outer()
def testGraphFunctionWithGradients(self):
v = resource_variable_ops.ResourceVariable(1.0, name='v')
@polymorphic_function.function
def step():
def inner():
return v * v
return backprop.implicit_grad(inner)()[0][0]
step_op = step.get_concrete_function()
self.assertEqual(step_op.output_dtypes, dtypes.float32)
self.assertEqual(step_op.output_shapes, tensor_shape.TensorShape([]))
self.assertAllEqual(step_op(), 2.0)
@test_util.run_in_graph_and_eager_modes()
def testDefunCondGradient(self):
@polymorphic_function.function
def f(x):
return cond.cond(x > 0.5, lambda: 2 * x, lambda: 3 * x)
with backprop.GradientTape() as t:
x = constant_op.constant(1.0)
t.watch(x)
y = f(x)
self.assertAllEqual(self.evaluate(t.gradient(y, x)), 2.0)
@test_util.run_in_graph_and_eager_modes()
def testGraphLoopGradient(self):
@polymorphic_function.function
def f(x):
return while_loop.while_loop(
lambda _, i: i < 2, lambda x, i: (2 * x, i + 1), [x, 0]
)[0]
with backprop.GradientTape() as t:
x = constant_op.constant(1.0)
t.watch(x)
y = f(x)
self.assertAllEqual(self.evaluate(t.gradient(y, x)), 4.0)
def testGraphLoopGradientInsideSession(self):
with ops.Graph().as_default():
n = constant_op.constant(2.0)
x = array_ops.placeholder(dtypes.float32, shape=None)
@polymorphic_function.function
def f():
c = lambda n: n < 10
b = lambda n: n * x
return while_loop.while_loop(c, b, [n], [tensor_shape.unknown_shape()])
l = f()
dx = gradients_impl.gradients(l, [x])[0]
with self.cached_session():
self.assertEqual(dx.eval(feed_dict={x: 2.0}), 24.0)
def testDefunDifferentiable(self):
v = resource_variable_ops.ResourceVariable(1.0)
@polymorphic_function.function
def f():
return v * v
self.assertAllEqual(backprop.implicit_grad(f)()[0][0], 2.0)
def testDefunCanBeDifferentiatedTwice(self):
v = resource_variable_ops.ResourceVariable(1.0)
@polymorphic_function.function
def f():
return v * v
self.assertAllEqual(backprop.implicit_grad(f)()[0][0], 2.0)
# Ensure that v is watched again.
self.assertAllEqual(backprop.implicit_grad(f)()[0][0], 2.0)
def testSymbolicGradientVariableNoneNotZerosLike(self):
with ops.Graph().as_default():
v = variables.Variable(1.0)
@polymorphic_function.function
def f(x, v):
v.read_value()
return x * x
x = constant_op.constant(1.0)
l = f(x, v)
_, dv = gradients_impl.gradients(l, [x, v])
with self.cached_session():
v.initializer.run()
self.assertEqual(dv, None)
def testDefunCallBackprop(self):
@polymorphic_function.function
def f(x):
return math_ops.add(x, x)
@polymorphic_function.function
def g(x):
return backprop.gradients_function(f, [0])(x)[0]
self.assertAllEqual(2, g(constant_op.constant(2.)))
@test_util.run_v1_only('b/120545219')
def testGraphModeEagerGradError(self):
with context.graph_mode():
def f():
x = variable_scope.get_variable(
'v', initializer=constant_op.constant(1.0))
return x * constant_op.constant(2.0)
with self.assertRaisesRegex(ValueError,
'No trainable variables were accessed'):
backprop.implicit_val_and_grad(f)()
def testDefunCallBackpropUsingSameObjectForMultipleArguments(self):
@polymorphic_function.function
def g(x):
return backprop.gradients_function(math_ops.multiply, [0, 1])(x, x)
def np_g(x):
return [d.numpy() for d in g(x)]
x = constant_op.constant(1.)
self.assertAllEqual([1., 1.], np_g(x))
self.assertAllEqual([1., 1.], np_g(1.))
def testGradientTensorConversionWithDefun(self):
three = resource_variable_ops.ResourceVariable(3.0, name='v')
@polymorphic_function.function
def f(x):
return math_ops.add(x, three)
def g(x):
return f(x)
g = backprop.implicit_grad(g)(constant_op.constant(1.0))[0][0]
self.assertAllEqual(g, 1.0)
def testGradient(self):
matmul = polymorphic_function.function(math_ops.matmul)
def sq(x):
return matmul(x, x, transpose_a=True)
t = constant_op.constant([[1.0, 2.0], [3.0, 4.0]])
grad_t, = backprop.gradients_function(sq, [0])(t)
self.assertAllEqual(grad_t, [[6, 6], [14, 14]])
def testGradientInFunction(self):
@polymorphic_function.function
def f(x):
return backprop.gradients_function(lambda y: y * y, [0])(x)[0]
self.assertAllEqual(f(constant_op.constant(1.0)), 2.0)
def testGradientOfGatherWithDefun(self):
v = resource_variable_ops.ResourceVariable([0.0, 1.0, 2.0])
def sum_gather():
return math_ops.reduce_sum(array_ops.gather(v, [1, 2]))
grad_fn = backprop.implicit_grad(sum_gather)
gradient = grad_fn()
defun_grad_fn = backprop.implicit_grad(
polymorphic_function.function(sum_gather))
defun_gradient = defun_grad_fn()
self.assertEqual(len(gradient), len(defun_gradient))
gradient = gradient[0][0]
defun_gradient = defun_gradient[0][0]
self.assertAllEqual(gradient.values, defun_gradient.values)
self.assertAllEqual(gradient.indices, defun_gradient.indices)
self.assertAllEqual(gradient.dense_shape, defun_gradient.dense_shape)
def testDifferentiableFunctionNoneOutputs(self):
@polymorphic_function.function
def my_function(x):
return x, None
def wrapper(x):
return my_function(x)[0]
g = backprop.gradients_function(wrapper, [0])(constant_op.constant(0.0))
self.assertAllEqual(g[0], 1.)
@polymorphic_function.function
def foo(a):
return None, a * a
x = constant_op.constant(5.0)
with backprop.GradientTape() as tp:
tp.watch(x)
none, r = foo(x)
g = tp.gradient(r, x)
self.assertIs(none, None)
self.assertAllEqual(r, 25.0)
self.assertAllEqual(g, 2 * 5.0)
@test_util.run_in_graph_and_eager_modes
def testNestedDifferentiableFunction(self):
@polymorphic_function.function
def inner_fn(a, b):
return a * math_ops.add(a, b)
@polymorphic_function.function
def outer_fn(x):
return inner_fn(x, 1.0)
x = constant_op.constant(5.0)
with backprop.GradientTape() as tp:
tp.watch(x)
result = outer_fn(x)
grad = tp.gradient(result, x)
self.assertAllEqual(grad, 2 * 5.0 + 1.0)
@test_util.run_in_graph_and_eager_modes
def testDeeplyNestedDifferentiableFunction(self):
@polymorphic_function.function
def inner_inner_fn(a, b):
return math_ops.add(a, b)
@polymorphic_function.function
def inner_fn(a, b):
return inner_inner_fn(a, b)
@polymorphic_function.function
def middle_fn(a, b):
return a * inner_fn(a, b)
@polymorphic_function.function
def outer_fn(x):
return middle_fn(x, 1.0)
x = constant_op.constant(5.0)
with backprop.GradientTape() as tp:
tp.watch(x)
result = outer_fn(x)
grad = tp.gradient(result, x)
self.assertAllEqual(grad, 2 * 5.0 + 1.0)
@test_util.run_in_graph_and_eager_modes
def testDeeplyNestedDifferentiableFunctionWithMultipleGradCalls(self):
@polymorphic_function.function
def inner_fn(a, b):
return math_ops.add(a, b)
@polymorphic_function.function
def middle_fn(a, b):
return math_ops.mul(a, inner_fn(a, b))
@polymorphic_function.function
def outer_fn(x):
return middle_fn(x, 3.0)
x = constant_op.constant(5.0)
self.assertAllEqual(outer_fn(x), 5.0 * (5.0 + 3.0))
with backprop.GradientTape() as tp:
tp.watch(x)
result = outer_fn(x)
grad = tp.gradient(result, x)
self.assertAllEqual(grad, 2 * 5.0 + 3.0)
self.assertAllEqual(outer_fn(x), 5.0 * (5.0 + 3.0))
self.assertAllEqual(middle_fn(3.0, x), 3.0 * (3.0 + 5.0))
with backprop.GradientTape() as tp:
tp.watch(x)
result = outer_fn(x)
grad = tp.gradient(result, x)
self.assertAllEqual(grad, 2 * 5.0 + 3.0)
y = constant_op.constant(4.0)
with backprop.GradientTape() as tp:
tp.watch(y)
result = outer_fn(y)
grad = tp.gradient(result, y)
self.assertAllEqual(grad, 2 * 4.0 + 3.0)
with backprop.GradientTape() as tp:
tp.watch(y)
result = inner_fn(y, y)
grad = tp.gradient(result, y)
self.assertAllEqual(grad, 2.0)
@test_util.run_in_graph_and_eager_modes
def testDeeplyNestedDifferentiableFunctionGradientTapeInDefun(self):
@polymorphic_function.function
def inner_inner_fn(a, b):
return math_ops.add(a, b)
@polymorphic_function.function
def inner_fn(a, b):
return inner_inner_fn(a, b)
@polymorphic_function.function
def middle_fn(a, b):
return a * inner_fn(a, b)
@polymorphic_function.function
def outer_fn(x):
with backprop.GradientTape() as tp:
tp.watch(x)
result = middle_fn(x, 1.0)
grad = tp.gradient(result, x)
return grad
x = constant_op.constant(5.0)
grad = outer_fn(x)
self.assertAllEqual(grad, 2 * 5.0 + 1.0)
@test_util.run_in_graph_and_eager_modes
def testDeeplyNestedDifferentiableFunctionGradientTapeInNestedDefun(self):
@polymorphic_function.function
def inner_inner_fn(a, b):
return math_ops.add(a, b)
@polymorphic_function.function
def inner_fn(a, b):
return inner_inner_fn(a, b)
@polymorphic_function.function
def middle_fn(a, b):
return a * inner_fn(a, b)
@polymorphic_function.function
def almost_outer_fn(x):
with backprop.GradientTape() as tp:
tp.watch(x)
result = middle_fn(x, 1.0)
grad = tp.gradient(result, x)
return grad
@polymorphic_function.function
def outer_fn(x):
return almost_outer_fn(x)
x = constant_op.constant(5.0)
grad = outer_fn(x)
self.assertAllEqual(grad, 2 * 5.0 + 1.0)
@test_util.run_in_graph_and_eager_modes
def testDeeplyNestedDifferentiableFunctionGradientTapeInMultNestedDefun(self):
@polymorphic_function.function
def inner_inner_fn(a, b):
return math_ops.add(a, b)
@polymorphic_function.function
def inner_fn(a, b):
return inner_inner_fn(a, b)
@polymorphic_function.function
def middle_fn(a, b):
return a * inner_fn(a, b)
@polymorphic_function.function
def almost_outer_fn(x):
with backprop.GradientTape() as tp:
tp.watch(x)
result = middle_fn(x, 1.0)
grad = tp.gradient(result, x)
return grad
@polymorphic_function.function
def outer_fn(x):
return almost_outer_fn(x)
@polymorphic_function.function
def outer_outer_fn(x):
return outer_fn(x)
x = constant_op.constant(5.0)
grad = outer_outer_fn(x)
self.assertAllEqual(grad, 2 * 5.0 + 1.0)
@test_util.run_in_graph_and_eager_modes
def testDeeplyNestedDifferentiableFunctionTFGradientInDefun(self):
@polymorphic_function.function
def inner_inner_fn(a, b):
return math_ops.add(a, b)
@polymorphic_function.function
def inner_fn(a, b):
return inner_inner_fn(a, b)
@polymorphic_function.function
def middle_fn(a, b):
return a * inner_fn(a, b)
@polymorphic_function.function
def outer_fn(x):
result = middle_fn(x, 1.0)
return gradients_impl.gradients(result, [x])[0]
x = constant_op.constant(5.0)
grad = outer_fn(x)
self.assertAllEqual(grad, 2 * 5.0 + 1.0)
@test_util.run_in_graph_and_eager_modes
def testDeeplyNestedDifferentiableFunctionTFGradientInNestedDefun(self):
@polymorphic_function.function
def inner_inner_fn(a, b):
return math_ops.add(a, b)
@polymorphic_function.function
def inner_fn(a, b):
return inner_inner_fn(a, b)
@polymorphic_function.function
def middle_fn(a, b):
return a * inner_fn(a, b)
@polymorphic_function.function
def almost_outer_fn(x):
result = middle_fn(x, 1.0)
return gradients_impl.gradients(result, [x])[0]
@polymorphic_function.function
def outer_fn(x):
return almost_outer_fn(x)
x = constant_op.constant(5.0)
grad = outer_fn(x)
self.assertAllEqual(grad, 2 * 5.0 + 1.0)
@test_util.run_in_graph_and_eager_modes
def testDeeplyNestedDifferentiableFunctionTFGradientInMultNestedDefun(self):
@polymorphic_function.function
def inner_inner_fn(a, b):
return math_ops.add(a, b)
@polymorphic_function.function
def inner_fn(a, b):
return inner_inner_fn(a, b)
@polymorphic_function.function
def middle_fn(a, b):
return a * inner_fn(a, b)
@polymorphic_function.function
def almost_outer_fn(x):
result = middle_fn(x, 1.0)
return gradients_impl.gradients(result, [x])[0]
@polymorphic_function.function
def outer_fn(x):
return almost_outer_fn(x)
@polymorphic_function.function
def outer_outer_fn(x):
return outer_fn(x)
x = constant_op.constant(5.0)
grad = outer_outer_fn(x)
self.assertAllEqual(grad, 2 * 5.0 + 1.0)
def testDeeplyNestedDifferentiableFunctionWithVariable(self):
var = variables.Variable(constant_op.constant(1.0))
@polymorphic_function.function
def inner_fn(a, b):
return math_ops.add(a, b)
@polymorphic_function.function
def middle_fn(a, b):
return a * inner_fn(a, b)
@polymorphic_function.function
def outer_fn(x):
return middle_fn(x, var)
x = constant_op.constant(5.0)
with backprop.GradientTape() as tp:
tp.watch(x)
result = outer_fn(x)
grad = tp.gradient(result, x)
self.assertAllEqual(grad, 2 * 5.0 + 1.0)
def testDeeplyNestedDifferentiableFunctionWithVariableMultipleGradCalls(self):
v = variables.Variable(constant_op.constant(3.0))
@polymorphic_function.function
def inner_fn(a, b):
return math_ops.add(a, b)
@polymorphic_function.function
def middle_fn(a, b):
return math_ops.mul(a, inner_fn(a, b))
@polymorphic_function.function
def outer_fn(x):
return middle_fn(x, v)
x = constant_op.constant(5.0)
self.assertAllEqual(outer_fn(x), 5.0 * (5.0 + 3.0))
with backprop.GradientTape() as tp:
tp.watch(x)
result = outer_fn(x)
grad = tp.gradient(result, x)
self.assertAllEqual(grad, 2 * 5.0 + 3.0)
self.assertAllEqual(outer_fn(x), 5.0 * (5.0 + 3.0))
self.assertAllEqual(middle_fn(v, x), 3.0 * (3.0 + 5.0))
with backprop.GradientTape() as tp:
tp.watch(x)
result = outer_fn(x)
grad = tp.gradient(result, x)
self.assertAllEqual(grad, 2 * 5.0 + 3.0)
y = constant_op.constant(4.0)
with backprop.GradientTape() as tp:
tp.watch(y)
result = outer_fn(y)
grad = tp.gradient(result, y)
self.assertAllEqual(grad, 2 * 4.0 + 3.0)
v.assign(constant_op.constant(1.5))
with backprop.GradientTape() as tp:
tp.watch(y)
result = outer_fn(y)
grad = tp.gradient(result, y)
self.assertAllEqual(grad, 2 * 4.0 + 1.5)
with backprop.GradientTape() as tp:
tp.watch(y)
result = inner_fn(y, v)
grad = tp.gradient(result, y)
self.assertAllEqual(grad, 1.0)
def testDeeplyNestedDifferentiableFunctionWithVariableMultipleTFGrads(self):
with context.graph_mode(), self.cached_session():
v = resource_variable_ops.ResourceVariable(3.0)
v.initializer.run()
@polymorphic_function.function
def inner_fn(a, b):
return math_ops.add(a, b)
@polymorphic_function.function
def middle_fn(a, b):
return math_ops.mul(a, inner_fn(a, b))
@polymorphic_function.function
def outer_fn(x):
return middle_fn(x, v)
x = constant_op.constant(5.0)
self.assertAllEqual(outer_fn(x), 5.0 * (5.0 + 3.0))
grad, = gradients_impl.gradients(outer_fn(x), x)
self.assertAllEqual(grad, 2 * 5.0 + 3.0)
self.assertAllEqual(outer_fn(x), 5.0 * (5.0 + 3.0))
self.assertAllEqual(middle_fn(v, x), 3.0 * (3.0 + 5.0))
grad, = gradients_impl.gradients(outer_fn(x), x)
self.assertAllEqual(grad, 2 * 5.0 + 3.0)
y = constant_op.constant(4.0)
grad, = gradients_impl.gradients(outer_fn(y), y)
self.assertAllEqual(grad, 2 * 4.0 + 3.0)
self.evaluate(v.assign(constant_op.constant(1.5)))
grad, = gradients_impl.gradients(outer_fn(y), y)
self.assertAllEqual(grad, 2 * 4.0 + 1.5)
grad, = gradients_impl.gradients(inner_fn(y, v), y)
self.assertAllEqual(grad, 1.0)
def testNestedDifferentiableFunctionNoneOutputs(self):
@polymorphic_function.function
def foo(a, b):
return None, a * math_ops.add(a, b), None, 2*a
@polymorphic_function.function
def bar(x):
return foo(x, 1.0)
x = constant_op.constant(5.0)
with backprop.GradientTape(persistent=True) as tp:
tp.watch(x)
none1, r1, none2, r2 = bar(x)
g1 = tp.gradient(r1, x)
g2 = tp.gradient(r2, x)
self.assertAllEqual(r1, 30.0)
self.assertAllEqual(r2, 10.0)
self.assertIs(none1, None)
self.assertIs(none2, None)
self.assertAllEqual(g1, 2 * 5.0 + 1.0)
self.assertAllEqual(g2, 2.0)
def testGradientWithKeywordArguments(self):
matmul = polymorphic_function.function(math_ops.matmul)
def sq(x):
return matmul(a=x, b=x, transpose_a=True)
t = constant_op.constant([[1.0, 2.0], [3.0, 4.0]])
grad_t, = backprop.gradients_function(sq, [0])(t)
self.assertAllEqual(grad_t, [[6, 6], [14, 14]])
with backprop.GradientTape(persistent=True) as tape:
tape.watch(t)
one = matmul(t, b=t, transpose_a=True)
two = matmul(b=t, a=t, transpose_a=True)
three = matmul(a=t, b=t, transpose_a=True)
for output in [one, two, three]:
self.assertAllEqual(tape.gradient(output, t), [[6, 6], [14, 14]])
def testGradientInFunctionWithKeywordArguments(self):
@polymorphic_function.function
def f(x):
return backprop.gradients_function(lambda y: y * y, [0])(x)[0]
self.assertAllEqual(f(x=constant_op.constant(1.0)), 2.0)
def testFunctionHasNoSecondOrderGradient(self):
# This test needs nn_grad imported. We could just disable the lint error,
# but this way if the test is deleted we'll know the import isn't needed.
_ = nn_grad
v = variables.Variable(1.)
@polymorphic_function.function
def f(labels, logits):
return polymorphic_function.function(
nn_ops.sparse_softmax_cross_entropy_with_logits)(
labels=labels, logits=logits + v)
@polymorphic_function.function
def f_grad():
with backprop.GradientTape() as tape:
logits = constant_op.constant([1., 2.])
tape.watch(logits)
out = f(constant_op.constant(1), logits)
return tape.gradient(out, logits)
# Mainly we want to check that the function builds despite
# sparse_softmax_cross_entropy_with_logits not having a second-order
# gradient defined.
self.assertAllEqual([2], f_grad().shape)
if __name__ == '__main__':
ops.enable_eager_execution()
test.main()
tensorflow/python/eager/polymorphic_function/polymorphic_function.py
+1752
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/eager/polymorphic_function/polymorphic_function_test.py
+5053
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/eager/polymorphic_function/polymorphic_function_test_cpu_only.py
+48
View File
@@ -0,0 +1,48 @@
# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
from absl.testing import parameterized
from tensorflow.python.eager.polymorphic_function import polymorphic_function
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow.python.platform import test
class FunctionCpuOnlyTest(test.TestCase, parameterized.TestCase):
"""Test that jit_compile=True correctly throws an exception if XLA is not available.
This test should only be run without `--config=cuda`, as that implicitly links
in XLA JIT.
"""
def testJitCompileRaisesExceptionWhenXlaIsUnsupported(self):
if test.is_built_with_rocm() or test_util.is_xla_enabled():
return
with self.assertRaisesRegex(errors.UnimplementedError,
'support for that platform linked in'):
@polymorphic_function.function(jit_compile=True)
def fn(x):
return x + x
fn([1, 1, 2, 3])
if __name__ == '__main__':
ops.enable_eager_execution()
test.main()
tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_jit_test.py
+1262
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_test.py
+45
View File
@@ -0,0 +1,45 @@
# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
from tensorflow.compiler.tests import xla_test
from tensorflow.python.eager.polymorphic_function import polymorphic_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
class FunctionTests(xla_test.XLATestCase):
def testVarInitializedInFunction(self):
with self.test_scope():
v_holder = []
@polymorphic_function.function
def add_var(x):
if not v_holder:
v = variables.Variable([1., 2.])
v_holder.append(v)
already_initialized = variables.Variable(3.)
with ops.init_scope():
already_initialized.assign(10.)
v_holder.append(already_initialized)
return v_holder[0] + v_holder[1] + x
self.assertAllClose([13., 14.], add_var(constant_op.constant(2.)))
if __name__ == "__main__":
ops.enable_eager_execution()
test.main()
tensorflow/python/eager/polymorphic_function/saved_model_exported_concrete.py
+121
View File
@@ -0,0 +1,121 @@
# 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.
# ==============================================================================
# pylint: disable=unidiomatic-typecheck
"""ExportedConcreteFunction class and its associated functions.
Part of saved model utils, a shim layer for working with
functions exported/restored from saved models.
This functionality should ultimately be moved into a first-class core API.
"""
import gc
from tensorflow.python.eager.polymorphic_function import function_type_utils
from tensorflow.python.trackable import base as trackable
# TODO(kathywu): Delete this class when ConcreteFunctions can be copied with new
# captures.
class ExportedConcreteFunction(trackable.Trackable):
"""A callable class that uses captures from the exported SavedModel graph."""
__slots__ = ("function", "tensor_map")
def __init__(self, function, tensor_map):
self.function = function
self.tensor_map = tensor_map
def __call__(self, *args, **kwargs):
bound_arguments = function_type_utils.canonicalize_function_inputs(
args, kwargs, self.function._function_type
)
filtered_flat_args = self.function._function_type.unpack_inputs(
bound_arguments
)
export_captures = _map_captures_to_created_tensors(
self.function.graph.captures, self.tensor_map, self.function)
return self.function._call_flat(filtered_flat_args, export_captures)
def _map_captures_to_created_tensors(original_captures, tensor_map, function):
"""Maps eager tensors captured by a function to Graph resources for export.
Args:
original_captures: A dictionary mapping from tensors captured by the
function to interior placeholders for those tensors (inside the function
body).
tensor_map: A dictionary mapping from resource tensors owned by the eager
context to resource tensors in the exported graph.
function: Function with the original captures. Only used when raising the
AssertionError.
Returns:
A list of stand-in tensors which belong to the exported graph, corresponding
to the function's captures.
Raises:
AssertionError: If the function references a resource which is not part of
`tensor_map`.
"""
export_captures = []
for exterior, interior in original_captures:
mapped_resource = tensor_map.get(exterior, None)
if mapped_resource is None:
_raise_untracked_capture_error(function.name, exterior, interior)
export_captures.append(mapped_resource)
return export_captures
def _raise_untracked_capture_error(function_name, capture,
internal_capture=None,
node_path=None):
"""Raises AssertionError due to being unable to export a function."""
msg = ("Tried to export a function which references an 'untracked' resource. "
"TensorFlow objects (e.g. tf.Variable) captured by functions must be "
"'tracked' by assigning them to an attribute of a tracked object or "
"assigned to an attribute of the main object directly. See the "
"information below:"
f"\n\tFunction name = {function_name}")
if node_path is not None:
msg += f"\n\tPath to Function = {node_path}"
msg += f"\n\tCaptured Tensor = {capture}"
msg += f"\n\t{_get_trackable_parent_error_string(capture)}"
if internal_capture is not None:
msg += f"\n\tInternal Tensor = {internal_capture}"
raise AssertionError(msg)
def _get_trackable_parent_error_string(capture):
"""Gets error string with the capture's parent object."""
parent = getattr(capture, "_parent_trackable", None)
if parent is not None:
return f"Trackable referencing this tensor = {parent()}"
# Try to figure out where the resource came from by iterating over objects
# which reference it. This is slow and doesn't help us figure out how to
# match it to other objects when loading the SavedModel as a checkpoint,
# so we can't continue saving. But we can at least tell the user what
# needs attaching.
trackable_referrers = []
for primary_referrer in gc.get_referrers(capture):
if isinstance(primary_referrer, trackable.Trackable):
trackable_referrers.append(primary_referrer)
for secondary_referrer in gc.get_referrers(primary_referrer):
if isinstance(secondary_referrer, trackable.Trackable):
trackable_referrers.append(secondary_referrer)
return ("Trackable Python objects referring to this tensor "
"(from gc.get_referrers, limited to two hops) = [\n\t\t{}]"
.format("\n\t\t".join([repr(obj) for obj in trackable_referrers])))
tensorflow/python/eager/polymorphic_function/saved_model_utils.py
+82
View File
@@ -0,0 +1,82 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
# pylint: disable=unidiomatic-typecheck
"""A shim layer for working with functions exported/restored from saved models.
This functionality should ultimately be moved into a first-class core API.
"""
import numpy
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_util
from tensorflow.python.saved_model import registration
from tensorflow.python.trackable import base as trackable
@registration.register_tf_serializable()
class TrackableConstant(trackable.Trackable):
"""Trackable class for captured constants."""
__slots__ = ("capture", "function", "_exported_tensor")
def __init__(self, capture, function):
self.capture = capture
self.function = function
self._exported_tensor = None
def _export_to_saved_model_graph(self, tensor_map, **unused_kwargs):
capture_constant_value = tensor_util.constant_value(self.capture)
if capture_constant_value is None:
raise ValueError(
f"Unable to save function {self.function.name} because it "
f"captures graph tensor {self.capture} from a parent function which "
"cannot be converted to a constant with `tf.get_static_value`.")
if numpy.prod(self.capture.shape.as_list()) > 1 and numpy.all(
capture_constant_value == capture_constant_value.flat[0]):
# For the common case of a constant array filled with the same
# value, rebuild the constant op specifically with the shape arg,
# since otherwise the whole array is written into the node def,
# causing performance and graph proto size issues (protos cannot be
# bigger than 2GB).
copied_tensor = constant_op.constant(
capture_constant_value.flat[0],
dtype=self.capture.dtype,
shape=self.capture.shape)
else:
copied_tensor = constant_op.constant(capture_constant_value)
tensor_map[self.capture] = copied_tensor
self._exported_tensor = copied_tensor
return [self.capture]
def _serialize_to_proto(self, object_proto=None, **kwargs):
object_proto.constant.operation = self._exported_tensor.op.name
@classmethod
def _deserialize_from_proto(cls, object_proto, operation_attributes,
**kwargs):
tensor_proto = (
operation_attributes[object_proto.constant.operation]["value"].tensor)
ndarray = tensor_util.MakeNdarray(tensor_proto)
if dtypes.as_dtype(tensor_proto.dtype) == dtypes.string:
with ops.device("CPU"):
# String operations should be done on the CPU.
imported_constant = constant_op.constant(ndarray)
else:
imported_constant = constant_op.constant(ndarray)
return imported_constant
tensorflow/python/eager/polymorphic_function/tf_method_target.py
+51
View File
@@ -0,0 +1,51 @@
# 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.
# ==============================================================================
"""Module for the TFMethodTarget Class."""
import weakref
from tensorflow.python.util import tf_inspect
# When a method is bound to objects of this type, it allows AutoGraph to
# recover a weak reference the original method's self pointer, so that it can
# execute it consistent with class_method_to_instance_method's
# bound_method_wrapper.
# TODO(b/119246461): This is not pretty. Use a descriptor instead?
class TfMethodTarget:
"""Binding target for methods replaced by function and defun."""
__slots__ = ("weakrefself_target__", "weakrefself_func__")
def __init__(self, target, original_python_function):
self.weakrefself_target__ = target
self.weakrefself_func__ = weakref.ref(original_python_function)
@property
def target(self):
return self.weakrefself_target__()
@property
def target_class(self):
true_self = self.weakrefself_target__()
if tf_inspect.isclass(true_self):
# Class method
return true_self
else:
return true_self.__class__
def call(self, args, kwargs):
wrapped_fn = self.weakrefself_func__()
return wrapped_fn(self.weakrefself_target__(), *args, **kwargs)
tensorflow/python/eager/polymorphic_function/tracing_compilation.py
+379
View File
@@ -0,0 +1,379 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Compile Python functions to TF graphs using tracing."""
import contextlib
import dataclasses
import enum
import threading
from typing import Any, Callable, Dict, Optional, Tuple
from tensorflow.core.function import trace_type
from tensorflow.core.function.capture import capture_container
from tensorflow.core.function.polymorphism import function_cache as function_cache_lib
from tensorflow.core.function.polymorphism import function_type as function_type_lib
from tensorflow.python.autograph.core import ag_ctx
from tensorflow.python.eager import monitoring
from tensorflow.python.eager.polymorphic_function import attributes as attributes_lib
from tensorflow.python.eager.polymorphic_function import concrete_function as concrete_function_lib
from tensorflow.python.eager.polymorphic_function import function_context
from tensorflow.python.eager.polymorphic_function import function_type_utils
from tensorflow.python.eager.polymorphic_function import transform
from tensorflow.python.framework import func_graph as func_graph_module
from tensorflow.python.framework import ops
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.profiler import trace
from tensorflow.python.util import compat
_graph_building_time_counter = monitoring.Counter(
"/tensorflow/core/tf_function/graph_building_time_usecs",
"Time for tf.function to build a graph (us).",
)
class ScopeType(enum.Enum):
"""Enumerate scopes under which functions might be traced."""
NO_SCOPE = 1
VARIABLE_CREATION = 2
NO_VARIABLE_CREATION = 3
@dataclasses.dataclass
class TracingOptions:
"""Configuration options for tracing."""
# Python function to trace.
python_function: Callable[[Any], Any] = lambda *args, **kwargs: None
# Name given to the traced function.
name: str = "function"
# Known FunctionType of the python function.
polymorphic_type: Optional[function_type_lib.FunctionType] = None
# Known default values for the python function parameters.
default_values: Optional[Dict[str, Any]] = None
# Identifies effecting scope under which the function is traced.
scope_type: ScopeType = ScopeType.NO_SCOPE
# FunctionDef attributes for traced function.
attributes: Optional[Dict[str, Any]] = None
# See https://www.tensorflow.org/guide/autograph for more information.
# If autograph is enabled.
autograph: bool = True
# Optional tuple of `tf.autograph.experimental.Feature` values.
autograph_options: Optional[Tuple[Any, ...]] = None
# Trace generalized functions where possible to avoid future retracing.
reduce_retracing: bool = False
# If true, graph of generated Function will be destroyed with the function.
bind_graph_to_function: bool = False
# A FunctionCache object that holds existing traced functions.
function_cache: Optional[function_cache_lib.FunctionCache] = None
# A FunctionCaptures object that tracks by-ref captures.
function_captures: Optional[capture_container.FunctionCaptures] = None
# If specified, guards tracing and function lookup
lock: Optional[threading.Lock] = None
def __post_init__(self):
if self.attributes:
for attribute in self.attributes:
if attribute not in attributes_lib.TRACING_COMPILATION_ALLOWLIST:
raise ValueError(
f"Tracing compilation does not support `{attribute}` as an"
" attribute."
)
if not self.polymorphic_type or self.default_values is None:
self.polymorphic_type = function_type_lib.FunctionType.from_callable(
self.python_function
)
self.default_values = function_type_lib.FunctionType.get_default_values(
self.python_function
)
self._input_signature = function_type_utils.to_input_signature(
self.polymorphic_type
)
@property
def is_pure(self):
return self.attributes and attributes_lib.IMPLEMENTS in self.attributes
@property
def input_signature(self):
return self._input_signature
def call_function(args=None, kwargs=None, tracing_options=None):
"""Traces a function for args and kwargs and calls it after."""
if not tracing_options:
tracing_options = TracingOptions()
args = args if args else ()
kwargs = kwargs if kwargs else {}
function = trace_function(
args=args, kwargs=kwargs, tracing_options=tracing_options
)
# Bind it ourselves to skip unnecessary canonicalization of default call.
bound_args = function.function_type.bind(*args, **kwargs)
flat_inputs = function.function_type.unpack_inputs(bound_args)
return function._call_flat( # pylint: disable=protected-access
flat_inputs, captured_inputs=function.captured_inputs
)
def trace_function(args=None, kwargs=None, tracing_options=None):
"""Returns a `ConcreteFunction` specialized to inputs and execution context.
Compiles a Graph corresponding to the Python function logic and uses that
to generate a differentiable ConcreteFunction.
Args:
args: inputs to specialize on. Can be concrete values (e.g. 1) or
`tf.Tensor` or `tf.TensorSpec`.
kwargs: keyword inputs to specialize on. Concrete values (e.g. 1) or
`tf.Tensor` or `tf.TensorSpec`.
tracing_options: TracingOptions for the tracing process.
"""
if not tracing_options:
tracing_options = TracingOptions()
args = args if args else ()
kwargs = kwargs if kwargs else {}
if tracing_options.input_signature and (args or kwargs):
# Check to see if a valid type can be generated from the args, kwargs
bound_args = function_type_utils.bind_function_inputs(
args,
kwargs,
tracing_options.polymorphic_type,
tracing_options.default_values,
)
args, kwargs = bound_args.args, bound_args.kwargs
with tracing_options.lock or contextlib.nullcontext():
if tracing_options.input_signature and not args and not kwargs:
args = tracing_options.input_signature
kwargs = {}
concrete_function = _maybe_define_function(
args, kwargs, tracing_options
)
if not tracing_options.bind_graph_to_function:
concrete_function._garbage_collector.release() # pylint: disable=protected-access
return concrete_function
def _maybe_define_function(args, kwargs, tracing_options):
"""Gets a function for these inputs, defining it if necessary.
Args:
args: The varargs for the Python function.
kwargs: The keyword args for the Python function.
tracing_options: TracingOptions for the tracing process.
Returns:
A ConcreteFunction generated based on args, kwargs and tracing_options.
Raises:
ValueError: If inputs are incompatible with the input signature.
TypeError: If the function inputs include non-hashable objects
RuntimeError: If there's an internal bug (inconsistency) in handling
shape relaxation retracing.
"""
bound_args = function_type_utils.canonicalize_function_inputs(
args,
kwargs,
tracing_options.polymorphic_type,
tracing_options.default_values,
tracing_options.is_pure,
)
args, kwargs = bound_args.args, bound_args.kwargs
if tracing_options.input_signature is not None:
args = (
*tracing_options.input_signature,
*args[len(tracing_options.input_signature) :],
)
current_func_context = function_context.make_function_context(
tracing_options.scope_type
)
capture_types = (
tracing_options.function_captures.capture_types
if tracing_options.function_captures
else {}
)
lookup_func_type, lookup_func_context = (
function_type_utils.make_canonicalized_monomorphic_type(
args,
kwargs,
capture_types,
tracing_options.polymorphic_type,
)
)
if tracing_options.function_cache is not None:
concrete_function = tracing_options.function_cache.lookup(
lookup_func_type, current_func_context
)
else:
concrete_function = None
if concrete_function is not None:
return concrete_function
# Use a timer for graph building only if not already inside a function. This
# avoids double counting graph building time for nested functions.
with monitoring.MonitoredTimer(
_graph_building_time_counter.get_cell()
) if not ops.inside_function() else contextlib.nullcontext():
with trace.Trace("tf.function-graph_building"):
logging.vlog(
1,
"Creating new FuncGraph for Python function %r (key: %r, %r)",
tracing_options.python_function,
current_func_context,
lookup_func_type,
)
logging.vlog(
2, "Python function signature [args: %s] [kwargs: %s]", args, kwargs
)
ag_status = (
ag_ctx.Status.ENABLED
if tracing_options.autograph
else ag_ctx.Status.DISABLED
)
with ag_ctx.ControlStatusCtx(
status=ag_status, options=tracing_options.autograph_options
):
func_graph = func_graph_module.FuncGraph(tracing_options.name)
if (
tracing_options.input_signature is None
and tracing_options.reduce_retracing
and tracing_options.function_cache
):
target_func_type = tracing_options.function_cache.generalize(
current_func_context, lookup_func_type
)
else:
target_func_type = lookup_func_type
concrete_function = _create_concrete_function(
target_func_type, lookup_func_context, func_graph, tracing_options
)
if tracing_options.function_cache is not None:
tracing_options.function_cache.add(
concrete_function, current_func_context
)
return concrete_function
def _create_concrete_function(
function_type, type_context, func_graph, tracing_options
):
"""Create a `ConcreteFunction` from `args`, `kwargs`, and `func_graph`."""
placeholder_context = trace_type.InternalPlaceholderContext(
func_graph, type_context.get_placeholder_mapping()
)
with func_graph.as_default():
placeholder_bound_args = function_type.placeholder_arguments(
placeholder_context
)
disable_acd = tracing_options.attributes and tracing_options.attributes.get(
attributes_lib.DISABLE_ACD, False
)
traced_func_graph = func_graph_module.func_graph_from_py_func(
tracing_options.name,
tracing_options.python_function,
placeholder_bound_args.args,
placeholder_bound_args.kwargs,
None,
func_graph=func_graph,
add_control_dependencies=not disable_acd,
arg_names=function_type_utils.to_arg_names(function_type),
create_placeholders=False,
)
transform.apply_func_graph_transforms(traced_func_graph)
graph_capture_container = traced_func_graph.function_captures
if tracing_options.function_captures:
# Maintain the list of all captures
tracing_options.function_captures.merge_by_ref_with(graph_capture_container)
# Create a new FunctionType including captures and outputs.
output_type = trace_type.from_value(
traced_func_graph.structured_outputs, type_context
)
traced_func_type = function_type_lib.FunctionType(
function_type.parameters.values(),
traced_func_graph.function_captures.capture_types,
return_annotation=output_type,
)
concrete_function = concrete_function_lib.ConcreteFunction.from_func_graph(
traced_func_graph,
traced_func_type,
tracing_options.attributes,
# Tell the ConcreteFunction to clean up its graph once it goes out of
# scope. This is not the default behavior since it gets used in some
# places (like Keras) where the FuncGraph lives longer than the
# ConcreteFunction.
shared_func_graph=False,
)
_set_arg_keywords(concrete_function)
transform.call_concrete_function_callbacks(concrete_function)
return concrete_function
def _set_arg_keywords(concrete_function):
"""Sets arg keywords for ConcreteFunction."""
seen_names = set()
concrete_function._arg_keywords = [] # pylint: disable=protected-access
prefix_counts = {}
graph = concrete_function.graph
num_captures = len(graph.internal_captures + graph.deferred_internal_captures)
num_positional = len(graph.inputs) - num_captures
for arg in concrete_function.graph.inputs[:num_positional]:
try:
user_arg_name = compat.as_str(arg.op.get_attr("_user_specified_name"))
except ValueError:
user_arg_name = "tensor_arg"
proposal = user_arg_name
while proposal in seen_names:
index = prefix_counts.get(user_arg_name, 1)
proposal = "{}_{}".format(user_arg_name, index)
prefix_counts[user_arg_name] = index + 1
seen_names.add(proposal)
concrete_function._arg_keywords.append(proposal) # pylint: disable=protected-access
# Anything can be a positional argument, in the same order as .inputs
concrete_function._num_positional_args = ( # pylint: disable=protected-access
num_positional
)
tensorflow/python/eager/polymorphic_function/tracing_compilation_test.py
+2277
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/eager/polymorphic_function/transform.py
+32
View File
@@ -0,0 +1,32 @@
# 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.
# ==============================================================================
"""tf.function transformations implementation."""
# TODO(fmuham): Move this logic to core/function when layered.
# TODO(fmuham): Deprecate and migrate these as AtomicFunction transformations.
FUNC_GRAPH_TRANSFORMS = []
CONCRETE_FUNCTION_CALLBACKS = []
def apply_func_graph_transforms(func_graph):
"""Applies registered transformations to FuncGraph."""
for transform in FUNC_GRAPH_TRANSFORMS:
transform(func_graph)
def call_concrete_function_callbacks(concrete_fn):
"""Calls registered callbacks against new ConcreteFunctions."""
for callback in CONCRETE_FUNCTION_CALLBACKS:
callback(concrete_fn)