Compatibility between tensorboard and keras3/keras-core to display functional model graphs

When using keras-nightly which is now a preview of the keras 3.0 to come, tensorboard fails to represent the conceptual graph of a functional keras model.

I know that tensorboard is intended to work with the keras of same version, so i put this as a feature request instead of a bug. I tried also with tensorboard-nightly with the same result: sequential models are well represented but not functional models, probably because of a change in json representation. The error comes from "inbound_nodes" which can now be a list of dictionaries/lists instead of a list of lists.

Here is a minimal example showing the bug:

import json

from keras import Input, Model, Sequential
from keras.layers import Dense

from tensorboard.plugins.graph.keras_util import keras_model_to_graph_def


input = Input((1,))
layer = Dense(3)

## sequential: working
model_seq = Sequential(layers=[input, layer])
graph_def = keras_model_to_graph_def(json.loads(model_seq.to_json()))

## functional: not working  (with keras-nightly aka keras3)
output = layer(input)
model_func = Model(input, output)
graph_def = keras_model_to_graph_def(json.loads(model_func.to_json()))

which results in

Traceback (most recent call last):
File "/home/nolwen/Projects/decomon/tensorboard/minimal_example.py", line 19, in
graph_def = keras_model_to_graph_def(json.loads(model_func.to_json()))
File "/home/nolwen/Projects/decomon/tensorboard/tb-keras3-venv/lib/python3.9/site-packages/tensorboard/plugins/graph/keras_util.py", line 248, in keras_model_to_graph_def
inbound_nodes = _norm_to_list_of_layers(maybe_inbound_node)
File "/home/nolwen/Projects/decomon/tensorboard/tb-keras3-venv/lib/python3.9/site-packages/tensorboard/plugins/graph/keras_util.py", line 116, in _norm_to_list_of_layers
maybe_layers if isinstance(maybe_layers[0], (list,)) else [maybe_layers]
KeyError: 0

Here is an extract of the json reprenstation showing that inbound_nodes members are not always lists:

{

    "inbound_nodes": [
          {
            "args": [
              {
                "class_name": "__keras_tensor__",
                "config": {
                  "shape": [
                    null,
                    1
                  ],
                  "dtype": "float32",
                  "keras_history": [
                    "input_layer",
                    0,
                    0
                  ]
                }
              }
            ],
            "kwargs": {}
          }
        ]
}

You have to run it in an environment with tensorboard (potentially tb-nightly) and keras-nightly. For instance with:

python -m venv tb-keras3-venv
. tb-keras3-venv/bin/activate
pip install tf-nightly
pip uninstall tf-keras-nightly keras-nightly -y  # be sure to avoid a mix of both keras
pip install keras-nightly

Start with tensorboard/plugins/graph/keras_util.py, especially keras_model_to_graph_def and _norm_to_list_of_layers. Reproduce the failure using the provided minimal functional-model example in an environment with tf-nightly and keras-nightly, then verify that functional model JSON with dictionary-based inbound_nodes produces a graph while sequential models continue to work.