code: def get_dsprites_pp(predicted_attribute, num_classes=None):
      """Data preprocess function for dsprites dataset."""
  
      attribute_to_classes = {
          "label_shape": 3,
          "label_scale": 6,
          "label_orientation": 40,
          "label_x_position": 32,
          "label_y_position": 32,
      }
  
      def _dsprites_pp(data):
          # For consistency with other datasets, image needs to have three channels
          # and be in [0, 255).   # pylint: disable=unused-argument
          # data["image"] = tf.tile(data["image"], [1, 1, 3]) * 255
          data["image"] = data["image"] * 255
  
          # If num_classes is set, we group together nearby integer values to arrive
          # at the desired number of classes. This is useful for example for grouping
          # together different spatial positions.
          num_original_classes = attribute_to_classes[
              predicted_attribute
          ]
          n_cls = (
              num_original_classes
              if num_classes is None
              else num_classes
          )
          if (
              not isinstance(n_cls, int)
              or n_cls <= 1
              or n_cls > num_original_classes
          ):
              raise ValueError(
                  "The number of classes should be None or in [2, ..., num_classes]."
              )
          class_division_factor = float(num_original_classes) / n_cls
  
          data["label"] = tf.cast(
              tf.math.floordiv(
                  tf.cast(data[predicted_attribute], tf.float32),
                  class_division_factor,
              ),
              data[predicted_attribute].dtype,
          )
          return data
  
      return _dsprites_pp
  
imports: import numpy
  from scenic.dataset_lib.big_transfer.registry import Registry
  import tensorflow.compat.v1
label_desc: Use image node_id 1 for calling a global function with example usage: get_dsprites_pp(predicted_attribute, num_classes) and returns: _dsprites_pp, data