code: def _pack_with_custom_ops(dataset, keys, length):
      """Helper-function for packing a dataset which has already been batched.
  
      See pack_dataset()
  
      Relies on custom ops which require a custom compiled binary.
      Faster than _pack_with_tf_ops(), and denser packing.
  
      Args:
        dataset: a dataset containing padded batches of examples.
        keys: a collection of strings (must have length 1 or 2)
        length: a dictionary from key to integer
  
      Returns:
        a dataset.
      """
      from tensor2tensor.data_generators.ops import (
          pack_sequences_ops,
      )  # pylint: disable=g-import-not-at-top
  
      # faster and better packing but requires custom-built binary.
      if len(keys) == 1:
          (k1,) = keys
          k2 = k1
      elif len(keys) == 2:
          k1, k2 = keys
      else:
          raise ValueError(
              f"Packing op requires 1 or 2 keys. Got {len(keys)}"
          )
  
      def custom_pack_batch(x):
          """Map-function."""
          (
              k1_packed,
              k1_segmentation,
              k1_position,
              k2_packed,
              k2_segmentation,
              k2_position,
          ) = pack_sequences_ops.pack_sequences2(
              # cast to int64 for compatibility with custom ops
              tf.cast(x[k1], tf.int64),
              tf.cast(x[k2], tf.int64),
              length[k1],
              length[k2],
          )
          packed = {
              k1: k1_packed,
              k1 + "_segmentation": k1_segmentation,
              k1 + "_position": k1_position,
          }
          if len(keys) == 2:
              packed.update(
                  {
                      k2: k2_packed,
                      k2 + "_segmentation": k2_segmentation,
                      k2 + "_position": k2_position,
                  }
              )
  
          # cast back to int32
          for k, v in packed.items():
              packed[k] = tf.cast(v, tf.int32)
  
          return packed
  
      dataset = dataset.map(
          custom_pack_batch,
          num_parallel_calls=tf.data.experimental.AUTOTUNE,
      )
      dataset = dataset.unbatch()
      return dataset
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import os
  import gin
  import tensorflow.compat.v1
  import tensorflow_datasets
label_desc: Use image node_id 16 for calling a global function with example usage: _pack_with_custom_ops(dataset, keys, length) and returns: dataset, packed