Migrate graph plugin to Keras 3 and remove dependencies on Keras 2 and tf-keras-nightly

.github/workflows/ci.yml

@@ -31,7 +31,6 @@ env:
   BUILDIFIER_SHA256SUM: 'e92a6793c7134c5431c58fbc34700664f101e5c9b1c1fcd93b97978e8b7f88db'
   BUILDOZER_SHA256SUM: '3d58a0b6972e4535718cdd6c12778170ea7382de7c75bc3728f5719437ffb84d'
   TENSORFLOW_VERSION: 'tf-nightly'
-  TF_KERAS_VERSION: 'tf-keras-nightly' # Keras 2
 
 jobs:
   build:
@@ -78,7 +77,6 @@ jobs:
         run: |
           python -m pip install -U pip
           pip install "${TENSORFLOW_VERSION}"
-          pip install "${TF_KERAS_VERSION}"
         if: matrix.tf_version_id != 'notf'
       - name: 'Install Python dependencies'
         run: |

DEVELOPMENT.md

@@ -19,7 +19,7 @@ TensorBoard at HEAD relies on the nightly installation of TensorFlow: this allow
 $ virtualenv -p python3 tf
 $ source tf/bin/activate
 (tf)$ pip install --upgrade pip
-(tf)$ pip install tf-nightly tf-keras-nightly -r tensorboard/pip_package/requirements.txt -r tensorboard/pip_package/requirements_dev.txt
+(tf)$ pip install tf-nightly -r tensorboard/pip_package/requirements.txt -r tensorboard/pip_package/requirements_dev.txt
 (tf)$ pip uninstall -y tb-nightly
 ```
 

tensorboard/plugins/graph/graphs_plugin_v2_test.py

@@ -24,12 +24,9 @@
 from tensorboard.compat.proto import graph_pb2
 from tensorboard.plugins.graph import graphs_plugin_test
 
-# Stay on Keras 2 for now: https://github.com/keras-team/keras/issues/18467.
-version_fn = getattr(tf.keras, "version", None)
-if version_fn and version_fn().startswith("3."):
-    import tf_keras as keras  # Keras 2
-else:
-    keras = tf.keras  # Keras 2
+
+# Graph plugin V2 Keras 3 is only supported in TensorFlow eager mode.
+tf.compat.v1.enable_eager_execution()
 
 
 class GraphsPluginV2Test(
@@ -41,13 +38,13 @@ def generate_run(
         x, y = np.ones((10, 10)), np.ones((10, 1))
         val_x, val_y = np.ones((4, 10)), np.ones((4, 1))
 
-        model = keras.Sequential(
+        model = tf.keras.Sequential(
             [
-                keras.layers.Dense(10, activation="relu"),
-                keras.layers.Dense(1, activation="sigmoid"),
+                tf.keras.layers.Dense(10, activation="relu"),
+                tf.keras.layers.Dense(1, activation="sigmoid"),
             ]
         )
-        model.compile("rmsprop", "binary_crossentropy")
+        model.compile(optimizer="rmsprop", loss="binary_crossentropy")
 
         model.fit(
             x,
@@ -56,7 +53,7 @@ def generate_run(
             batch_size=2,
             epochs=1,
             callbacks=[
-                keras.callbacks.TensorBoard(
+                tf.keras.callbacks.TensorBoard(
                     log_dir=os.path.join(logdir, run_name),
                     write_graph=include_graph,
                 )

tensorboard/plugins/graph/keras_util.py

@@ -117,6 +117,24 @@ def _norm_to_list_of_layers(maybe_layers):
     )
 
 
+def _get_inbound_nodes(layer):
+    """Returns a list of [name, size, index] for all inbound nodes of the given layer."""
+    inbound_nodes = []
+    if layer.get("inbound_nodes") is not None:
+        for maybe_inbound_node in layer.get("inbound_nodes", []):
+            for inbound_node_args in maybe_inbound_node.get("args", []):
+                # Sometimes this field is a list when there are multiple inbound nodes
+                # for the given layer.
+                if not isinstance(inbound_node_args, list):
+                    inbound_node_args = [inbound_node_args]
+                for arg in inbound_node_args:
+                    history = arg.get("config", {}).get("keras_history", [])
+                    if len(history) < 3:
+                        continue
+                    inbound_nodes.append(history[:3])
+    return inbound_nodes
+
+
 def _update_dicts(
     name_scope,
     model_layer,
@@ -149,7 +167,7 @@ def _update_dicts(
     node_name = _scoped_name(name_scope, layer_config.get("name"))
     input_layers = layer_config.get("input_layers")
     output_layers = layer_config.get("output_layers")
-    inbound_nodes = model_layer.get("inbound_nodes")
+    inbound_nodes = _get_inbound_nodes(model_layer)
 
     is_functional_model = bool(input_layers and output_layers)
     # In case of [1] and the parent model is functional, current layer
@@ -164,7 +182,7 @@ def _update_dicts(
     elif is_parent_functional_model and not is_functional_model:
         # Sequential model can take only one input. Make sure inbound to the
         # model is linked to the first layer in the Sequential model.
-        prev_node_name = _scoped_name(name_scope, inbound_nodes[0][0][0])
+        prev_node_name = _scoped_name(name_scope, inbound_nodes[0][0])
     elif (
         not is_parent_functional_model
         and prev_node_name
@@ -244,33 +262,31 @@ def keras_model_to_graph_def(keras_layer):
             tf_dtype = dtypes.as_dtype(layer_config.get("dtype"))
             node_def.attr["dtype"].type = tf_dtype.as_datatype_enum
         if layer.get("inbound_nodes") is not None:
-            for maybe_inbound_node in layer.get("inbound_nodes"):
-                inbound_nodes = _norm_to_list_of_layers(maybe_inbound_node)
-                for [name, size, index, _] in inbound_nodes:
-                    inbound_name = _scoped_name(name_scope, name)
-                    # An input to a layer can be output from a model. In that case, the name
-                    # of inbound_nodes to a layer is a name of a model. Remap the name of the
-                    # model to output layer of the model. Also, since there can be multiple
-                    # outputs in a model, make sure we pick the right output_layer from the model.
-                    inbound_node_names = model_name_to_output.get(
-                        inbound_name, [inbound_name]
-                    )
-                    # There can be multiple inbound_nodes that reference the
-                    # same upstream layer. This causes issues when looking for
-                    # a particular index in that layer, since the indices
-                    # captured in `inbound_nodes` doesn't necessarily match the
-                    # number of entries in the `inbound_node_names` list. To
-                    # avoid IndexErrors, we just use the last element in the
-                    # `inbound_node_names` in this situation.
-                    # Note that this is a quick hack to avoid IndexErrors in
-                    # this situation, and might not be an appropriate solution
-                    # to this problem in general.
-                    input_name = (
-                        inbound_node_names[index]
-                        if index < len(inbound_node_names)
-                        else inbound_node_names[-1]
-                    )
-                    node_def.input.append(input_name)
+            for name, size, index in _get_inbound_nodes(layer):
+                inbound_name = _scoped_name(name_scope, name)
+                # An input to a layer can be output from a model. In that case, the name
+                # of inbound_nodes to a layer is a name of a model. Remap the name of the
+                # model to output layer of the model. Also, since there can be multiple
+                # outputs in a model, make sure we pick the right output_layer from the model.
+                inbound_node_names = model_name_to_output.get(
+                    inbound_name, [inbound_name]
+                )
+                # There can be multiple inbound_nodes that reference the
+                # same upstream layer. This causes issues when looking for
+                # a particular index in that layer, since the indices
+                # captured in `inbound_nodes` doesn't necessarily match the
+                # number of entries in the `inbound_node_names` list. To
+                # avoid IndexErrors, we just use the last element in the
+                # `inbound_node_names` in this situation.
+                # Note that this is a quick hack to avoid IndexErrors in
+                # this situation, and might not be an appropriate solution
+                # to this problem in general.
+                input_name = (
+                    inbound_node_names[index]
+                    if index < len(inbound_node_names)
+                    else inbound_node_names[-1]
+                )
+                node_def.input.append(input_name)
         elif prev_node_name is not None:
             node_def.input.append(prev_node_name)