code: def get_input_dataset(in_data, out_data=None):
      batch_size = in_data.shape[0]
      if out_data is None:
          return tf.data.Dataset.from_tensor_slices(in_data).batch(
              batch_size
          )
  
      return tf.data.Dataset.from_tensor_slices(
          (in_data, out_data)
      ).batch(batch_size)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import numpy
  import tensorflow
  from tensorflow.python.framework import test_util
  from tensorflow_text.python.keras.layers.todense import ToDense
label_desc: Use image node_id 1 for calling a global function with example usage: get_input_dataset(in_data, out_data) and returns: tf, tf