def map_output(x, axis):
      """Maps output of tf.vectorized_map to the final output."""
      x = tf_np.asarray(x)
      if axis is None:
        # Note that `tf.vectorized_map always batches the outputs.
        # Here we unbatch it again.
        return x[0, ...]
      elif axis == 0:
        return x
      else:
        # Need to transpose the output.
        return tf_np.moveaxis(x, 0, axis)
Maps output of tf.vectorized_map to the final output.
trax/tf_numpy/extensions/extensions.py