code: def test_dataset_functional(relay_server, wandb_init):
      with relay_server() as relay:
          run = wandb_init()
  
          data = tf.data.Dataset.range(5).map(lambda x: (x, 1)).batch(1)
          inputs = tf.keras.Input(shape=(1,))
          outputs = tf.keras.layers.Dense(1)(inputs)
  
          model = tf.keras.Model(
              inputs=inputs,
              outputs=outputs,
          )
          model.compile(
              optimizer=tf.keras.optimizers.Adam(),
              loss="mse",
          )
          model.fit(
              data,
              callbacks=[
                  WandbCallback(save_model=False),
              ],
          )
  
          run_dir = run.dir
          run.finish()
  
      summary = relay.context.get_run_summary(run.id)
      assert (
          graph_json(run_dir, summary["graph"])["nodes"][0][
              "class_name"
          ]
          == "InputLayer"
      )
  
imports: import glob
  import json
  import os
  import pytest
  import tensorflow
  from tensorflow.keras import backend
  from tensorflow.keras.layers import Dense, Flatten, Reshape
  from tensorflow.keras.models import Sequential
  from wandb.keras import WandbCallback
label_desc: Use image node_id 12 for calling a global function with example usage: test_dataset_functional(relay_server, wandb_init) without return types