code: def get_onehot(
      depth, key="labels", key_result=None, multi=True, on=1.0, off=0.0
  ):
      """One-hot encodes the input.
  
      Args:
        depth: Length of the one-hot vector (how many classes).
        key: Key of the data to be one-hot encoded.
        key_result: Key under which to store the result (same as `key` if None).
        multi: If there are multiple labels, whether to merge them into the same
          "multi-hot" vector (True) or keep them as an extra dimension (False).
        on: Value to fill in for the positive label (default: 1).
        off: Value to fill in for negative labels (default: 0).
  
      Returns:
        Data dictionary.
      """
  
      def _onehot(data):
          # When there's more than one label, this is significantly more efficient
          # than using tf.one_hot followed by tf.reduce_max; we tested.
          labels = data[key]
          if labels.shape.rank > 0 and multi:
              # Currently, the assertion below is only used for datasets with single
              # labels. In a multi-label dataset either `on` or `off` should be computed
              # dynamically to yield the correct sum, when the number of labels varies.
              x = tf.scatter_nd(
                  labels[:, None],
                  tf.ones(tf.shape(labels)[0]),
                  (depth,),
              )
              x = tf.clip_by_value(x, 0, 1) * (on - off) + off
          else:
              assert np.isclose(
                  on + off * (depth - 1), 1
              ), "All on and off values must sum to 1"
              x = tf.one_hot(labels, depth, on_value=on, off_value=off)
          data[key_result or key] = x
          return data
  
      return _onehot
  
imports: from typing import Optional, Tuple
  import numpy
  from scenic.dataset_lib.big_transfer.preprocessing import autoaugment
  from scenic.dataset_lib.big_transfer.preprocessing import utils
  from scenic.dataset_lib.big_transfer.registry import Registry
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_addons import image
label_desc: Use image node_id 30 for calling a global function with example usage: get_onehot(depth, key, key_result, multi, on, off) and returns: _onehot, data