jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _predict_in_batches(
            self, images, intrinsic_matrix, distortion_coeffs, camspace_up, boxes,
            internal_batch_size, aug_should_flip, aug_rotflipmat, aug_gammas, aug_scales,
            antialias_factor):

        num_aug = tf.shape(aug_gammas)[0]
        boxes_per_batch = internal_batch_size // num_aug
        boxes_flat = boxes.flat_values
        image_id_per_box = boxes.value_rowids()

        # Gamma decoding for correct image rescaling later on
        images = (tf.cast(images, tf.float32) / np.float32(255)) ** 2.2

        if boxes_per_batch == 0:
            # Run all as a single batch
            return self._predict_single_batch(
                images, intrinsic_matrix, distortion_coeffs, camspace_up, boxes_flat,
                image_id_per_box, aug_rotflipmat, aug_should_flip, aug_scales, aug_gammas,
                antialias_factor)
        else:
            # Chunk the image crops into batches and predict them one by one
            n_total_boxes = tf.shape(boxes_flat)[0]
            n_batches = tf.cast(tf.math.ceil(n_total_boxes / boxes_per_batch), tf.int32)
            poses3d_batches = tf.TensorArray(
                tf.float32, size=n_batches, element_shape=(None, None, self.joint_info.n_joints, 3),
                infer_shape=False)
            # CROP
            # crop_batches = tf.TensorArray(
            #    tf.float32, size=n_batches,
            #    element_shape=(None, None, None, None, 3),
            #    infer_shape=False)

            for i in tf.range(n_batches):
                # tf.autograph.experimental.set_loop_options(parallel_iterations=1)
                batch_slice = slice(i * boxes_per_batch, (i + 1) * boxes_per_batch)
                # CROP
                # crops, poses3d = self._predict_single_batch(
                poses3d = self._predict_single_batch(
                    images, intrinsic_matrix[batch_slice], distortion_coeffs[batch_slice],
                    camspace_up[batch_slice], boxes_flat[batch_slice],
                    image_id_per_box[batch_slice], aug_rotflipmat, aug_should_flip, aug_scales,
                    aug_gammas, antialias_factor)
                poses3d_batches = poses3d_batches.write(i, poses3d)

                # CROP
                # crop_batches = crop_batches.write(i, crops)

            # CROP
            # return crop_batches.concat(), poses3d_batches.concat()
            return poses3d_batches.concat()
//  [python]
def compute_mistakes(box_x, box_y, box_w, box_c_sim, target_x, target_y, target_w, target_is_plane, grid_nn):
    DETECTION_TRESHOLD = 0.5  # plane "detected" if predicted C>0.5 TODO: refactor this
    ERROR_TRESHOLD = 0.3  # plane correctly localized if predicted x,y,w within % of ground truth
    detect_correct = tf.logical_not(tf.logical_xor(tf.greater(box_c_sim, DETECTION_TRESHOLD), target_is_plane))
    ones = tf.ones(tf.shape(target_w))
    nonzero_target_w = tf.where(target_is_plane, target_w, ones)
    # true if correct size where there is a plane, nonsense value where there is no plane
    size_correct = tf.less(tf.abs(box_w - target_w) / nonzero_target_w, ERROR_TRESHOLD)
    # true if correct position where there is a plane, nonsense value where there is no plane
    position_correct = tf.less(tf.sqrt(tf.square(box_x - target_x) + tf.square(box_y - target_y)) / nonzero_target_w / grid_nn, ERROR_TRESHOLD)
    truth_no_plane = tf.logical_not(target_is_plane)
    size_correct = tf.logical_or(size_correct, truth_no_plane)
    position_correct = tf.logical_or(position_correct, truth_no_plane)
    size_correct = tf.logical_and(detect_correct, size_correct)
    position_correct = tf.logical_and(detect_correct, position_correct)
    all_correct = tf.logical_and(size_correct, position_correct)
    mistakes = tf.reduce_sum(tf.to_int32(tf.logical_not(all_correct)), axis=[1,2,3])  # shape [batch]
    return mistakes, size_correct, position_correct, all_correct
//  [python]
def compute_mistakes(box_x, box_y, box_w, box_c_sim, target_x, target_y, target_w, target_is_plane, grid_nn):
    DETECTION_TRESHOLD = 0.5  # plane "detected" if predicted C>0.5 TODO: refactor this
    ERROR_TRESHOLD = 0.3  # plane correctly localized if predicted x,y,w within % of ground truth
    detect_correct = tf.logical_not(tf.logical_xor(tf.greater(box_c_sim, DETECTION_TRESHOLD), target_is_plane))
    ones = tf.ones(tf.shape(target_w))
    nonzero_target_w = tf.where(target_is_plane, target_w, ones)
    # true if correct size where there is a plane, nonsense value where there is no plane
    size_correct = tf.less(tf.abs(box_w - target_w) / nonzero_target_w, ERROR_TRESHOLD)
    # true if correct position where there is a plane, nonsense value where there is no plane
    position_correct = tf.less(tf.sqrt(tf.square(box_x - target_x) + tf.square(box_y - target_y)) / nonzero_target_w / grid_nn, ERROR_TRESHOLD)
    truth_no_plane = tf.logical_not(target_is_plane)
    size_correct = tf.logical_or(size_correct, truth_no_plane)
    position_correct = tf.logical_or(position_correct, truth_no_plane)
    size_correct = tf.logical_and(detect_correct, size_correct)
    position_correct = tf.logical_and(detect_correct, position_correct)
    all_correct = tf.logical_and(size_correct, position_correct)
    mistakes = tf.reduce_sum(tf.cast(tf.logical_not(all_correct), tf.int32), axis=[1,2,3])  # shape [batch]
    return mistakes, size_correct, position_correct, all_correct
//  [python]
def convert_imgs_from_uint8(images, drange = [-1,1], nhwc_to_nchw = False):
    # Convert a batch of images from uint8 to float32 with configurable dynamic range
    # Can be used as an input transformation for Network.run()
    images = t.cast(images, tf.float32)
    if nhwc_to_nchw:
        images = tf.transpose(images, [0, 3, 1, 2])
    return images * ((drange[1] - drange[0]) / 255) + drange[0]
//  [python]
def convert_imgs_from_uint8(images, drange = [-1,1], nhwc_to_nchw = False):
    # Convert a batch of images from uint8 to float32 with configurable dynamic range
    # Can be used as an input transformation for Network.run()
    images = tf.cast(images, tf.float32)
    if nhwc_to_nchw:
        images = tf.transpose(images, [0, 3, 1, 2])
    return images * ((drange[1] - drange[0]) / 255) + drange[0]
//  [python]
def load_mean_theta():
    """loads mean theta values

    Returns:
        mean: [batch x 85]
    """
    config = Config()
    mean_values = h5py.File(config.SMPL_MEAN_THETA_PATH, mode='r')
    mean = np.zeros((1, config.NUM_SMPL_PARAMS))

    mean_pose = mean_values.get('pose')[()]  # 72
    mean_pose[:3] = 0.  # Ignore the global rotation.
    mean_pose[0] = np.pi  # This initializes the global pose to be up-right when projected

    mean_shape = mean_values.get('shape')[()]  # 10

    # init scale is 0.9
    mean[0, 0] = 0.9
    mean[:, config.NUM_CAMERA_PARAMS:] = np.hstack((mean_pose, mean_shape))
    mean = t.cast(mean, dtype=t.float32)
    return mean
//  [python]
def load_mean_theta():
    """loads mean theta values

    Returns:
        mean: [batch x 85]
    """
    config = Config()
    mean_values = h5py.File(config.SMPL_MEAN_THETA_PATH, mode='r')
    mean = np.zeros((1, config.NUM_SMPL_PARAMS))

    mean_pose = mean_values.get('pose')[()]  # 72
    mean_pose[:3] = 0.  # Ignore the global rotation.
    mean_pose[0] = np.pi  # This initializes the global pose to be up-right when projected

    mean_shape = mean_values.get('shape')[()]  # 10

    # init scale is 0.9
    mean[0, 0] = 0.9
    mean[:, config.NUM_CAMERA_PARAMS:] = np.hstack((mean_pose, mean_shape))
    mean = tf.cast(mean, dtype=tf.float32)
    return mean
//  [python]
def data_augment(self, data):
        # data_aug = data.copy()
        data_xyz = data[:, :, 0:3]
        data_f = data[:, :, 3:]
        batch_size = data_f.get_shape()[0]
        mirror_opt = np.random.choice([0, 1, 2])
        if mirror_opt == 0:
            data_xyz = tf.stack([data_xyz[:, :, 0], -data_xyz[:, :, 1], data_xyz[:, :, 2]], 2)
        elif mirror_opt == 1:
            theta = 2*3.14592653*np.random.rand() 
            R = np.array([[np.cos(theta), 0, -np.sin(theta)], [0, 1, 0], [np.sin(theta), 0, np.cos(theta)]])
            R = tf.convert_to_tensor(R, dtype=tf.float32)
            data_xyz = tf.reshape(data_xyz, [-1, 3])
            data_xyz = tf.matmul(data_xyz, R)
            data_xyz = tf.reshape(data_xyz, [-1,self.config.num_points, 3])
        elif mirror_opt == 2:
            sigma = 0.01
            clip = 0.05
            jittered_point = tf.clip_by_value(sigma * np.random.randn(self.config.num_points, 3), -1 * clip, clip)
            jittered_point = tf.tile(tf.expand_dims(jittered_point, axis=0), [self.config.batch_size,1,1])
            data_xyz = data_xyz + tf.to_float(jittered_point)
        data_aug = tf.concat([data_xyz, data_f],axis = -1)
        data_aug_t = tf.transpose(data_aug, [0, 2, 1])
        data_aug_t = tf.reshape(data_aug_t,[-1, 6, self.config.num_points])
        att_activation = tf.layers.dense(data_aug_t, 1, activation= None, use_bias = False, name='channel_attention' )
        att_activation = tf.transpose(att_activation, [0, 2, 1])
        att_scores = tf.nn.softmax(att_activation, axis=-1)
        data_aug = tf.multiply(data_aug, att_scores)
        return data_aug
//  [python]
def data_augment(self, data):
        # data_aug = data.copy()
        data_xyz = data[:, :, 0:3]
        data_f = data[:, :, 3:]
        batch_size = data_f.get_shape()[0]
        mirror_opt = np.random.choice([0, 1, 2])
        if mirror_opt == 0:
            data_xyz = tf.stack([data_xyz[:, :, 0], -data_xyz[:, :, 1], data_xyz[:, :, 2]], 2)
        elif mirror_opt == 1:
            theta = 2*3.14592653*np.random.rand() 
            R = np.array([[np.cos(theta), 0, -np.sin(theta)], [0, 1, 0], [np.sin(theta), 0, np.cos(theta)]])
            R = tf.convert_to_tensor(R, dtype=tf.float32)
            data_xyz = tf.reshape(data_xyz, [-1, 3])
            data_xyz = tf.matmul(data_xyz, R)
            data_xyz = tf.reshape(data_xyz, [-1,self.config.num_points, 3])
        elif mirror_opt == 2:
            sigma = 0.01
            clip = 0.05
            jittered_point = tf.clip_by_value(sigma * np.random.randn(self.config.num_points, 3), -1 * clip, clip)
            jittered_point = tf.tile(tf.expand_dims(jittered_point, axis=0), [self.config.batch_size,1,1])
            data_xyz = data_xyz + tf.cast(jittered_point, tf.float32)
        data_aug = tf.concat([data_xyz, data_f],axis = -1)
        data_aug_t = tf.transpose(data_aug, [0, 2, 1])
        data_aug_t = tf.reshape(data_aug_t,[-1, 6, self.config.num_points])
        att_activation = tf.layers.dense(data_aug_t, 1, activation= None, use_bias = False, name='channel_attention' )
        att_activation = tf.transpose(att_activation, [0, 2, 1])
        att_scores = tf.nn.softmax(att_activation, axis=-1)
        data_aug = tf.multiply(data_aug, att_scores)
        return data_aug
//  [python]
def convert_imgs_to_uint8(images, drange = [-1,1], nchw_to_nhwc = False, shrink = 1, lst = False):
    # Convert a batch of images from float32 to uint8 with configurable dynamic range
    # Can be used as an output transformation for Network.run()
    if lst:
        images = images[0]
    images = t.cast(images, t.float32)
    if shrink > 1:
        ksize = [1, 1, shrink, shrink]
        images = tf.nn.avg_pool(images, ksize = ksize, strides = ksize, padding="VALID", data_format="NCHW")
    if nchw_to_nhwc:
        images = tf.transpose(images, [0, 2, 3, 1])
    scale = 255 / (drange[1] - drange[0])
    images = images * scale + (0.5 - drange[0] * scale)
    return tf.saturate_cast(images, tf.uint8)
//  [python]
def convert_imgs_to_uint8(images, drange = [-1,1], nchw_to_nhwc = False, shrink = 1, lst = False):
    # Convert a batch of images from float32 to uint8 with configurable dynamic range
    # Can be used as an output transformation for Network.run()
    if lst:
        images = images[0]
    images = tf.cast(images, tf.float32)
    if shrink > 1:
        ksize = [1, 1, shrink, shrink]
        images = tf.nn.avg_pool(images, ksize = ksize, strides = ksize, padding="VALID", data_format="NCHW")
    if nchw_to_nhwc:
        images = tf.transpose(images, [0, 2, 3, 1])
    scale = 255 / (drange[1] - drange[0])
    images = images * scale + (0.5 - drange[0] * scale)
    return tf.saturate_cast(images, tf.uint8)
//  [python]
def myolo_mask_loss_graph(target_masks, target_class_ids, pred_masks):
    """Mask binary cross-entropy loss for the masks head.

    target_masks: [batch, num_rois, height, width].
        A float32 tensor of values 0 or 1. Uses zero padding to fill array.
    target_class_ids: [batch, num_rois]. Integer class IDs. Zero padded.
    pred_masks: [batch, proposals, height, width, num_classes] float32 tensor
                with values from 0 to 1.
    """
    # Reshape for simplicity. Merge first two dimensions into one.
    target_class_ids = K.reshape(target_class_ids, (-1,))
    mask_shape = tf.shape(target_masks)
    target_masks = K.reshape(target_masks, (-1, mask_shape[2], mask_shape[3]))
    pred_shape = tf.shape(pred_masks)
    pred_masks = K.reshape(pred_masks,
                           (-1, pred_shape[2], pred_shape[3], pred_shape[4]))
    # Permute predicted masks to [N, num_classes, height, width]
    pred_masks = tf.transpose(pred_masks, [0, 3, 1, 2])

    # Only positive ROIs contribute to the loss. And only
    # the class specific mask of each ROI.
    positive_ix = tf.where(target_class_ids > 0)[:, 0]
    positive_class_ids = tf.to_int64(
        tf.gather(target_class_ids, positive_ix))
    indices = tf.stack([positive_ix, positive_class_ids], axis=1)

    # Gather the masks (predicted and true) that contribute to loss
    y_true = tf.gather(target_masks, positive_ix)
    y_pred = tf.gather_nd(pred_masks, indices)

    # Compute binary cross entropy. If no positive ROIs, then return 0.
    # shape: [batch, roi, num_classes]
    loss = K.switch(tf.size(y_true) > 0,
                    K.binary_crossentropy(target=y_true, output=y_pred),
                    tf.constant(0.0))
    loss = K.mean(loss)
    return loss
//  [python]
def myolo_mask_loss_graph(target_masks, target_class_ids, pred_masks):
    """Mask binary cross-entropy loss for the masks head.

    target_masks: [batch, num_rois, height, width].
        A float32 tensor of values 0 or 1. Uses zero padding to fill array.
    target_class_ids: [batch, num_rois]. Integer class IDs. Zero padded.
    pred_masks: [batch, proposals, height, width, num_classes] float32 tensor
                with values from 0 to 1.
    """
    # Reshape for simplicity. Merge first two dimensions into one.
    target_class_ids = K.reshape(target_class_ids, (-1,))
    mask_shape = tf.shape(target_masks)
    target_masks = K.reshape(target_masks, (-1, mask_shape[2], mask_shape[3]))
    pred_shape = tf.shape(pred_masks)
    pred_masks = K.reshape(pred_masks,
                           (-1, pred_shape[2], pred_shape[3], pred_shape[4]))
    # Permute predicted masks to [N, num_classes, height, width]
    pred_masks = tf.transpose(pred_masks, [0, 3, 1, 2])

    # Only positive ROIs contribute to the loss. And only
    # the class specific mask of each ROI.
    positive_ix = tf.where(target_class_ids > 0)[:, 0]
    positive_class_ids = tf.cast(
        tf.gather(target_class_ids, positive_ix), tf.int64)
    indices = tf.stack([positive_ix, positive_class_ids], axis=1)

    # Gather the masks (predicted and true) that contribute to loss
    y_true = tf.gather(target_masks, positive_ix)
    y_pred = tf.gather_nd(pred_masks, indices)

    # Compute binary cross entropy. If no positive ROIs, then return 0.
    # shape: [batch, roi, num_classes]
    loss = K.switch(tf.size(y_true) > 0,
                    K.binary_crossentropy(target=y_true, output=y_pred),
                    tf.constant(0.0))
    loss = K.mean(loss)
    return loss
//  [python]
def hf_compute_loss(self, labels, logits):
        """CrossEntropyLoss that ignores pad tokens"""
        loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction=keras.losses.Reduction.NONE)
        if self.config.tf_legacy_loss:
            melted_labels = tf.reshape(labels, (-1,))
            active_loss = tf.not_equal(melted_labels, self.config.pad_token_id)
            reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, shape_list(logits)[2])), active_loss)
            labels = tf.boolean_mask(melted_labels, active_loss)
            return loss_fn(labels, reduced_logits)

        # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway
        unmasked_loss = loss_fn(tf.nn.relu(labels), logits)
        # make sure only non-padding labels affect the loss
        loss_mask = tf.to_float(tf.cast(labels != self.config.pad_token_id, dtype=tf.bool))
        masked_loss = unmasked_loss * loss_mask
        reduced_masked_loss = tf.reduce_sum(masked_loss) / tf.reduce_sum(loss_mask)
        return tf.reshape(reduced_masked_loss, (1,))
//  [python]
def hf_compute_loss(self, labels, logits):
        """CrossEntropyLoss that ignores pad tokens"""
        loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction=keras.losses.Reduction.NONE)
        if self.config.tf_legacy_loss:
            melted_labels = tf.reshape(labels, (-1,))
            active_loss = tf.not_equal(melted_labels, self.config.pad_token_id)
            reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, shape_list(logits)[2])), active_loss)
            labels = tf.boolean_mask(melted_labels, active_loss)
            return loss_fn(labels, reduced_logits)

        # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway
        unmasked_loss = loss_fn(tf.nn.relu(labels), logits)
        # make sure only non-padding labels affect the loss
        loss_mask = tf.cast(labels != self.config.pad_token_id, dtype=unmasked_loss.dtype)
        masked_loss = unmasked_loss * loss_mask
        reduced_masked_loss = tf.reduce_sum(masked_loss) / tf.reduce_sum(loss_mask)
        return tf.reshape(reduced_masked_loss, (1,))
//  [python]
def test_attention(self):
    num_heads = 12
    key_dim = 64
    seq_length = 1024
    batch_size = 2
    block_size = 64
    mask_layer = attention.BigBirdMasks(block_size=block_size)
    encoder_inputs_mask = tf.zeros((batch_size, seq_length), dtype=tf.int32)
    test_layer = attention.BigBirdAttention(
        num_heads=num_heads,
        key_dim=key_dim,
        from_block_size=block_size,
        to_block_size=block_size,
        seed=0)
    query = tf.random.normal(
        shape=(batch_size, seq_length, key_dim))
    masks = mask_layer(query, tf.to_float(encoder_inputs_mask))
    value = query
    output = test_layer(
        query=query,
        value=value,
        attention_mask=masks)
    self.assertEqual(output.shape, [batch_size, seq_length, key_dim])
//  [python]
def test_attention(self):
    num_heads = 12
    key_dim = 64
    seq_length = 1024
    batch_size = 2
    block_size = 64
    mask_layer = attention.BigBirdMasks(block_size=block_size)
    encoder_inputs_mask = tf.zeros((batch_size, seq_length), dtype=tf.int32)
    test_layer = attention.BigBirdAttention(
        num_heads=num_heads,
        key_dim=key_dim,
        from_block_size=block_size,
        to_block_size=block_size,
        seed=0)
    query = tf.random.normal(
        shape=(batch_size, seq_length, key_dim))
    masks = mask_layer(query, tf.cast(encoder_inputs_mask, dtype=tf.float64))
    value = query
    output = test_layer(
        query=query,
        value=value,
        attention_mask=masks)
    self.assertEqual(output.shape, [batch_size, seq_length, key_dim])
//  [python]
def minibatch(self, dataset, subset):
    with tf.compat.v1.name_scope('batch_processing'):
      images = [[] for i in range(self.device_count)]
      labels = [[] for i in range(self.device_count)]
      record_input = data_flow_ops.RecordInput(
          file_pattern=dataset.tf_record_pattern(subset),
          seed=randint(0, 9000),
          parallelism=64,
          buffer_size=10000,
          batch_size=self.batch_size,
          name='record_input')
      records = record_input.get_yield_op()
      records = tf.split(records, self.batch_size, 0)
      records = [tf.reshape(record, []) for record in records]
      for i in xrange(self.batch_size):
        value = records[i]
        image_buffer, label_index, bbox, _ = parse_example_proto(value)
        image = self.preprocess(image_buffer, bbox, i % 4)

        device_index = i % self.device_count
        images[device_index].append(image)
        labels[device_index].append(label_index)
      label_index_batch = [None] * self.device_count
      for device_index in xrange(self.device_count):
        images[device_index] = tf.parallel_stack(images[device_index])
        label_index_batch[device_index] = tf.concat(labels[device_index], 0)

        # dynamic_pad=True) # HACK TESTING dynamic_pad=True
        images[device_index] = tf.to_float(images[device_index])
        depth = 3
        images[device_index] = tf.reshape(
            images[device_index],
            shape=[self.batch_size_per_device, self.height, self.width, depth])
        label_index_batch[device_index] = tf.reshape(
            label_index_batch[device_index], [self.batch_size_per_device])
        # Display the training images in the visualizer.
        # tf.summary.image('images', images)

      return images, label_index_batch
//  [python]
def minibatch(self, dataset, subset):
    with tf.compat.v1.name_scope('batch_processing'):
      images = [[] for i in range(self.device_count)]
      labels = [[] for i in range(self.device_count)]
      record_input = data_flow_ops.RecordInput(
          file_pattern=dataset.tf_record_pattern(subset),
          seed=randint(0, 9000),
          parallelism=64,
          buffer_size=10000,
          batch_size=self.batch_size,
          name='record_input')
      records = record_input.get_yield_op()
      records = tf.split(records, self.batch_size, 0)
      records = [tf.reshape(record, []) for record in records]
      for i in xrange(self.batch_size):
        value = records[i]
        image_buffer, label_index, bbox, _ = parse_example_proto(value)
        image = self.preprocess(image_buffer, bbox, i % 4)

        device_index = i % self.device_count
        images[device_index].append(image)
        labels[device_index].append(label_index)
      label_index_batch = [None] * self.device_count
      for device_index in xrange(self.device_count):
        images[device_index] = tf.parallel_stack(images[device_index])
        label_index_batch[device_index] = tf.concat(labels[device_index], 0)

        # dynamic_pad=True) # HACK TESTING dynamic_pad=True
        images[device_index] = tf.cast(images[device_index], self.dtype)
        depth = 3
        images[device_index] = tf.reshape(
            images[device_index],
            shape=[self.batch_size_per_device, self.height, self.width, depth])
        label_index_batch[device_index] = tf.reshape(
            label_index_batch[device_index], [self.batch_size_per_device])
        # Display the training images in the visualizer.
        # tf.summary.image('images', images)

      return images, label_index_batch
//  [python]
def call(self, x, training=False):
        y = self.model(x)

        x_shape = tf.to_float(tf.shape(x))
        y_shape = tf.to_float(tf.shape(y))

        x_end, y_end = y_shape[1], y_shape[2]
        x_range = tf.range(x_end, dtype="float") * self.strides[0]
        y_range = tf.range(y_end, dtype="float") * self.strides[1]
        Y_mat, X_mat = tf.meshgrid(y_range, x_range)

        pred_x = y[..., 0]
        pred_y = y[..., 1]
        weights = tf.keras.activations.sigmoid(y[..., -1:])

        pred_x = pred_x + X_mat
        pred_y = pred_y + Y_mat

        pred = tf.stack((pred_x, pred_y), axis=-1)

        if training:
            pred = pred - x_shape[1:3] / 2

        pred = tf.concat([pred, y[..., 2:-1]], axis=-1)

        return pred, weights
//  [python]
def call(self, x, training=False):
        y = self.model(x)

        x_shape = tf.cast(tf.shape(x), "float")
        y_shape = tf.cast(tf.shape(y), "float")

        x_end, y_end = y_shape[1], y_shape[2]
        x_range = tf.range(x_end, dtype="float") * self.strides[0]
        y_range = tf.range(y_end, dtype="float") * self.strides[1]
        Y_mat, X_mat = tf.meshgrid(y_range, x_range)

        pred_x = y[..., 0]
        pred_y = y[..., 1]
        weights = tf.keras.activations.sigmoid(y[..., -1:])

        pred_x = pred_x + X_mat
        pred_y = pred_y + Y_mat

        pred = tf.stack((pred_x, pred_y), axis=-1)

        if training:
            pred = pred - x_shape[1:3] / 2

        pred = tf.concat([pred, y[..., 2:-1]], axis=-1)

        return pred, weights
//  [python]
def prediction(self):
        # input features
        with tf.variable_scope('input_layer',
                               partitioner=self.input_layer_partitioner,
                               reuse=tf.AUTO_REUSE):

            fm_cols = {}
            wide_input = tf.feature_column.input_layer(
                self.feature, self.wide_column, cols_to_output_tensors=fm_cols)
            fm_input = tf.stack([fm_cols[cols] for cols in self.fm_column], 1)
            dnn_input = tf.feature_column.input_layer(self.feature,
                                                      self.deep_column)

        if self.bf16:
            wide_input = tf.cast(wide_input, dtype=tf.bfloat16)
            fm_input = tf.cast(fm_input, dtype=tf.bfloat16)
            dnn_input = tf.cast(dnn_input, dtype=tf.bfloat16)

        # DNN part
        if self.bf16:
            with tf.variable_scope('dnn').keep_weights():
                dnn_output = self.dnn(dnn_input, self.dnn_hidden_units,
                                      'dnn_layer')
        else:
            with tf.variable_scope('dnn'):
                dnn_output = self.dnn(dnn_input, self.dnn_hidden_units,
                                      'dnn_layer')

        # linear / fisrt order part
        with tf.variable_scope('linear', reuse=tf.AUTO_REUSE) as linear:
            linear_output = tf.reduce_sum(wide_input, axis=1, keepdims=True)

        # FM second order part
        with tf.variable_scope('fm', reuse=tf.AUTO_REUSE) as fm:
            sum_square = tf.square(tf.reduce_sum(fm_input, axis=1))
            square_sum = tf.reduce_sum(tf.square(fm_input), axis=1)
            fm_output = 0.5 * tf.subtract(sum_square, square_sum)

        # Final dnn layer
        all_input = tf.concat([dnn_output, linear_output, fm_output], 1)
        if self.bf16:
            with tf.variable_scope('final_dnn').keep_weights():
                net = self.dnn(all_input, self.final_hidden_units, 'final_dnn')
            net = tf.cast(net, dtype=tf.float32)
        else:
            with tf.variable_scope('final_dnn'):
                net = self.dnn(all_input, self.final_hidden_units, 'final_dnn')

        net = tf.layers.dense(net, units=1)
        net = tf.math.sigmoid(net)

        return net
//  [python]
def bboxes_crop_or_pad(bboxes, xs, ys,
                       height, width,
                       offset_y, offset_x,
                       target_height, target_width):
    """Adapt bounding boxes to crop or pad operations.
    Coordinates are always supposed to be relative to the image.

    Arguments:
      bboxes: Tensor Nx4 with bboxes coordinates [y_min, x_min, y_max, x_max];
      height, width: Original image dimension;
      offset_y, offset_x: Offset to apply,
        negative if cropping, positive if padding;
      target_height, target_width: Target dimension after cropping / padding.
    """
    with tf.name_scope('bboxes_crop_or_pad'):
        # Rescale bounding boxes in pixels.
        scale = t.cast(t.stack([height, width, height, width]), bboxes.dtype)
        bboxes = bboxes * scale
        xs *= t.cast(width, bboxes.dtype)
        ys *= t.cast(height, bboxes.dtype)
        # Add offset.
        offset = t.cast(t.stack([offset_y, offset_x, offset_y, offset_x]), bboxes.dtype)
        bboxes = bboxes + offset
        xs += t.cast(offset_x, bboxes.dtype)
        ys += t.cast(offset_y, bboxes.dtype)
        
        # Rescale to target dimension.
        scale = t.cast(t.stack([target_height, target_width,
                                  target_height, target_width]), bboxes.dtype)
        bboxes = bboxes / scale
        xs = xs / t.cast(target_width, xs.dtype)
        ys = ys / t.cast(target_height, ys.dtype)
        return bboxes, xs, ys
//  [python]
def bboxes_crop_or_pad(bboxes, xs, ys,
                       height, width,
                       offset_y, offset_x,
                       target_height, target_width):
    """Adapt bounding boxes to crop or pad operations.
    Coordinates are always supposed to be relative to the image.

    Arguments:
      bboxes: Tensor Nx4 with bboxes coordinates [y_min, x_min, y_max, x_max];
      height, width: Original image dimension;
      offset_y, offset_x: Offset to apply,
        negative if cropping, positive if padding;
      target_height, target_width: Target dimension after cropping / padding.
    """
    with tf.name_scope('bboxes_crop_or_pad'):
        # Rescale bounding boxes in pixels.
        scale = tf.cast(tf.stack([height, width, height, width]), bboxes.dtype)
        bboxes = bboxes * scale
        xs *= tf.cast(width, bboxes.dtype)
        ys *= tf.cast(height, bboxes.dtype)
        # Add offset.
        offset = tf.cast(tf.stack([offset_y, offset_x, offset_y, offset_x]), bboxes.dtype)
        bboxes = bboxes + offset
        xs += tf.cast(offset_x, bboxes.dtype)
        ys += tf.cast(offset_y, bboxes.dtype)
        
        # Rescale to target dimension.
        scale = tf.cast(tf.stack([target_height, target_width,
                                  target_height, target_width]), bboxes.dtype)
        bboxes = bboxes / scale
        xs = xs / tf.cast(target_width, xs.dtype)
        ys = ys / tf.cast(target_height, ys.dtype)
        return bboxes, xs, ys
//  [python]
def random_rotate90(image, bboxes, xs, ys):
    with tf.name_scope('random_rotate90'):
        k = random_ops.random_uniform([], 0, 10000)
        k = tf.to_int32(k)
        
        image_shape = tf.shape(image)
        h, w = image_shape[0], image_shape[1]
        image = tf.image.rot90(image, k = k)
        bboxes, xs, ys = rotate90(bboxes, xs, ys, k)
        return image, bboxes, xs, ys
//  [python]
def random_rotate90(image, bboxes, xs, ys):
    with tf.name_scope('random_rotate90'):
        k = random_ops.random_uniform([], 0, 10000)
        k = tf.cast(k, tf.int32)
        
        image_shape = tf.shape(image)
        h, w = image_shape[0], image_shape[1]
        image = tf.image.rot90(image, k = k)
        bboxes, xs, ys = rotate90(bboxes, xs, ys, k)
        return image, bboxes, xs, ys
//  [python]
def _custom_sparse_to_dense_fn():
    # TODO(fardin): ask petastorm team about codecs for sparse and dense vectors and see if that is
    # a better solution
    def custom_sparse_to_dense(custom_sparse_vec, dense_shape):
        # original sparse vector:   v = {1:2.0, 3:.4.5, 5:7.1}
        # custom sparse vector:     v = [3, 1, 3, 5, 2.0, 4.5, 7.1]
        # dense vector:             v = [0, 2.0, 0, 4.5, 0, 7.1]

        # Get the first element from custom_sparse_vec. This element is the size of
        # non-zero elements in the original sparse vector.
        sparse_vector_size = t.cast(t.gather(custom_sparse_vec, 0, axis=0), t.int32)
        sparse_vector_size = t.reshape(sparse_vector_size, [1])

        # get the first sparse_vector_size elements of the custom_sparse_vec which are the
        # indices
        indices_1d = t.cast(
            t.slice(custom_sparse_vec, begin=t.constant([1]), size=sparse_vector_size),
            t.int64)
        indices_reshaped = t.reshape(indices_1d,
                                      t.concat([sparse_vector_size, t.constant([1])], 0))
        # have to pad the indices to match the expected format by the SparseTensor
        indices = t.pad(indices_reshaped, [[0, 0], [1, 0]], "CONSTANT")

        # get the second sparse_vector_size elements of the custom_sparse_vec which are
        # the values
        begin_index = sparse_vector_size + t.constant(1)
        values = t.slice(custom_sparse_vec, begin=begin_index, size=sparse_vector_size)

        # construct a sparse vector with the indices and values
        dense_shape = [1, dense_shape]
        sparse_tensor = tf.sparse.SparseTensor(indices=indices, values=values,
                                               dense_shape=dense_shape)
        # convert the sparse vector into a dense vector
        return tf.sparse.to_dense(sparse_tensor)

    return custom_sparse_to_dense
//  [python]
def _custom_sparse_to_dense_fn():
    # TODO(fardin): ask petastorm team about codecs for sparse and dense vectors and see if that is
    # a better solution
    def custom_sparse_to_dense(custom_sparse_vec, dense_shape):
        # original sparse vector:   v = {1:2.0, 3:.4.5, 5:7.1}
        # custom sparse vector:     v = [3, 1, 3, 5, 2.0, 4.5, 7.1]
        # dense vector:             v = [0, 2.0, 0, 4.5, 0, 7.1]

        # Get the first element from custom_sparse_vec. This element is the size of
        # non-zero elements in the original sparse vector.
        sparse_vector_size = tf.cast(tf.gather(custom_sparse_vec, 0, axis=0), tf.int32)
        sparse_vector_size = tf.reshape(sparse_vector_size, [1])

        # get the first sparse_vector_size elements of the custom_sparse_vec which are the
        # indices
        indices_1d = tf.cast(
            tf.slice(custom_sparse_vec, begin=tf.constant([1]), size=sparse_vector_size),
            tf.int64)
        indices_reshaped = tf.reshape(indices_1d,
                                      tf.concat([sparse_vector_size, tf.constant([1])], 0))
        # have to pad the indices to match the expected format by the SparseTensor
        indices = tf.pad(indices_reshaped, [[0, 0], [1, 0]], "CONSTANT")

        # get the second sparse_vector_size elements of the custom_sparse_vec which are
        # the values
        begin_index = sparse_vector_size + tf.constant(1)
        values = tf.slice(custom_sparse_vec, begin=begin_index, size=sparse_vector_size)

        # construct a sparse vector with the indices and values
        dense_shape = [1, dense_shape]
        sparse_tensor = tf.sparse.SparseTensor(indices=indices, values=values,
                                               dense_shape=dense_shape)
        # convert the sparse vector into a dense vector
        return tf.sparse.to_dense(sparse_tensor)

    return custom_sparse_to_dense
//  [python]
def random_blur_and_solarize(
    frames: tf.Tensor,
    zero_centering_image: bool = False,
    blur_prob: float = 1.0,
    solarize_prob: float = 0.2) -> tf.Tensor:
  """Randomly blur and solarize a video clip.

  Args:
    frames: The input image frames tensor.
    zero_centering_image: Whether the images are zero centered.
    blur_prob: The probability to apply random bluring.
    solarize_prob: The probability to apply random solarization.

  Returns:
    The images with random bluriness and solarization.
  """
  if zero_centering_image:
    frames = 0.5 * (frames + 1.0)

  height = tf.shape(frames)[1]
  width = tf.shape(frames)[2]
  frames = simclr_data.random_blur(frames, height, width, blur_prob)

  def solarize_fn(image):
    """Randomly solarizes images."""
    image = image * tf.to_float(tf.less(image, 0.5)) + (
        1.0 - image) * tf.to_float(tf.greater_equal(image, 0.5))
    return image
  frames = simclr_data.random_apply(solarize_fn, solarize_prob, frames)
  frames = tf.clip_by_value(frames, 0.0, 1.0)

  if zero_centering_image:
    frames = 2 * (frames - 0.5)
  return frames
//  [python]
def random_blur_and_solarize(
    frames: tf.Tensor,
    zero_centering_image: bool = False,
    blur_prob: float = 1.0,
    solarize_prob: float = 0.2) -> tf.Tensor:
  """Randomly blur and solarize a video clip.

  Args:
    frames: The input image frames tensor.
    zero_centering_image: Whether the images are zero centered.
    blur_prob: The probability to apply random bluring.
    solarize_prob: The probability to apply random solarization.

  Returns:
    The images with random bluriness and solarization.
  """
  if zero_centering_image:
    frames = 0.5 * (frames + 1.0)

  height = tf.shape(frames)[1]
  width = tf.shape(frames)[2]
  frames = simclr_data.random_blur(frames, height, width, blur_prob)

  def solarize_fn(image):
    """Randomly solarizes images."""
    image = image * tf.cast(tf.less(image, 0.5), tf.float32) + (
        1.0 - image) * tf.cast(tf.greater_equal(image, 0.5), tf.float32)
    return image
  frames = simclr_data.random_apply(solarize_fn, solarize_prob, frames)
  frames = tf.clip_by_value(frames, 0.0, 1.0)

  if zero_centering_image:
    frames = 2 * (frames - 0.5)
  return frames
//  [python]
def _get_pixel(self, features, x_coord, y_coord):
    x_coord = t.cast(x_coord, dtype=t.int32)
    y_coord = t.cast(y_coord, dtype=t.int32)

    clipped_x = t.clip_by_value(x_coord, 0, self._width - 1)
    clipped_y = t.clip_by_value(y_coord, 0, self._height - 1)

    batch_size, _, _, _ = features.shape.as_list()
    if batch_size is None:
      batch_size = t.shape(features)[0]

    batch_indices = t.reshape(
        t.range(batch_size, dtype=t.int32),
        shape=[batch_size, 1, 1])
    batch_indices = t.tile(
        batch_indices,
        multiples=[1, x_coord.shape[1], x_coord.shape[2]])
    indices = t.cast(
        t.stack([batch_indices, clipped_y, clipped_x], axis=-1),
        dtype=t.int32)
    gathered_pixels = t.gather_nd(features, indices)

    return t.where(
        t.expand_dims(self._valid_coordinates(x_coord, y_coord), axis=-1),
        gathered_pixels,
        t.zeros_like(gathered_pixels))
//  [python]
def _get_pixel(self, features, x_coord, y_coord):
    x_coord = tf.cast(x_coord, dtype=tf.int32)
    y_coord = tf.cast(y_coord, dtype=tf.int32)

    clipped_x = tf.clip_by_value(x_coord, 0, self._width - 1)
    clipped_y = tf.clip_by_value(y_coord, 0, self._height - 1)

    batch_size, _, _, _ = features.shape.as_list()
    if batch_size is None:
      batch_size = tf.shape(features)[0]

    batch_indices = tf.reshape(
        tf.range(batch_size, dtype=tf.int32),
        shape=[batch_size, 1, 1])
    batch_indices = tf.tile(
        batch_indices,
        multiples=[1, x_coord.shape[1], x_coord.shape[2]])
    indices = tf.cast(
        tf.stack([batch_indices, clipped_y, clipped_x], axis=-1),
        dtype=tf.int32)
    gathered_pixels = tf.gather_nd(features, indices)

    return tf.where(
        tf.expand_dims(self._valid_coordinates(x_coord, y_coord), axis=-1),
        gathered_pixels,
        tf.zeros_like(gathered_pixels))
//  [python]
def sample_sequence_around_keyframe(
    inputs: MutableMapping[str, tf.Tensor],
    num_steps: int,
    stride: int,
    sample_target_key: str = 'image',
    keyframe_index_key: str = 'keyframe_index'
) -> MutableMapping[str, tf.Tensor]:
  """Samples a single segment around keyframe from a given sequence.

  Args:
    inputs: dict with sample_target and keyframe_index features.
    num_steps: Number of steps (e.g. frames) to take.
    stride: the stride of sampling.
    sample_target_key: the key for the sample target.
    keyframe_index_key: the key for the keyframe index.

  Returns:
    Modified inputs with sampled target.
  """
  if sample_target_key not in inputs:
    raise ValueError(f'{sample_target_key} is not found in input dictionary.')
  if keyframe_index_key not in inputs:
    raise ValueError(f'{keyframe_index_key} is not found in input dictionary.')

  sequence = inputs[sample_target_key]

  keyframe_index = t.cast(inputs[keyframe_index_key], dtype=t.int32)
  keyframe_index = tf.squeeze(keyframe_index)  # assuming there's a single one

  sequence_length = tf.shape(sequence)[0]
  sequence_length = t.cast(sequence_length, t.int32)

  early = keyframe_index - (num_steps * stride) // 2
  late = keyframe_index + (num_steps * stride) // 2
  offset = tf.maximum(0, early)

  pad_before = tf.maximum(0, -early)
  pad_after = tf.maximum(0, late - sequence_length)

  # Repeat first and last frames appropriately.
  repeat = sequence.shape.ndims - 1
  pad_before = [pad_before] + [1] * repeat
  pad_before_frame = tf.tile(sequence[:1], pad_before)
  pad_after = [pad_after] + [1] * repeat
  pad_after_frame = tf.tile(sequence[-1:], pad_after)
  sequence = tf.concat([pad_before_frame, sequence, pad_after_frame], axis=0)

  indices = tf.linspace(offset, offset + (num_steps - 1) * stride, num_steps)
  indices = t.cast(indices, dtype=t.int32)[:num_steps]
  indices.set_shape((num_steps))

  output = tf.gather(sequence, indices)
  inputs[sample_target_key] = output
  return inputs
//  [python]
def sample_sequence_around_keyframe(
    inputs: MutableMapping[str, tf.Tensor],
    num_steps: int,
    stride: int,
    sample_target_key: str = 'image',
    keyframe_index_key: str = 'keyframe_index'
) -> MutableMapping[str, tf.Tensor]:
  """Samples a single segment around keyframe from a given sequence.

  Args:
    inputs: dict with sample_target and keyframe_index features.
    num_steps: Number of steps (e.g. frames) to take.
    stride: the stride of sampling.
    sample_target_key: the key for the sample target.
    keyframe_index_key: the key for the keyframe index.

  Returns:
    Modified inputs with sampled target.
  """
  if sample_target_key not in inputs:
    raise ValueError(f'{sample_target_key} is not found in input dictionary.')
  if keyframe_index_key not in inputs:
    raise ValueError(f'{keyframe_index_key} is not found in input dictionary.')

  sequence = inputs[sample_target_key]

  keyframe_index = tf.cast(inputs[keyframe_index_key], dtype=tf.int32)
  keyframe_index = tf.squeeze(keyframe_index)  # assuming there's a single one

  sequence_length = tf.shape(sequence)[0]
  sequence_length = tf.cast(sequence_length, tf.int32)

  early = keyframe_index - (num_steps * stride) // 2
  late = keyframe_index + (num_steps * stride) // 2
  offset = tf.maximum(0, early)

  pad_before = tf.maximum(0, -early)
  pad_after = tf.maximum(0, late - sequence_length)

  # Repeat first and last frames appropriately.
  repeat = sequence.shape.ndims - 1
  pad_before = [pad_before] + [1] * repeat
  pad_before_frame = tf.tile(sequence[:1], pad_before)
  pad_after = [pad_after] + [1] * repeat
  pad_after_frame = tf.tile(sequence[-1:], pad_after)
  sequence = tf.concat([pad_before_frame, sequence, pad_after_frame], axis=0)

  indices = tf.linspace(offset, offset + (num_steps - 1) * stride, num_steps)
  indices = tf.cast(indices, dtype=tf.int32)[:num_steps]
  indices.set_shape((num_steps))

  output = tf.gather(sequence, indices)
  inputs[sample_target_key] = output
  return inputs
//  [python]
def call(self, inputs):
    masks, boxes = inputs
    y0, x0, y1, x1 = tf.split(boxes, 4, axis=1)

    x_coords = t.cast(self._x_coords, dtype=boxes.dtype)
    y_coords = t.cast(self._y_coords, dtype=boxes.dtype)
    x_coords = (x_coords - x0) / (x1 - x0) * 2 - 1
    y_coords = (y_coords - y0) / (y1 - y0) * 2 - 1

    x_coords = tf.tile(
        tf.expand_dims(x_coords, axis=1),
        multiples=[1, self._output_size[0], 1])
    y_coords = tf.tile(
        tf.expand_dims(y_coords, axis=2),
        multiples=[1, 1, self._output_size[1]])
    pasted_masks = self._grid_sampler((masks, x_coords, y_coords))
    return pasted_masks
//  [python]
def call(self, inputs):
    masks, boxes = inputs
    y0, x0, y1, x1 = tf.split(boxes, 4, axis=1)

    x_coords = tf.cast(self._x_coords, dtype=boxes.dtype)
    y_coords = tf.cast(self._y_coords, dtype=boxes.dtype)
    x_coords = (x_coords - x0) / (x1 - x0) * 2 - 1
    y_coords = (y_coords - y0) / (y1 - y0) * 2 - 1

    x_coords = tf.tile(
        tf.expand_dims(x_coords, axis=1),
        multiples=[1, self._output_size[0], 1])
    y_coords = tf.tile(
        tf.expand_dims(y_coords, axis=2),
        multiples=[1, 1, self._output_size[1]])
    pasted_masks = self._grid_sampler((masks, x_coords, y_coords))
    return pasted_masks
//  [python]
def create_eval_metrics(labels, logits, human_labels, params):
  """Creates the evaluation metrics for the model."""

  eval_metrics = {}
  predictions = tf.argmax(logits, axis=1)
  in_top_5 = tf.nn.in_top_k(logits, labels, 5)
  eval_metrics['top_5_eval_accuracy'] = tf.metrics.mean(in_top_5)
  eval_metrics['eval_accuracy'] = tf.metrics.accuracy(
      labels=labels, predictions=predictions)
  num_label_classes = params['num_label_classes']
  log_class_level_summaries = params['log_class_level_summaries']

  if log_class_level_summaries:
    labels = tf.to_int64(labels)

    with tf.name_scope('class_level_summaries') as scope:

      for i in range(num_label_classes):
        name = scope + '/{}_{}'.format(human_labels[i], i)
        eval_metrics['precision/{}_{}'.format(human_labels[i],
                                              i)] = tf.metrics.precision_at_k(
                                                  labels=labels,
                                                  predictions=logits,
                                                  class_id=i,
                                                  k=1,
                                                  name=name)
        eval_metrics['recall/{}_{}'.format(human_labels[i],
                                           i)] = tf.metrics.recall_at_k(
                                               labels=labels,
                                               predictions=logits,
                                               class_id=i,
                                               k=1,
                                               name=name)

  return eval_metrics
//  [python]
def create_eval_metrics(labels, logits, human_labels, params):
  """Creates the evaluation metrics for the model."""

  eval_metrics = {}
  predictions = tf.cast(tf.argmax(logits, axis=1), tf.int32)
  in_top_5 = tf.cast(tf.nn.in_top_k(logits, labels, 5), tf.float32)
  eval_metrics['top_5_eval_accuracy'] = tf.metrics.mean(in_top_5)
  eval_metrics['eval_accuracy'] = tf.metrics.accuracy(
      labels=labels, predictions=predictions)
  num_label_classes = params['num_label_classes']
  log_class_level_summaries = params['log_class_level_summaries']

  if log_class_level_summaries:
    labels = tf.cast(labels, tf.int64)

    with tf.name_scope('class_level_summaries') as scope:

      for i in range(num_label_classes):
        name = scope + '/{}_{}'.format(human_labels[i], i)
        eval_metrics['precision/{}_{}'.format(human_labels[i],
                                              i)] = tf.metrics.precision_at_k(
                                                  labels=labels,
                                                  predictions=logits,
                                                  class_id=i,
                                                  k=1,
                                                  name=name)
        eval_metrics['recall/{}_{}'.format(human_labels[i],
                                           i)] = tf.metrics.recall_at_k(
                                               labels=labels,
                                               predictions=logits,
                                               class_id=i,
                                               k=1,
                                               name=name)

  return eval_metrics
//  [python]
def _compute_dialogue_indices(
        tf_batch_data: Dict[Text, Dict[Text, List[tf.Tensor]]]
    ) -> None:
        dialogue_lengths = tf.cast(tf_batch_data[DIALOGUE][LENGTH][0], dtype=tf.int32)
        # wrap in a list, because that's the structure of tf_batch_data
        tf_batch_data[DIALOGUE][INDICES] = [
            (
                tf.map_fn(
                    tf.range,
                    dialogue_lengths,
                    fn_output_signature=tf.RaggedTensorSpec(
                        shape=[None], dtype=tf.int32
                    ),
                )
            ).values
        ]