Python tensorflow.map_fn() Examples

def preprocess(self, inputs):
    """Feature-extractor specific preprocessing.

    See base class.

    Args:
      inputs: a [batch, height_in, width_in, channels] float tensor representing
        a batch of images with values between 0 and 255.0.

    Returns:
      preprocessed_inputs: a [batch, height_out, width_out, channels] float
        tensor representing a batch of images.
    Raises:
      ValueError: if inputs tensor does not have type tf.float32
    """
    if inputs.dtype is not tf.float32:
      raise ValueError('`preprocess` expects a tf.float32 tensor')
    with tf.name_scope('Preprocessor'):
      # TODO: revisit whether to always use batch size as  the number of
      # parallel iterations vs allow for dynamic batching.
      resized_inputs = tf.map_fn(self._image_resizer_fn,
                                 elems=inputs,
                                 dtype=tf.float32)
      return self._feature_extractor.preprocess(resized_inputs) 

def main():
    args = parse_args()

    with tf.Session(graph=tf.Graph()) as session:
        input_var = tf.placeholder(
            tf.uint8, (None, 128, 64, 3), name="images")
        image_var = tf.map_fn(
            lambda x: _preprocess(x), tf.cast(input_var, tf.float32),
            back_prop=False)

        factory_fn = _network_factory()
        features, _ = factory_fn(image_var, reuse=None)
        features = tf.identity(features, name="features")

        saver = tf.train.Saver(slim.get_variables_to_restore())
        saver.restore(session, args.checkpoint_in)

        output_graph_def = tf.graph_util.convert_variables_to_constants(
            session, tf.get_default_graph().as_graph_def(),
            [features.name.split(":")[0]])
        with tf.gfile.GFile(args.graphdef_out, "wb") as file_handle:
            file_handle.write(output_graph_def.SerializeToString()) 

def _preprocess_frames(self, example, indices):
    """Instantiates the ops used to preprocess the frames data."""
    frames = tf.concat(example['frames'], axis=0)
    frames = tf.gather(frames, indices, axis=1)
    frames = tf.map_fn(
        _convert_frame_data, tf.reshape(frames, [-1]),
        dtype=tf.float32, back_prop=False)
    dataset_image_dimensions = tuple(
        [self._dataset_info.frame_size] * 2 + [_NUM_CHANNELS])
    frames = tf.reshape(
        frames, (-1, self._example_size) + dataset_image_dimensions)
    if (self._custom_frame_size and
        self._custom_frame_size != self._dataset_info.frame_size):
      frames = tf.reshape(frames, (-1,) + dataset_image_dimensions)
      new_frame_dimensions = (self._custom_frame_size,) * 2 + (_NUM_CHANNELS,)
      frames = tf.image.resize_bilinear(
          frames, new_frame_dimensions[:2], align_corners=True)
      frames = tf.reshape(
          frames, (-1, self._example_size) + new_frame_dimensions)
    return frames 

def _vectorised_get_cum_graph_size(nodes, graph_sizes):
    """
    Takes a list of node ids and graph sizes ordered by segment ID and returns the
    number of nodes contained in graphs with smaller segment ID.

    :param nodes: List of node ids of shape (nodes)
    :param graph_sizes: List of graph sizes (i.e. tf.math.segment_sum(tf.ones_like(I), I) where I are the
    segment IDs).
    :return: A list of shape (nodes) where each entry corresponds to the number of nodes contained in graphs
    with smaller segment ID for each node.
    """
    def get_cum_graph_size(node):
        cum_graph_sizes = tf.cumsum(graph_sizes, exclusive=True)
        indicator_if_smaller = tf.cast(node - cum_graph_sizes >= 0, tf.int32)
        graph_id = tf.reduce_sum(indicator_if_smaller) - 1
        return tf.cumsum(graph_sizes, exclusive=True)[graph_id]

    return tf.map_fn(get_cum_graph_size, nodes) 

def _encoded_image_string_tensor_input_placeholder():
  """Returns input that accepts a batch of PNG or JPEG strings.

  Returns:
    a tuple of input placeholder and the output decoded images.
  """
  batch_image_str_placeholder = tf.placeholder(
      dtype=tf.string,
      shape=[None],
      name='encoded_image_string_tensor')
  def decode(encoded_image_string_tensor):
    image_tensor = tf.image.decode_image(encoded_image_string_tensor,
                                         channels=3)
    image_tensor.set_shape((None, None, 3))
    return image_tensor
  return (batch_image_str_placeholder,
          tf.map_fn(
              decode,
              elems=batch_image_str_placeholder,
              dtype=tf.uint8,
              parallel_iterations=32,
              back_prop=False)) 

def build_detection(self, reuse=False):
    with tf.variable_scope('detection', reuse=reuse):
      def _translation_match(x, z):  # translation match for one example within a batch
        x = tf.expand_dims(x, 0)  # [1, in_height, in_width, in_channels]
        z = tf.expand_dims(z, -1)  # [filter_height, filter_width, in_channels, 1]
        return tf.nn.conv2d(x, z, strides=[1, 1, 1, 1], padding='VALID', name='translation_match')

      output = tf.map_fn(lambda x: _translation_match(x[0], x[1]),
                         (self.instance_embeds, self.templates),
                         dtype=self.instance_embeds.dtype)
      output = tf.squeeze(output, [1, 4])  # of shape e.g., [8, 15, 15]

      # Adjust score, this is required to make training possible.
      config = self.model_config['adjust_response_config']
      bias = tf.get_variable('biases', [1],
                             dtype=tf.float32,
                             initializer=tf.constant_initializer(0.0, dtype=tf.float32),
                             trainable=config['train_bias'])
      response = config['scale'] * output + bias
      self.response = response 

def from_config(cls, config, custom_objects=None):
    model = super(CalibratedLinear, cls).from_config(
        config, custom_objects=custom_objects)
    try:
      model_config = tf.keras.utils.deserialize_keras_object(
          config.get('model_config'), custom_objects=custom_objects)
      premade_lib.verify_config(model_config)
      model.model_config = model_config
    except ValueError:
      logging.warning(
          'Could not load model_config. Constructing model without it: %s',
          str(config.get('model_config')))
    return model


# TODO: add support for tf.map_fn and inputs of shape (B, ?, input_dim)
# as well as non-ragged inputs using padding/mask. 

def main():
    args = parse_args()

    with tf.Session(graph=tf.Graph()) as session:
        input_var = tf.placeholder(
            tf.uint8, (None, 128, 64, 3), name="images")
        image_var = tf.map_fn(
            lambda x: _preprocess(x), tf.cast(input_var, tf.float32),
            back_prop=False)

        factory_fn = _network_factory()
        features, _ = factory_fn(image_var, reuse=None)
        features = tf.identity(features, name="features")

        saver = tf.train.Saver(slim.get_variables_to_restore())
        saver.restore(session, args.checkpoint_in)

        output_graph_def = tf.graph_util.convert_variables_to_constants(
            session, tf.get_default_graph().as_graph_def(),
            [features.name.split(":")[0]])
        with tf.gfile.GFile(args.graphdef_out, "wb") as file_handle:
            file_handle.write(output_graph_def.SerializeToString()) 

def build_detection(self):
    self.embeds = self.get_image_embedding(self.search_images, reuse=True)
    with tf.variable_scope('detection'):
      def _translation_match(x, z):
        x = tf.expand_dims(x, 0)  # [batch, in_height, in_width, in_channels]
        z = tf.expand_dims(z, -1)  # [filter_height, filter_width, in_channels, out_channels]
        return tf.nn.conv2d(x, z, strides=[1, 1, 1, 1], padding='VALID', name='translation_match')

      output = tf.map_fn(
        lambda x: _translation_match(x[0], x[1]),
        (self.embeds, self.templates), dtype=self.embeds.dtype)  # of shape [3, 1, 17, 17, 1]
      output = tf.squeeze(output, [1, 4])  # of shape e.g. [3, 17, 17]

      bias = tf.get_variable('biases', [1],
                             dtype=tf.float32,
                             initializer=tf.constant_initializer(0.0, dtype=tf.float32),
                             trainable=False)
      response = self.model_config['adjust_response_config']['scale'] * output + bias
      self.response = response 

def _tf_example_input_placeholder():
  """Returns input that accepts a batch of strings with tf examples.

  Returns:
    a tuple of input placeholder and the output decoded images.
  """
  batch_tf_example_placeholder = tf.placeholder(
      tf.string, shape=[None], name='tf_example')
  def decode(tf_example_string_tensor):
    tensor_dict = tf_example_decoder.TfExampleDecoder().decode(
        tf_example_string_tensor)
    image_tensor = tensor_dict[fields.InputDataFields.image]
    return image_tensor
  return (batch_tf_example_placeholder,
          tf.map_fn(decode,
                    elems=batch_tf_example_placeholder,
                    dtype=tf.uint8,
                    parallel_iterations=32,
                    back_prop=False)) 

def _encoded_image_string_tensor_input_placeholder():
  """Returns input that accepts a batch of PNG or JPEG strings.

  Returns:
    a tuple of input placeholder and the output decoded images.
  """
  batch_image_str_placeholder = tf.placeholder(
      dtype=tf.string,
      shape=[None],
      name='encoded_image_string_tensor')
  def decode(encoded_image_string_tensor):
    image_tensor = tf.image.decode_image(encoded_image_string_tensor,
                                         channels=3)
    image_tensor.set_shape((None, None, 3))
    return image_tensor
  return (batch_image_str_placeholder,
          tf.map_fn(
              decode,
              elems=batch_image_str_placeholder,
              dtype=tf.uint8,
              parallel_iterations=32,
              back_prop=False)) 

def preprocess(self, inputs):
    """Feature-extractor specific preprocessing.

    See base class.

    Args:
      inputs: a [batch, height_in, width_in, channels] float tensor representing
        a batch of images with values between 0 and 255.0.

    Returns:
      preprocessed_inputs: a [batch, height_out, width_out, channels] float
        tensor representing a batch of images.
    Raises:
      ValueError: if inputs tensor does not have type tf.float32
    """
    if inputs.dtype is not tf.float32:
      raise ValueError('`preprocess` expects a tf.float32 tensor')
    with tf.name_scope('Preprocessor'):
      # TODO: revisit whether to always use batch size as  the number of
      # parallel iterations vs allow for dynamic batching.
      resized_inputs = tf.map_fn(self._image_resizer_fn,
                                 elems=inputs,
                                 dtype=tf.float32)
      return self._feature_extractor.preprocess(resized_inputs) 

def add_noise(ids, sequence_length):
  """Wraps add_noise_python for a batch of tensors."""

  def _add_noise_single(ids, sequence_length):
    noisy_ids = add_noise_python(ids[:sequence_length])
    noisy_sequence_length = len(noisy_ids)
    ids[:noisy_sequence_length] = noisy_ids
    ids[noisy_sequence_length:] = 0
    return ids, np.int32(noisy_sequence_length)

  noisy_ids, noisy_sequence_length = tf.map_fn(
      lambda x: tf.py_func(_add_noise_single, x, [ids.dtype, tf.int32]),
      [ids, sequence_length],
      dtype=[ids.dtype, tf.int32],
      back_prop=False)

  noisy_ids.set_shape(ids.get_shape())
  noisy_sequence_length.set_shape(sequence_length.get_shape())

  return noisy_ids, noisy_sequence_length


# Step 3 

def test_input_with_unknown_leading_dimension(self):

    def get_random_shape_input():
      # Returns a Tensor of shape (?, 6)
      return tf.map_fn(lambda x: x * tf.random.normal([6]),
                       test_utils.get_tensor_with_random_shape())

    # Validate the premise of the test.
    assert get_random_shape_input().shape.as_list() == [None, 6]
    test_data = self.run_one_to_many_encode_decode(
        self.default_encoding_stage(), get_random_shape_input)
    self.common_asserts_for_test_data(test_data)
    encoded_shape = test_data.encoded_x[
        kashin.KashinHadamardEncodingStage.ENCODED_VALUES_KEY].shape
    self.assertEqual(test_data.x.shape[0], encoded_shape[0])
    self.assertEqual(8, encoded_shape[1]) 

def get_logits(self, x):
    if x.name in self._logits_dict:
      return self._logits_dict[x.name]

    x = tf.map_fn(tf.image.per_image_standardization, x)

    logits = self._model_fn(
        {
            "inputs": x
        },
        None,
        "attack",
        params=self._params,
        config=self._config)
    self._logits_dict[x.name] = logits

    return tf.squeeze(logits) 

def image_serving_input_fn():
  """Serving input fn for raw images."""

  def _preprocess_image(image_bytes):
    """Preprocess a single raw image."""
    image = resnet_preprocessing.preprocess_image(
        image_bytes=image_bytes, is_training=False)
    return image

  image_bytes_list = tf.placeholder(
      shape=[None],
      dtype=tf.string,
  )
  images = tf.map_fn(
      _preprocess_image, image_bytes_list, back_prop=False, dtype=tf.float32)
  return tf.estimator.export.ServingInputReceiver(
      images, {'image_bytes': image_bytes_list}) 

def image_serving_input_fn():
  """Serving input fn for raw images."""

  def _preprocess_image(image_bytes):
    """Preprocess a single raw image."""
    image = resnet_preprocessing.preprocess_image(
        image_bytes=image_bytes, is_training=False)
    return image

  image_bytes_list = tf.placeholder(
      shape=[None],
      dtype=tf.string,
  )
  images = tf.map_fn(
      _preprocess_image, image_bytes_list, back_prop=False, dtype=tf.float32)
  return tf.estimator.export.ServingInputReceiver(
      images, {'image_bytes': image_bytes_list}) 

def test_input_with_unknown_leading_dimension(self):

    def get_random_shape_input():
      # Returns a Tensor of shape (?, 6)
      return tf.map_fn(lambda x: x * tf.random.normal([6]),
                       test_utils.get_tensor_with_random_shape())

    # Validate the premise of the test.
    assert get_random_shape_input().shape.as_list() == [None, 6]
    test_data = self.run_one_to_many_encode_decode(
        self.default_encoding_stage(), get_random_shape_input)
    self.common_asserts_for_test_data(test_data)
    encoded_shape = test_data.encoded_x[
        stages_impl.HadamardEncodingStage.ENCODED_VALUES_KEY].shape
    self.assertEqual(test_data.x.shape[0], encoded_shape[0])
    self.assertEqual(8, encoded_shape[1]) 

def image_serving_input_fn():
  """Serving input fn for raw images."""

  def _preprocess_image(image_bytes):
    """Preprocess a single raw image."""
    image = resnet_preprocessing.preprocess_image(
        image_bytes=image_bytes, is_training=False)
    return image

  image_bytes_list = tf.placeholder(
      shape=[None],
      dtype=tf.string,
  )
  images = tf.map_fn(
      _preprocess_image, image_bytes_list, back_prop=False, dtype=tf.float32)
  return tf.estimator.export.ServingInputReceiver(
      images, {'image_bytes': image_bytes_list}) 

def _build(self, X):
        """Build the graph of this layer."""
        n_samples, input_shape = self._get_X_dims(X)
        Wdim = input_shape + [self.output_dim]

        W_init = initialise_weights(Wdim, self.init_fn)
        W = tf.Variable(W_init, name="W_map")
        summary_histogram(W)

        # Tiling W is much faster than mapping (tf.map_fn) the matmul
        Net = tf.matmul(X, _tile2samples(n_samples, W))

        # Regularizers
        penalty = self.l2 * tf.nn.l2_loss(W) + self.l1 * _l1_loss(W)

        # Optional Bias
        if self.use_bias is True:
            b_init = initialise_weights((1, self.output_dim), self.init_fn)
            b = tf.Variable(b_init, name="b_map")
            summary_histogram(b)

            Net += b
            penalty += self.l2 * tf.nn.l2_loss(b) + self.l1 * _l1_loss(b)

        return Net, penalty 

def _invert_permutation(tensor):
    '''wrapper for matrix'''
    return tf.cast(tf.map_fn(tf.invert_permutation, tensor), tf.float32) 

def build_valid_rl(self, shuffle=False):
    print("-" * 80)
    print("Build valid graph on shuffled data")
    with tf.device("/cpu:0"):
      # shuffled valid data: for choosing validation model
      if not shuffle and self.data_format == "NCHW":
        self.images["valid_original"] = np.transpose(
          self.images["valid_original"], [0, 3, 1, 2])
      x_valid_shuffle, y_valid_shuffle = tf.train.shuffle_batch(
        [self.images["valid_original"], self.labels["valid_original"]],
        batch_size=self.batch_size,
        capacity=25000,
        enqueue_many=True,
        min_after_dequeue=0,
        num_threads=16,
        seed=self.seed,
        allow_smaller_final_batch=True,
      )

      def _pre_process(x):
        x = tf.pad(x, [[4, 4], [4, 4], [0, 0]])
        x = tf.random_crop(x, [32, 32, 3], seed=self.seed)
        x = tf.image.random_flip_left_right(x, seed=self.seed)
        if self.data_format == "NCHW":
          x = tf.transpose(x, [2, 0, 1])
        return x

      if shuffle:
        x_valid_shuffle = tf.map_fn(
          _pre_process, x_valid_shuffle, back_prop=False)

    logits = self._model(x_valid_shuffle, is_training=True, reuse=True)
    valid_shuffle_preds = tf.argmax(logits, axis=1)
    valid_shuffle_preds = tf.to_int32(valid_shuffle_preds)
    self.valid_shuffle_acc = tf.equal(valid_shuffle_preds, y_valid_shuffle)
    self.valid_shuffle_acc = tf.to_int32(self.valid_shuffle_acc)
    self.valid_shuffle_acc = tf.reduce_sum(self.valid_shuffle_acc) 

def build_valid_rl(self, shuffle=False):
    print("-" * 80)
    print("Build valid graph on shuffled data")
    with tf.device("/cpu:0"):
    #with tf.device("/device:GPU:0"):
      # shuffled valid data: for choosing validation model
      if not shuffle and self.data_format == "NCHW":
        self.images["valid_original"] = np.transpose(
          self.images["valid_original"], [0, 3, 1, 2])
      x_valid_shuffle, y_valid_shuffle = tf.train.shuffle_batch(
        [self.images["valid_original"], self.labels["valid_original"]],
        batch_size=self.batch_size,
        capacity=25000,
        enqueue_many=True,
        min_after_dequeue=0,
        num_threads=16,
        seed=self.seed,
        allow_smaller_final_batch=True,
      )

      def _pre_process(x):
        x = tf.pad(x, [[4, 4], [4, 4], [0, 0]])
        x = tf.random_crop(x, [32, 32, 3], seed=self.seed)
        x = tf.image.random_flip_left_right(x, seed=self.seed)
        if self.data_format == "NCHW":
          x = tf.transpose(x, [2, 0, 1])

        return x

      if shuffle:
        x_valid_shuffle = tf.map_fn(
          _pre_process, x_valid_shuffle, back_prop=False)

    logits = self._model(x_valid_shuffle, False, reuse=True)
    valid_shuffle_preds = tf.argmax(logits, axis=1)
    valid_shuffle_preds = tf.to_int32(valid_shuffle_preds)
    self.valid_shuffle_acc = tf.equal(valid_shuffle_preds, y_valid_shuffle)
    self.valid_shuffle_acc = tf.to_int32(self.valid_shuffle_acc)
    self.valid_shuffle_acc = tf.reduce_sum(self.valid_shuffle_acc) 

def _gather(param, indices):
    '''wrapper for matrix'''
    return tf.map_fn(lambda x : tf.gather(x[0], x[1]), (param, indices), dtype=param.dtype) 

def call(self,docs):
            
            #input shape: batch x lines x words        
            doc_embeds = tf.map_fn(self._attention_step,docs,dtype=tf.float32)
            logits = []
            for l in self.classify_layers:
                logits.append(l(doc_embeds))
            return logits 

def train_step(self, positions, pi, z, learning_rate=0):
        # Record the operations used to compute the loss, so that the gradient
        # of the loss with respect to the variables can be computed.
        #         metrics = 0

        with tf.GradientTape() as tape:
            policy_head, value_head = self.model(positions, training=True)
            loss = self.compute_loss(pi, z, policy_head, value_head)
            # self.ComputeMetrics(y, logits)
            metrics = self.compute_metrics(pi, policy_head)
        grads = tape.gradient(loss, self.model.trainable_variables)

        # grads = self.average_gradients(tower_grads)
        # grads = self.optimizer.compute_gradients(self.loss)
        # defensive step 2 to clip norm
        # grads0_lst = tf.map_fn(lambda x: x[0], grads)  # [g for g, _ in grads]
        clipped_grads, self.norm = tf.clip_by_global_norm(grads, self.global_norm)

        # defensive step 3 check NaN
        # See: https://stackoverflow.com/questions/40701712/how-to-check-nan-in-gradients-in-tensorflow-when-updating
        grad_check = [tf.debugging.check_numerics(g, message='NaN Found!') for g in clipped_grads]
        with tf.control_dependencies(grad_check):
            self.optimizer.apply_gradients(
                zip(clipped_grads, self.model.trainable_variables),  # [v for _, v in grads]
                global_step=self.global_step, name='train_step')

        if self.is_logging:
            for grad, var in zip(grads, self.model.trainable_variables):
                if grad is not None:
                    summary_ops_v2.histogram(var.name + '/gradients', grad)
            for var in self.model.trainable_variables:
                summary_ops_v2.histogram(var.name, var)

        return metrics, loss, self.global_step

    #@profile 

def _graph_fn_get_records(self, num_records=1):
        # Sum total mass.
        current_size = self.read_variable(self.size)
        stored_elements_prob_sum = self.sum_segment_tree.reduce(start=0, limit=current_size - 1)

        # Sample the entire batch.
        sample = stored_elements_prob_sum * tf.random_uniform(shape=(num_records, ))

        # Sample by looking up prefix sum.
        sample_indices = tf.map_fn(fn=self.sum_segment_tree.index_of_prefixsum, elems=sample, dtype=tf.int32)
        # sample_indices = self.sum_segment_tree.index_of_prefixsum(sample)

        # Importance correction.
        total_prob = self.sum_segment_tree.reduce(start=0, limit=self.priority_capacity - 1)
        min_prob = self.min_segment_tree.get_min_value() / total_prob
        max_weight = tf.pow(x=min_prob * tf.cast(current_size, tf.float32), y=-self.beta)

        def importance_sampling_fn(sample_index):
            sample_prob = self.sum_segment_tree.get(sample_index) / stored_elements_prob_sum
            weight = tf.pow(x=sample_prob * tf.cast(current_size, tf.float32), y=-self.beta)

            return weight / max_weight

        corrected_weights = tf.map_fn(
            fn=importance_sampling_fn,
            elems=sample_indices,
            dtype=tf.float32
        )
        # sample_indices = tf.Print(sample_indices, [sample_indices, self.sum_segment_tree.values], summarize=1000,
        #                           message='sample indices, segment tree values = ')
        return self._read_records(indices=sample_indices), sample_indices, corrected_weights 

def subpixel_reshuffle_1D(X, m, name=None):
    """
    maps over the batch dimension
    """
    return tf.map_fn(lambda x: subpixel_reshuffle_1D_impl(x, m), X, name=name) 

def build_inputs_and_outputs(self):

    if self.frame_features:

      serialized_examples = tf.placeholder(tf.string, shape=(None,))

      fn = lambda x: self.build_prediction_graph(x)
      video_id_output, top_indices_output, top_predictions_output = (
          tf.map_fn(fn, serialized_examples, 
                    dtype=(tf.string, tf.int32, tf.float32)))

    else:

      serialized_examples = tf.placeholder(tf.string, shape=(None,))

      video_id_output, top_indices_output, top_predictions_output = (
          self.build_prediction_graph(serialized_examples))

    inputs = {"example_bytes": 
              saved_model_utils.build_tensor_info(serialized_examples)}

    outputs = {
        "video_id": saved_model_utils.build_tensor_info(video_id_output),
        "class_indexes": saved_model_utils.build_tensor_info(top_indices_output),
        "predictions": saved_model_utils.build_tensor_info(top_predictions_output)}

    return inputs, outputs 

def get_exemplar_images(images, exemplar_size, targets_pos=None):
  """Crop exemplar image from input images"""
  with tf.name_scope('get_exemplar_image'):
    batch_size, x_height, x_width = images.get_shape().as_list()[:3]
    z_height, z_width = exemplar_size

    if targets_pos is None:
      target_pos_single = [[get_center(x_height), get_center(x_width)]]
      targets_pos_ = tf.tile(target_pos_single, [batch_size, 1])
    else:
      targets_pos_ = targets_pos

    # convert to top-left corner based coordinates
    top = tf.to_int32(tf.round(targets_pos_[:, 0] - get_center(z_height)))
    bottom = tf.to_int32(top + z_height)
    left = tf.to_int32(tf.round(targets_pos_[:, 1] - get_center(z_width)))
    right = tf.to_int32(left + z_width)

    def _slice(x):
      f, t, l, b, r = x
      c = f[t:b, l:r]
      return c

    exemplar_img = tf.map_fn(_slice, (images, top, left, bottom, right), dtype=images.dtype)
    exemplar_img.set_shape([batch_size, z_height, z_width, 3])
    return exemplar_img