code: def _cat_dataset(
      n_samples, input_dim, n_classes, batch_size, generator=False
  ):
      """Creates a categorically encoded dataset.
  
      Creates a categorically encoded dataset (y is categorical).
      returns the specified dataset either as a static array or as a generator.
      Will have evenly split samples across each output class.
      Each output class will be a different point in the input space.
  
      Args:
        n_samples: number of rows
        input_dim: input dimensionality
        n_classes: output dimensionality
        batch_size: The desired batch_size
        generator: False for array, True for generator
  
      Returns:
        X as (n_samples, input_dim), Y as (n_samples, n_outputs)
      """
      x_stack = []
      y_stack = []
      for i_class in range(n_classes):
          x_stack.append(
              tf.constant(
                  1 * i_class, tf.float32, (n_samples, input_dim)
              )
          )
          y_stack.append(
              tf.constant(i_class, tf.float32, (n_samples, n_classes))
          )
      x_set, y_set = tf.stack(x_stack), tf.stack(y_stack)
      if generator:
          dataset = tf.data.Dataset.from_tensor_slices((x_set, y_set))
          dataset = dataset.batch(batch_size=batch_size)
          return dataset
      return x_set, y_set
  
imports: from absl.testing import parameterized
  import tensorflow.compat.v1
  from tensorflow.compat.v1.python.framework import ops
  from tensorflow.compat.v1.python.keras import keras_parameterized
  from tensorflow.compat.v1.python.keras import losses
  from tensorflow.compat.v1.python.keras.optimizer_v2.optimizer_v2 import OptimizerV2
  from tensorflow.compat.v1.python.keras.regularizers import L1L2
  from tensorflow_privacy.privacy.bolt_on import models
  from tensorflow_privacy.privacy.bolt_on.losses import StrongConvexMixin
  from tensorflow_privacy.privacy.bolt_on.optimizers import BoltOn
label_desc: Use image node_id 1 for calling a global function with example usage: _cat_dataset(n_samples, input_dim, n_classes, batch_size, generator) and returns: x_set,  y_set, dataset