Enable the KerasTensors refactoring of the Keras functional API. Replaces symbolic tensors during Functional API construction with lightweight tensor-like objects that have a very similar API, but do not require a global graph workspace.

This should improve the reliability and performance of the Keras functional API and of automatic TF op -> Lambda layer conversions during functional API construction.
E.g. ~8-10% faster functional model construction time, dramatically lower memory usage

If this causes a breakage in your code, you are likely experiencing one of the following:
* code that uses map_fn/tf.cond/tf.while_loop/control flow as op layers and happens to maybe work right now (I wouldn't trust any model that uses those during functional model construction)
* code that uses get_concrete_function to trace keras inputs directly.
* code that uses isinstance(x, tf.Tensor) instead of tf.is_tensor
* Any code already susceptible to leaking tensors outside of a graph becomes slightly more susceptible now (though it would be an issue anyway)
* Any code that is overly dependent on the exact names attached to symbolic tensors (e.g. assumes there will be `:0` at the end of the inputs, treating names as unique identifiers instead of using .ref(), etc.)
* Your code relies on the exact # and names of the op layers not changing from what it was
* You have code that has very tricky shape manipulation via automatically converted tf op layers, and the KerasTensor shape inference is insufficient.
* You have code that tries manually walking a model layer by layer and assumes layers only ever have 1 positional argument (This doesn't hold true in head either, but it becomes marginally more likely to cause issues w/ the newer op lambda layers)
* Your code manually enters keras.backend.get_graph() before building your functional model (Usually this was done to work around fragility in the functional api/ op -> layer conversions. This should no longer be needed, and in fact will now cause tensors to leak out of a graph if you do it)
* direct asserts might break in case where an op like tf.rank used to return a static or symbolic value depending on if the input had a fully static shape or not. Now it is always symbolic.

To see if this refactoring is causing you problems, you can disable the refactoring as follows for as long as we have an internal flag available:

```
from tensorflow.python.keras.engine import keras_tensor
keras_tensor.disable_keras_tensors()
```

PiperOrigin-RevId: 324032210
Change-Id: Ib5e33030f6fc62a84870078046e4f1acdf800f95

RELEASE.md

@@ -14,6 +14,18 @@
 * Removed `tf.distribute.Strategy.experimental_run_v2` method, which was deprecated in TF 2.2.
 * `tensorflow.python`, `tensorflow.core` and `tensorflow.compiler` modules are
     now hidden. These modules are not part of TensorFlow public API.
+* A major refactoring of the internals of the Keras Functional API may affect code that is relying on certain internal details:
+    * Code that uses `isinstance(x, tf.Tensor)` instead of `tf.is_tensor` when checking Keras symbolic inputs/outputs should switch to using `tf.is_tensor`.
+    * Code that is overly dependent on the exact names attached to symbolic tensors (e.g. assumes there will be ":0" at the end of the inputs, treats names as unique identifiers instead of using `tensor.ref()`, etc.)
+    * Code that uses `get_concrete_function` to trace Keras symbolic inputs directly should switch to building matching `tf.TensorSpec`s directly and tracing the `TensorSpec` objects.
+    * Code that relies on the exact number and names of the op layers that TensorFlow operations were converted into. These may have changed.
+    * Code that uses `tf.map_fn`/`tf.cond`/`tf.while_loop`/control flow as op layers and happens to work before TF 2.4. These will explicitly be unsupported now. Converting these ops to Functional API op layers was unreliable before TF 2.4, and prone to erroring incomprehensibly or being silently buggy.
+    * Code that directly asserts on a Keras symbolic value in cases where ops like `tf.rank` used to return a static or symbolic value depending on if the input had a fully static shape or not. Now these ops always return symbolic values.
+    * Code already susceptible to leaking tensors outside of graphs becomes slightly more likely to do so now.
+    * Code that requires very tricky shape manipulation via converted op layers in order to work, where the Keras symbolic shape inference proves insufficient.
+    * Code that tries manually walking a `tf.keras.Model` layer by layer and assumes layers only ever have one positional argument. This assumption doesn't hold true before TF 2.4 either, but is more likely to cause issues know.
+    * Code that manually enters `keras.backend.get_graph()` before building a functional model. This is no longer needed.
+
 
 ## Known Caveats
 
@@ -24,6 +36,7 @@
 * <INSERT MAJOR FEATURE HERE, USING MARKDOWN SYNTAX>
 * <IF RELEASE CONTAINS MULTIPLE FEATURES FROM SAME AREA, GROUP THEM TOGETHER>
 * A new module named `tf.experimental.numpy` is added, which is a NumPy-compatible API for writing TF programs. This module provides class `ndarray`, which mimics the `ndarray` class in NumPy, and wraps an immutable `tf.Tensor` under the hood. A subset of NumPy functions (e.g. `numpy.add`) are provided. Their inter-operation with TF facilities is seamless in most cases. See tensorflow/python/ops/numpy_ops/README.md for details of what are supported and what are the differences with NumPy.
+* A major refactoring of the internals of the Keras Functional API has been completed, that should improve the reliability, stability, and performance of constructing Functional models.
 
 ## Bug Fixes and Other Changes
 
@@ -65,7 +78,13 @@
       called, and independent of global seed settings.
 *   `tf.distribute`:
     * <ADD RELEASE NOTES HERE>
-*   `tf.keras`:
+* `tf.keras`:
+    * Improvements from the functional API refactoring:
+      * Functional model construction does not need to maintain a global workspace graph, removing memory leaks especially when building many models or very large models.
+      * Functional model construction should be ~8-10% faster on average.
+      * Functional models can now contain non-symbolic values in their call inputs inside of the first positional argument.
+      * Several classes of TF ops that were not reliably converted to Keras layers during functional API construction should now work, e.g. `tf.image.ssim_multiscale`
+      * Error messages when Functional API construction goes wrong (and when ops cannot be converted to Keras layers automatically) should be clearer and easier to understand.
     * `Optimizer.minimize` can now accept a loss `Tensor` and a `GradientTape`
       as an alternative to accepting a `callable` loss.
 * `tf.function` / AutoGraph:

tensorflow/python/keras/backend.py

@@ -1181,7 +1181,7 @@ def placeholder(shape=None,
 
   >>> input_ph = tf.keras.backend.placeholder(shape=(2, 4, 5))
   >>> input_ph
-  <tf.Tensor 'Placeholder_...' shape=(2, 4, 5) dtype=float32>
+  <KerasTensor: shape=(2, 4, 5) dtype=float32 (Symbolic value ...)>
 
   """
   if sparse and ragged:
@@ -1282,7 +1282,7 @@ def shape(x):
   <tf.Tensor: shape=(2,), dtype=int32, numpy=array([2, 2], dtype=int32)>
   >>> input = tf.keras.backend.placeholder(shape=(2, 4, 5))
   >>> tf.keras.backend.shape(input)
-  <tf.Tensor 'Shape_...' shape=(3,) dtype=int32>
+  <KerasTensor: shape=(3,) dtype=int32 inferred_value=[2, 4, 5] ...>
 
   """
   return array_ops.shape(x)
@@ -1797,13 +1797,13 @@ def dot(x, y):
   >>> y = tf.keras.backend.placeholder(shape=(3, 4))
   >>> xy = tf.keras.backend.dot(x, y)
   >>> xy
-  <tf.Tensor ... shape=(2, 4) dtype=float32>
+  <KerasTensor: shape=(2, 4) dtype=float32 ...>
 
   >>> x = tf.keras.backend.placeholder(shape=(32, 28, 3))
   >>> y = tf.keras.backend.placeholder(shape=(3, 4))
   >>> xy = tf.keras.backend.dot(x, y)
   >>> xy
-  <tf.Tensor ... shape=(32, 28, 4) dtype=float32>
+  <KerasTensor: shape=(32, 28, 4) dtype=float32 ...>
 
   >>> x = tf.keras.backend.random_uniform_variable(shape=(2, 3), low=0, high=1)
   >>> y = tf.keras.backend.ones((4, 3, 5))
@@ -2053,10 +2053,10 @@ def transpose(x):
          [3.,  6.]], dtype=float32)
   >>> input = tf.keras.backend.placeholder((2, 3))
   >>> input
-  <tf.Tensor 'Placeholder_...' shape=(2, 3) dtype=float32>
+  <KerasTensor: shape=(2, 3) dtype=float32 ...>
   >>> input_transposed = tf.keras.backend.transpose(input)
   >>> input_transposed
-  <tf.Tensor 'Transpose_...' shape=(3, 2) dtype=float32>
+  <KerasTensor: shape=(3, 2) dtype=float32 ...>
   """
   return array_ops.transpose(x)
 

tensorflow/python/keras/engine/keras_tensor.py

@@ -30,7 +30,7 @@ from tensorflow.python.util import object_identity
 
 # pylint: disable=g-classes-have-attributes
 
-_KERAS_TENSORS_ENABLED = False
+_KERAS_TENSORS_ENABLED = True
 
 
 def enable_keras_tensors():