chore: import upstream snapshot with attribution

tensorflow/core/config/BUILD

@@ -0,0 +1,111 @@
+load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
+load("//tensorflow:tensorflow.bzl", "py_test", "tf_cc_test")
+load("//tensorflow:tensorflow.default.bzl", "filegroup")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+exports_files(
+    ["flags_api_wrapper.cc"],
+    visibility = ["//tensorflow:internal"],
+)
+
+cc_library(
+    name = "flags_headers",
+    textual_hdrs = [
+        "flag_defs.h",
+        "flags.h",
+    ],
+    visibility = ["//tensorflow:internal"],
+    deps = [
+        "//tensorflow/core/platform:stringpiece",
+        "@com_google_absl//absl/strings:string_view",
+    ],
+)
+
+filegroup(
+    name = "flags_headers_filegroup",
+    srcs = [
+        "flag_defs.h",
+        "flags.h",
+    ],
+    visibility = ["//tensorflow:internal"],
+)
+
+filegroup(
+    name = "mobile_srcs_only_runtime",
+    srcs = [
+        "flag_defs.cc",
+        "flags.cc",
+    ],
+    visibility = ["//tensorflow:internal"],
+)
+
+filegroup(
+    name = "mobile_hdrs_only_runtime",
+    srcs = [
+        "flag_defs.h",
+        "flags.h",
+    ],
+    visibility = ["//tensorflow:internal"],
+)
+
+cc_library(
+    name = "flag_defs",
+    srcs = ["flag_defs.cc"],
+    hdrs = ["flag_defs.h"],
+    visibility = ["//tensorflow:internal"],
+    deps = [
+        ":flags",
+    ],
+)
+
+cc_library(
+    name = "flags",
+    srcs = ["flags.cc"],
+    hdrs = ["flags.h"],
+    visibility = ["//tensorflow:internal"],
+    deps = [
+        "//tensorflow/core/platform:stringpiece",
+        "//tensorflow/core/util:env_var",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+    ],
+)
+
+tf_cc_test(
+    name = "flags_test",
+    srcs = ["flags_test.cc"],
+    deps = [
+        ":flag_defs",
+        ":flags",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+    ],
+)
+
+pytype_strict_library(
+    name = "flags_py",
+    srcs = ["flags.py"],
+    visibility = ["//tensorflow:internal"],
+    deps = [
+        "//tensorflow/python:flags_pybind",
+        "//tensorflow/python:pywrap_tensorflow",
+    ],
+)
+
+py_test(
+    name = "flags_py_test",
+    size = "small",
+    srcs = ["flags_test.py"],
+    main = "flags_test.py",
+    strict_deps = True,
+    deps = [
+        ":flags_py",
+        #internal proto upb dep
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)

tensorflow/core/config/flag_defs.cc

@@ -0,0 +1,30 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/core/config/flag_defs.h"
+
+namespace tensorflow {
+namespace flags {
+
+// Note: this is subject to change. Currently the flags are stored as part
+// of a global object. In the future the container might be moved to a different
+// location.
+Flags& Global() {
+  static Flags* registry = new Flags();
+  return *registry;
+}
+
+}  // namespace flags
+}  // namespace tensorflow

tensorflow/core/config/flag_defs.h

@@ -0,0 +1,83 @@
+/* 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_CONFIG_FLAG_DEFS_H_
+#define TENSORFLOW_CORE_CONFIG_FLAG_DEFS_H_
+
+#include "tensorflow/core/config/flags.h"
+
+namespace tensorflow {
+namespace flags {
+
+class Flags {
+ public:
+  // Test only flags. See flags_test.cc for example usage.
+  TF_DECLARE_FLAG(test_only_experiment_1, true, "Test only experiment 1.");
+  TF_DECLARE_FLAG(test_only_experiment_2, false, "Test only experiment 2.");
+
+  // Declare flags below here.
+  // LINT.IfChange
+  TF_DECLARE_FLAG(enable_nested_function_shape_inference, false,
+                  "Allow ops such as tf.cond to invoke the ShapeRefiner on "
+                  "their nested functions.");
+  TF_DECLARE_FLAG(enable_quantized_dtypes_training, false,
+                  "Set quantized dtypes, like tf.qint8, to be trainable.");
+  TF_DECLARE_FLAG(graph_building_optimization, false,
+                  "Optimize graph building for faster tf.function tracing.");
+  TF_DECLARE_FLAG(
+      op_building_optimization, true,
+      "Optimize tf.Operation building for faster tf.function tracing.");
+  TF_DECLARE_FLAG(saved_model_fingerprinting, true,
+                  "Add fingerprint to SavedModels.");
+  TF_DECLARE_FLAG(
+      tf_shape_default_int64, false,
+      "The default output of tf.shape (i.e. when out_type is not specified) is "
+      "int64 when this flag is true and int32 otherwise. Setting this to true "
+      "is an unsupported, experimental setting that causes known breakages.");
+  TF_DECLARE_FLAG(more_stack_traces, false,
+                  "Enable experimental code that preserves and propagates "
+                  "graph node stack traces in C++.");
+  TF_DECLARE_FLAG(publish_function_graphs, true,
+                  "Enables the publication of partitioned function graphs "
+                  "via StatsPublisherInterface. Disabling this flag can "
+                  "reduce memory consumption.");
+  TF_DECLARE_FLAG(enable_aggressive_constant_replication, true,
+                  "Replicate constants across CPU devices and even for local "
+                  "CPUs within the same task if available.")
+  TF_DECLARE_FLAG(enable_colocation_key_propagation_in_while_op_lowering, false,
+                  "If true, colocation key attributes for the ops will be "
+                  "propagated during while op lowering to switch/merge ops.")
+  TF_DECLARE_FLAG(enable_tf2min_ici_weight, false,
+                  "If true, ici weight optimization will be used in tf2/min.")
+  // TODO(b/341325107): Make this behavior the default and remove the flag.
+  TF_DECLARE_FLAG(enable_function_pruning_before_inlining, false,
+                  "If true, functions will be pruned before inlining.")
+  TF_DECLARE_FLAG(enable_skip_encapsulation_for_non_tpu_graphs, false,
+                  "If true, TF2XLA encapsulation will be skipped for non-TPU "
+                  "graphs.")
+  TF_DECLARE_FLAG(enable_graph_debug_info_caching_for_stack_frames, true,
+                  "If true, graph debug info will cache the stack frames.")
+  TF_DECLARE_FLAG(
+      enable_fatal_error_on_collective_abort, false,
+      "If true, a fatal error will be raised when a collective is aborted.")
+  // LINT.ThenChange(//tensorflow/core/config/flags_api_wrapper.cc)
+};
+
+Flags& Global();
+
+}  // namespace flags
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_CONFIG_FLAG_DEFS_H_

tensorflow/core/config/flags.cc

@@ -0,0 +1,37 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/core/config/flags.h"
+
+#include "absl/strings/ascii.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/core/platform/stringpiece.h"
+#include "tensorflow/core/util/env_var.h"
+
+namespace tensorflow {
+namespace config {
+
+Flag::Flag(absl::string_view flag, bool default_value) {
+  bool val = default_value;
+  if (ReadBoolFromEnvVar(absl::AsciiStrToUpper(flag), default_value, &val)
+          .ok()) {
+    value_ = val;
+    return;
+  }
+  value_ = default_value;
+}
+
+}  // namespace config
+}  // namespace tensorflow

tensorflow/core/config/flags.h

@@ -0,0 +1,48 @@
+/* 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_CONFIG_FLAGS_H_
+#define TENSORFLOW_CORE_CONFIG_FLAGS_H_
+
+#include "absl/strings/string_view.h"
+#include "tensorflow/core/platform/stringpiece.h"
+
+namespace tensorflow {
+namespace config {
+
+// Container class for a single feature flag.
+// Note: this class is not thread safe.
+class Flag {
+ public:
+  explicit Flag(absl::string_view flag_name, bool default_value);
+  bool value() { return value_; }
+  void reset(bool value) { value_ = value; }
+
+ private:
+  bool value_;
+};
+
+// Macro to declare new flags. Declare all flags in core/config/flag_defs.h
+// These flags can be overridden by setting the associated environment variable
+// TF_FLAG_* flag to true or false. E.g. setting TF_FLAG_MY_FLAG=false will
+// override the default value for a flag named `my_flag` to false.
+#define TF_DECLARE_FLAG(flag_name, default_value, doc) \
+  ::tensorflow::config::Flag flag_name =               \
+      ::tensorflow::config::Flag("TF_FLAG_" #flag_name, default_value);
+
+}  // namespace config
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_CONFIG_FLAGS_H_

tensorflow/core/config/flags.py

@@ -0,0 +1,29 @@
+# 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.
+# ==============================================================================
+"""Code to enable/disable feature flags."""
+
+# pylint: disable=invalid-import-order, g-bad-import-order, unused-import
+from tensorflow.python import pywrap_tensorflow
+from tensorflow.python import flags_pybind
+
+
+FLAGS = None
+
+
+def config():
+  global FLAGS
+  if FLAGS is None:
+    FLAGS = flags_pybind.Flags()
+  return FLAGS

tensorflow/core/config/flags_api_wrapper.cc

@@ -0,0 +1,61 @@
+/* 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.
+==============================================================================*/
+
+#include <memory>
+
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/pytypes.h"  // from @pybind11
+#include "pybind11/stl.h"  // from @pybind11
+#include "tensorflow/core/config/flag_defs.h"
+#include "tensorflow/core/config/flags.h"
+
+namespace py = pybind11;
+namespace flags = tensorflow::flags;
+
+using tensorflow::config::Flag;
+using tensorflow::flags::Flags;
+
+// Macro to expose the declared flag to python.
+#define TF_PY_DECLARE_FLAG(flag_name) \
+  flags.def_readwrite(#flag_name, &Flags::flag_name);
+
+PYBIND11_MODULE(flags_pybind, m) {
+  py::class_<Flag>(m, "Flag")
+      .def("value", &Flag::value)
+      .def("reset", &Flag::reset);
+
+  py::class_<Flags, std::unique_ptr<Flags, py::nodelete>> flags(m, "Flags");
+  flags.def(py::init(
+      []() { return std::unique_ptr<Flags, py::nodelete>(&flags::Global()); }));
+  // LINT.IfChange
+  TF_PY_DECLARE_FLAG(test_only_experiment_1);
+  TF_PY_DECLARE_FLAG(test_only_experiment_2);
+  TF_PY_DECLARE_FLAG(enable_nested_function_shape_inference);
+  TF_PY_DECLARE_FLAG(enable_quantized_dtypes_training);
+  TF_PY_DECLARE_FLAG(graph_building_optimization);
+  TF_PY_DECLARE_FLAG(op_building_optimization);
+  TF_PY_DECLARE_FLAG(saved_model_fingerprinting);
+  TF_PY_DECLARE_FLAG(tf_shape_default_int64);
+  TF_PY_DECLARE_FLAG(more_stack_traces);
+  TF_PY_DECLARE_FLAG(publish_function_graphs);
+  TF_PY_DECLARE_FLAG(enable_aggressive_constant_replication);
+  TF_PY_DECLARE_FLAG(enable_colocation_key_propagation_in_while_op_lowering);
+  TF_PY_DECLARE_FLAG(enable_tf2min_ici_weight)
+  TF_PY_DECLARE_FLAG(enable_function_pruning_before_inlining)
+  TF_PY_DECLARE_FLAG(enable_skip_encapsulation_for_non_tpu_graphs)
+  TF_PY_DECLARE_FLAG(enable_graph_debug_info_caching_for_stack_frames)
+  TF_PY_DECLARE_FLAG(enable_fatal_error_on_collective_abort)
+  // LINT.ThenChange(//tensorflow/core/config/flag_defs.h)
+};

tensorflow/core/config/flags_test.cc

@@ -0,0 +1,36 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/core/config/flags.h"
+
+#include "tensorflow/core/config/flag_defs.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace {
+
+TEST(TFFlags, ReadFlagValue) {
+  EXPECT_TRUE(flags::Global().test_only_experiment_1.value());
+  EXPECT_FALSE(flags::Global().test_only_experiment_2.value());
+}
+
+TEST(TFFlags, ResetFlagValue) {
+  EXPECT_TRUE(flags::Global().test_only_experiment_1.value());
+  flags::Global().test_only_experiment_1.reset(false);
+  EXPECT_FALSE(flags::Global().test_only_experiment_1.value());
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/core/config/flags_test.py

@@ -0,0 +1,47 @@
+# 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 flags."""
+
+from tensorflow.core.config import flags
+from tensorflow.python.platform import test
+
+
+class FlagsTest(test.TestCase):
+
+  def test_experiment_flag(self):
+    self.assertTrue(flags.config().test_only_experiment_1.value())
+    self.assertFalse(flags.config().test_only_experiment_2.value())
+
+    flags.config().test_only_experiment_1.reset(False)
+    flags.config().test_only_experiment_2.reset(True)
+
+    self.assertFalse(flags.config().test_only_experiment_1.value())
+    self.assertTrue(flags.config().test_only_experiment_2.value())
+
+  def test_flags_singleton(self):
+    flags.config().test_only_experiment_1.reset(False)
+    self.assertFalse(flags.config().test_only_experiment_1.value())
+
+    # Get second reference to underlying Flags singleton.
+    flag = flags.flags_pybind.Flags()
+    flag.test_only_experiment_1.reset(True)
+
+    # check that both references are correctly updated.
+    self.assertTrue(flags.config().test_only_experiment_1.value())
+    self.assertTrue(flag.test_only_experiment_1.value())
+
+
+if __name__ == '__main__':
+  test.main()