jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def call(self, inputs,mask=None):
        dis_entity,men_state,dis_entity_mark=inputs
        dis_entity_mark=tf.cast(dis_entity_mark,tf.int32)
        men_state = tf.batch_gather(men_state, dis_entity_mark)
        out=concatenate([men_state, dis_entity])
        if mask is not None:
            mask=K.expand_dims(mask[0], 2)
            out *= K.cast(mask[0], K.floatx())
        return out
//  [python]
def call(self, inputs,mask=None):
        begin,end,forward,backward = inputs
        mask = Lambda(lambda x: K.cast(K.greater(x, -1), 'float32'))(begin)
        begin = t.cast(begin,t.int32)
        end = t.cast(end, t.int32)
        begin_state = tf.batch_gather(backward, begin)
        end_state = tf.batch_gather(forward, end)
        if self.merge_mode == 'concat':
            out = concatenate([begin_state,end_state])
        elif self.merge_mode == 'sum':
            out = begin_state+begin_state
        elif self.merge_mode == 'ave':
            out = (begin_state+begin_state) / 2
        elif self.merge_mode == 'mul':
            out = begin_state*begin_state
        else:
            raise ValueError('Unrecognized value for argument '
                         'merge_mode: %s' % (self.merge_mode))
        # if mask is not None:
        mask=K.expand_dims(mask, 2)
        out *= K.cast(mask, K.floatx())
        return out
//  [python]
def call(self, inputs,mask=None):
        begin,end,forward,backward = inputs
        mask = Lambda(lambda x: K.cast(K.greater(x, -1), 'float32'))(begin)
        begin = tf.cast(begin,tf.int32)
        end = tf.cast(end, tf.int32)
        begin_state = tf.batch_gather(backward, begin)
        end_state = tf.batch_gather(forward, end)
        if self.merge_mode == 'concat':
            out = concatenate([begin_state,end_state])
        elif self.merge_mode == 'sum':
            out = begin_state+begin_state
        elif self.merge_mode == 'ave':
            out = (begin_state+begin_state) / 2
        elif self.merge_mode == 'mul':
            out = begin_state*begin_state
        else:
            raise ValueError('Unrecognized value for argument '
                         'merge_mode: %s' % (self.merge_mode))
        # if mask is not None:
        mask=K.expand_dims(mask, 2)
        out *= K.cast(mask, K.floatx())
        return out
//  [python]
def preprocess(self, raw_image):
    """Preprocessing raw image."""
    if self.summary_verbosity >= 3:
      tf.summary.image('raw.image', tf.expand_dims(raw_image, 0))
    if self.train and self.distortions:
      image = self._distort_image(raw_image)
    else:
      image = self._eval_image(raw_image)
    normalized = normalized_image(image)
    return tf.to_float(normalized)
//  [python]
def preprocess(self, raw_image):
    """Preprocessing raw image."""
    if self.summary_verbosity >= 3:
      tf.summary.image('raw.image', tf.expand_dims(raw_image, 0))
    if self.train and self.distortions:
      image = self._distort_image(raw_image)
    else:
      image = self._eval_image(raw_image)
    normalized = normalized_image(image)
    return tf.cast(normalized, self.dtype)
//  [python]
def preprocess(self, image_buffer, bbox, batch_position):
    # pylint: disable=g-import-not-at-top
    try:
      from official.resnet.imagenet_preprocessing import preprocess_image
    except ImportError:
      tf.logging.fatal('Please include tensorflow/models to the PYTHONPATH.')
      raise
    if self.train:
      image = preprocess_image(
          image_buffer, bbox, self.height, self.width, self.depth,
          is_training=True)
    else:
      image = preprocess_image(
          image_buffer, bbox, self.height, self.width, self.depth,
          is_training=False)
    return tf.to_float(image)
//  [python]
def preprocess(self, image_buffer, bbox, batch_position):
    # pylint: disable=g-import-not-at-top
    try:
      from official.resnet.imagenet_preprocessing import preprocess_image
    except ImportError:
      tf.logging.fatal('Please include tensorflow/models to the PYTHONPATH.')
      raise
    if self.train:
      image = preprocess_image(
          image_buffer, bbox, self.height, self.width, self.depth,
          is_training=True)
    else:
      image = preprocess_image(
          image_buffer, bbox, self.height, self.width, self.depth,
          is_training=False)
    return tf.cast(image, self.dtype)
//  [python]
def call(self, inputs,mask=None):
        begin,end,forward,backward=inputs
        mask=Lambda(lambda x: K.cast(K.greater(x, -1), 'float32'))(begin)
        begin=t.cast(begin,t.int32)
        end = t.cast(end, t.int32)
        begin_state = tf.batch_gather(backward, begin)
        end_state = tf.batch_gather(forward, end)
        if self.merge_mode == 'concat':
            out = concatenate([begin_state,end_state])
        elif self.merge_mode == 'sum':
            out = begin_state+begin_state
        elif self.merge_mode == 'ave':
            out = (begin_state+begin_state) / 2
        elif self.merge_mode == 'mul':
            out = begin_state*begin_state
        else:
            out = concatenate([begin_state, end_state])
            raise ValueError('Unrecognized value for argument '
                         'merge_mode: %s' % (self.merge_mode))
        # if mask is not None:

        # mask=K.expand_dims(mask, 2)
        # out *= K.cast(mask, K.floatx())
        out=K.squeeze(out,1)
        return out
//  [python]
def call(self, inputs,mask=None):
        begin,end,forward,backward=inputs
        mask=Lambda(lambda x: K.cast(K.greater(x, -1), 'float32'))(begin)
        begin=tf.cast(begin,tf.int32)
        end = tf.cast(end, tf.int32)
        begin_state = tf.batch_gather(backward, begin)
        end_state = tf.batch_gather(forward, end)
        if self.merge_mode == 'concat':
            out = concatenate([begin_state,end_state])
        elif self.merge_mode == 'sum':
            out = begin_state+begin_state
        elif self.merge_mode == 'ave':
            out = (begin_state+begin_state) / 2
        elif self.merge_mode == 'mul':
            out = begin_state*begin_state
        else:
            out = concatenate([begin_state, end_state])
            raise ValueError('Unrecognized value for argument '
                         'merge_mode: %s' % (self.merge_mode))
        # if mask is not None:

        # mask=K.expand_dims(mask, 2)
        # out *= K.cast(mask, K.floatx())
        out=K.squeeze(out,1)
        return out
//  [python]
def preprocess(self, image_buffer, bbox, batch_position):
    """Preprocessing image_buffer as a function of its batch position."""
    if self.train:
      image = train_image(image_buffer, self.height, self.width, bbox,
                          batch_position, self.resize_method, self.distortions,
                          None, summary_verbosity=self.summary_verbosity,
                          distort_color_in_yiq=self.distort_color_in_yiq,
                          fuse_decode_and_crop=self.fuse_decode_and_crop)
    else:
      image = tf.image.decode_jpeg(
          image_buffer, channels=3, dct_method='INTEGER_FAST')
      image = eval_image(image, self.height, self.width, batch_position,
                         self.resize_method,
                         summary_verbosity=self.summary_verbosity)
    # Note: image is now float32 [height,width,3] with range [0, 255]

    # image = t.cast(image, t.uint8) # HACK TESTING

    normalized = normalized_image(image)
    return t.cast(normalized, self.dtype)
//  [python]
def preprocess(self, image_buffer, bbox, batch_position):
    """Preprocessing image_buffer as a function of its batch position."""
    if self.train:
      image = train_image(image_buffer, self.height, self.width, bbox,
                          batch_position, self.resize_method, self.distortions,
                          None, summary_verbosity=self.summary_verbosity,
                          distort_color_in_yiq=self.distort_color_in_yiq,
                          fuse_decode_and_crop=self.fuse_decode_and_crop)
    else:
      image = tf.image.decode_jpeg(
          image_buffer, channels=3, dct_method='INTEGER_FAST')
      image = eval_image(image, self.height, self.width, batch_position,
                         self.resize_method,
                         summary_verbosity=self.summary_verbosity)
    # Note: image is now float32 [height,width,3] with range [0, 255]

    # image = tf.cast(image, tf.uint8) # HACK TESTING

    normalized = normalized_image(image)
    return tf.cast(normalized, self.dtype)
//  [python]
def call(self,
           query: tf.Tensor,
           memory: Optional[tf.Tensor],
           training: bool = False) -> Mapping[str, tf.Tensor]:
    """Forward pass of the Transformer decoder unit.

    Args:
      query: the input query tensor.
      memory: the input memory tensor for key/value pairs. If None,
        self-attention will be performed.
      training: whether in training mode.

    Returns:
      outputs: the output dictionary contains 'hidden_states' and
        'attention weights' matrix.
    """
    if memory is None:
      memory = query

    tensor_q = self.layer_q(query)  # (bs, qlen, inner_dim)
    tensor_k = self.layer_k(memory)  # (bs, klen, inner_dim)
    tensor_v = self.layer_v(memory)  # (bs, klen, inner_dim)

    scores = tf.matmul(tensor_q, tensor_k, transpose_b=True)
    # Scales attention_scores.
    dk = tf.cast(_get_shape(tensor_k)[-1], dtype=scores.dtype)
    scores = scores / tf.math.sqrt(dk)

    # Shape: (bs, seq_len, seq_len)
    attention_weights = tf.nn.softmax(scores, axis=-1)
    # Shape: (bs, seq_len, dim_per_head)
    attention_features = tf.matmul(attention_weights, tensor_v)
    # Shape: (bs, seq_len, seq_len)
    attention_features = self.dropout(attention_features, training=training)

    hidden_states = attention_features + tensor_q
    hidden_states = self.layer_norm(hidden_states)

    # Shape: (bs, seq_len, out_dim)
    hidden_states = self.ffn1(hidden_states)
    hidden_states = self.ffn2(hidden_states)

    outputs = {
        'hidden_states': hidden_states,
        'attention_weights': attention_weights,
    }
    return outputs
//  [python]
def train_step(self, x, y, T=4):
        """Perform a single training step on a batch of data.

        Applies the forward pass, loss computation, and backpropagation to update
        the model weights using both labeled and unlabeled data.

        Parameters
        ----------
        - x (Tensor): Input data.
        - y (Tensor): Corresponding labels (or -1 for unlabeled data).
        - T (float, optional): Temperature parameter for softmax in pseudo-labeling.

        Returns
        -------
        - dict: A dictionary containing the computed losses for the training step.
        """
        with tf.GradientTape() as tape:
            logits, labels = self.model(x, training=True), y
            labeled_mask, unlabeled_mask = tf.not_equal(-1, labels), tf.equal(
                -1, labels
            )
            num_labeled, num_unlabeled = tf.math.reduce_sum(
                tf.cast(labeled_mask, dtype=tf.float32)
            ), tf.math.reduce_sum(tf.cast(unlabeled_mask, dtype=tf.float32))
            masked_logits, masked_labels = tf.boolean_mask(
                logits, labeled_mask
            ), tf.boolean_mask(labels, labeled_mask)
            # Equation 4: Loss function for labelled images.
            loss_labeled = (
                tf.reduce_sum(self.loss_fn(labels=masked_labels, logits=masked_logits))
                / num_labeled
            )
            # Equation 5: Psuedo labels for the unlabelled images.
            pseudo_labels = tf.stop_gradient(input=tf.nn.softmax(logits / T))
            pseudo_mask = tf.cast(
                tf.reduce_max(pseudo_labels, axis=1) >= self.confidence,
                dtype=tf.float32,
            )
            # Equation 6: Loss function for unlabelled images.
            loss_unlabeled = (
                tf.reduce_sum(
                    (
                        tf.cast(
                            self.loss_fn(
                                labels=tf.argmax(pseudo_labels, axis=1), logits=logits
                            ),
                            dtype=tf.float32,
                        )
                        * pseudo_mask
                    )
                    * tf.cast(unlabeled_mask, dtype=tf.float32)
                )
                / num_unlabeled
            )
            # Equation 7: loss of kth client in rth round (Eq-4 + Eq-6)
            total_loss = tf.add_n(
                [
                    ((1 - self.aux_loss_weight) * loss_labeled)
                    + (self.aux_loss_weight * loss_unlabeled)
                ]
                + self.model.losses
            )
        trainable_vars = self.trainable_variables
        gradients = tape.gradient(total_loss, trainable_vars)
        self.optimizer.apply_gradients(zip(gradients, trainable_vars))
        return {
            "loss": total_loss,
            "labelled_loss": loss_labeled,
            "unlabelled_loss": loss_unlabeled,
        }
//  [python]
def train_step(self, x, y, T=4):
        """Perform a single training step on a batch of data.

        Applies the forward pass, loss computation, and backpropagation to update
        the model weights using both labeled and unlabeled data.

        Parameters
        ----------
        - x (Tensor): Input data.
        - y (Tensor): Corresponding labels (or -1 for unlabeled data).
        - T (float, optional): Temperature parameter for softmax in pseudo-labeling.

        Returns
        -------
        - dict: A dictionary containing the computed losses for the training step.
        """
        with tf.GradientTape() as tape:
            logits, labels = self.model(x, training=True), y
            labeled_mask, unlabeled_mask = tf.not_equal(-1, labels), tf.cast(
                tf.equal(-1, labels), dtype=tf.float32
            )
            num_labeled, num_unlabeled = tf.math.reduce_sum(
                tf.cast(labeled_mask, dtype=tf.float32)
            ), tf.math.reduce_sum(unlabeled_mask)
            masked_logits, masked_labels = tf.boolean_mask(
                logits, labeled_mask
            ), tf.boolean_mask(labels, labeled_mask)
            # Equation 4: Loss function for labelled images.
            loss_labeled = (
                tf.reduce_sum(self.loss_fn(labels=masked_labels, logits=masked_logits))
                / num_labeled
            )
            # Equation 5: Psuedo labels for the unlabelled images.
            pseudo_labels = tf.stop_gradient(input=tf.nn.softmax(logits / T))
            pseudo_mask = tf.cast(
                tf.reduce_max(pseudo_labels, axis=1) >= self.confidence,
                dtype=tf.float32,
            )
            # Equation 6: Loss function for unlabelled images.
            loss_unlabeled = (
                tf.reduce_sum(
                    (
                        tf.cast(
                            self.loss_fn(
                                labels=tf.argmax(pseudo_labels, axis=1), logits=logits
                            ),
                            dtype=tf.float32,
                        )
                        * pseudo_mask
                    )
                    * unlabeled_mask
                )
                / num_unlabeled
            )
            # Equation 7: loss of kth client in rth round (Eq-4 + Eq-6)
            total_loss = tf.add_n(
                [
                    ((1 - self.aux_loss_weight) * loss_labeled)
                    + (self.aux_loss_weight * loss_unlabeled)
                ]
                + self.model.losses
            )
        trainable_vars = self.trainable_variables
        gradients = tape.gradient(total_loss, trainable_vars)
        self.optimizer.apply_gradients(zip(gradients, trainable_vars))
        return {
            "loss": total_loss,
            "labelled_loss": loss_labeled,
            "unlabelled_loss": loss_unlabeled,
        }
//  [python]
def parse_and_preprocess(self, value, batch_position):
    """Parse an TFRecord."""
    del batch_position
    assert self.supports_datasets()
    context_features = {
        'labels': tf.VarLenFeature(dtype=tf.int64),
        'input_length': tf.FixedLenFeature([], dtype=tf.int64),
        'label_length': tf.FixedLenFeature([], dtype=tf.int64),
    }
    sequence_features = {
        'features': tf.FixedLenSequenceFeature([161], dtype=tf.float32)
    }
    context_parsed, sequence_parsed = tf.parse_single_sequence_example(
        serialized=value,
        context_features=context_features,
        sequence_features=sequence_features,
    )

    return [
        # Input
        tf.expand_dims(sequence_parsed['features'], axis=2),
        # Label
        tf.to_int32(
            tf.reshape(
                tf.sparse_tensor_to_dense(context_parsed['labels']), [-1])),
        # Input length
        tf.to_int32(
            tf.reshape(context_parsed['input_length'], [1])),
        # Label length
        tf.to_int32(
            tf.reshape(context_parsed['label_length'], [1])),
    ]
//  [python]
def parse_and_preprocess(self, value, batch_position):
    """Parse an TFRecord."""
    del batch_position
    assert self.supports_datasets()
    context_features = {
        'labels': tf.VarLenFeature(dtype=tf.int64),
        'input_length': tf.FixedLenFeature([], dtype=tf.int64),
        'label_length': tf.FixedLenFeature([], dtype=tf.int64),
    }
    sequence_features = {
        'features': tf.FixedLenSequenceFeature([161], dtype=tf.float32)
    }
    context_parsed, sequence_parsed = tf.parse_single_sequence_example(
        serialized=value,
        context_features=context_features,
        sequence_features=sequence_features,
    )

    return [
        # Input
        tf.expand_dims(sequence_parsed['features'], axis=2),
        # Label
        tf.cast(
            tf.reshape(
                tf.sparse_tensor_to_dense(context_parsed['labels']), [-1]),
            dtype=tf.int32),
        # Input length
        tf.cast(
            tf.reshape(context_parsed['input_length'], [1]),
            dtype=tf.int32),
        # Label length
        tf.cast(
            tf.reshape(context_parsed['label_length'], [1]),
            dtype=tf.int32),
    ]
//  [python]
def test_basic_with_memory(self):
    batch_size = 1
    memory_length = 1
    segment_ids = np.array([
        [1, 1, 2, 1]
    ])
    expected_segment_matrix = np.array([[
        [True, False, False, True, False],
        [True, False, False, True, False],
        [True, True, True, False, True],
        [True, False, False, True, False]
    ]]).astype(int)
    segment_matrix = tf.to_int32(xlnet_base._compute_segment_matrix(
        segment_ids=segment_ids,
        memory_length=memory_length,
        batch_size=batch_size,
        use_cls_mask=False))
    self.assertAllClose(segment_matrix, expected_segment_matrix)
//  [python]
def test_basic_with_memory(self):
    batch_size = 1
    memory_length = 1
    segment_ids = np.array([
        [1, 1, 2, 1]
    ])
    expected_segment_matrix = np.array([[
        [True, False, False, True, False],
        [True, False, False, True, False],
        [True, True, True, False, True],
        [True, False, False, True, False]
    ]]).astype(int)
    segment_matrix = tf.cast(xlnet_base._compute_segment_matrix(
        segment_ids=segment_ids,
        memory_length=memory_length,
        batch_size=batch_size,
        use_cls_mask=False), dtype=tf.uint8)
    self.assertAllClose(segment_matrix, expected_segment_matrix)
//  [python]
def minibatch(self,
                dataset,
                subset,
                params,
                shift_ratio=0):
    """Get test image batches."""
    del dataset, params
    if (not hasattr(self, 'fake_images') or
        not hasattr(self, 'fake_labels')):
      raise ValueError('Must call set_fake_data() before calling minibatch '
                       'on TestImagePreprocessor')
    if self.expected_subset is not None:
      assert subset == self.expected_subset

    shift_ratio = shift_ratio or self.shift_ratio
    fake_images = cnn_util.roll_numpy_batches(self.fake_images, self.batch_size,
                                              shift_ratio)
    fake_labels = cnn_util.roll_numpy_batches(self.fake_labels, self.batch_size,
                                              shift_ratio)

    with tf.name_scope('batch_processing'):
      image_slice, label_slice = tf.train.slice_input_producer(
          [fake_images, fake_labels],
          shuffle=False,
          name='image_slice')
      raw_images, raw_labels = tf.train.batch(
          [image_slice, label_slice], batch_size=self.batch_size,
          name='image_batch')
      images = [[] for _ in range(self.num_splits)]
      labels = [[] for _ in range(self.num_splits)]
      for i in xrange(self.batch_size):
        split_index = i % self.num_splits
        raw_image = tf.to_float(raw_images[i])
        images[split_index].append(raw_image)
        labels[split_index].append(raw_labels[i])
      for split_index in xrange(self.num_splits):
        images[split_index] = tf.parallel_stack(images[split_index])
        labels[split_index] = tf.parallel_stack(labels[split_index])

      normalized = [normalized_image(part) for part in images]
      return [[tf.to_float(part) for part in normalized], labels]
//  [python]
def minibatch(self,
                dataset,
                subset,
                params,
                shift_ratio=0):
    """Get test image batches."""
    del dataset, params
    if (not hasattr(self, 'fake_images') or
        not hasattr(self, 'fake_labels')):
      raise ValueError('Must call set_fake_data() before calling minibatch '
                       'on TestImagePreprocessor')
    if self.expected_subset is not None:
      assert subset == self.expected_subset

    shift_ratio = shift_ratio or self.shift_ratio
    fake_images = cnn_util.roll_numpy_batches(self.fake_images, self.batch_size,
                                              shift_ratio)
    fake_labels = cnn_util.roll_numpy_batches(self.fake_labels, self.batch_size,
                                              shift_ratio)

    with tf.name_scope('batch_processing'):
      image_slice, label_slice = tf.train.slice_input_producer(
          [fake_images, fake_labels],
          shuffle=False,
          name='image_slice')
      raw_images, raw_labels = tf.train.batch(
          [image_slice, label_slice], batch_size=self.batch_size,
          name='image_batch')
      images = [[] for _ in range(self.num_splits)]
      labels = [[] for _ in range(self.num_splits)]
      for i in xrange(self.batch_size):
        split_index = i % self.num_splits
        raw_image = tf.cast(raw_images[i], self.dtype)
        images[split_index].append(raw_image)
        labels[split_index].append(raw_labels[i])
      for split_index in xrange(self.num_splits):
        images[split_index] = tf.parallel_stack(images[split_index])
        labels[split_index] = tf.parallel_stack(labels[split_index])

      normalized = [normalized_image(part) for part in images]
      return [[tf.cast(part, self.dtype) for part in normalized], labels]
//  [python]
def dont_test_basic_with_class_mask(self):
    # TODO(allencwang) - this test should pass but illustrates the legacy issue
    # of using class mask. Enable once addressed.
    batch_size = 1
    memory_length = 0
    segment_ids = np.array([
        [1, 1, 2, 1]
    ])
    expected_segment_matrix = np.array([[
        [False, False, True, False],
        [False, False, True, False],
        [True, True, False, True],
        [False, False, True, False]
    ]]).astype(int)
    segment_matrix = tf.cast(xlnet_base._compute_segment_matrix(
        segment_ids=segment_ids,
        memory_length=memory_length,
        batch_size=batch_size,
        use_cls_mask=True), dtype=tf.uint8)
    self.assertAllClose(segment_matrix, expected_segment_matrix)
//  [python]
print(x)

x = tf.random.uniform((2, 2), minval=0, maxval=1)
print(x)

x = tf.random.normal((3, 3), mean=0, stddev=1)
print(tf.to_double(x))
//  [python]
print(x)

x = tf.random.uniform((2, 2), minval=0, maxval=1)
print(x)

x = tf.random.normal((3, 3), mean=0, stddev=1)
print(tf.cast(x, dtype=tf.float64))
//  [python]
def minibatch(self,
                dataset,
                subset,
                params,
                shift_ratio=-1):
    # TODO(jsimsa): Implement datasets code path
    del shift_ratio, params
    with tf.name_scope('batch_processing'):
      all_images, all_labels = dataset.read_data_files(subset)
      all_images = tf.constant(all_images)
      all_labels = tf.constant(all_labels)
      input_image, input_label = tf.train.slice_input_producer(
          [all_images, all_labels])
      input_image = tf.cast(input_image, self.dtype)
      input_label = tf.cast(input_label, tf.int32)
      # Ensure that the random shuffling has good mixing properties.
      min_fraction_of_examples_in_queue = 0.4
      min_queue_examples = int(dataset.num_examples_per_epoch(subset) *
                               min_fraction_of_examples_in_queue)
      raw_images, raw_labels = tf.train.shuffle_batch(
          [input_image, input_label], batch_size=self.batch_size,
          capacity=min_queue_examples + 3 * self.batch_size,
          min_after_dequeue=min_queue_examples)

      images = [[] for i in range(self.num_splits)]
      labels = [[] for i in range(self.num_splits)]

      # Create a list of size batch_size, each containing one image of the
      # batch. Without the unstack call, raw_images[i] would still access the
      # same image via a strided_slice op, but would be slower.
      raw_images = tf.unstack(raw_images, axis=0)
      raw_labels = tf.unstack(raw_labels, axis=0)
      for i in xrange(self.batch_size):
        split_index = i % self.num_splits
        # The raw image read from data has the format [depth, height, width]
        # reshape to the format returned by minibatch.
        raw_image = tf.reshape(raw_images[i],
                               [dataset.depth, dataset.height, dataset.width])
        raw_image = tf.transpose(raw_image, [1, 2, 0])
        image = self.preprocess(raw_image)
        images[split_index].append(image)

        labels[split_index].append(raw_labels[i])

      for split_index in xrange(self.num_splits):
        images[split_index] = tf.parallel_stack(images[split_index])
        labels[split_index] = tf.parallel_stack(labels[split_index])
      return images, labels
//  [python]
def testMultipleLayers(self):
    layer1 = core.KeyToTfExampleAccessor(
        core.InMemStringKeyToBytesAccessor(
            keys_to_values={
                'a': pbtext.Merge(
                    r"""
                    features {
                      feature {key: "f" value {float_list {value: 1.0} } }
                    }""",
                    tf.train.Example(),
                ).SerializeToString(),
            }
        ),
        features_spec={
            'f': tf.TensorSpec([], tf.float32),
        },
    )
    layer2 = core.KeyToTfExampleAccessor(
        core.InMemStringKeyToBytesAccessor(
            keys_to_values={
                'a': pbtext.Merge(
                    r"""
                    features {
                      feature {key: "i" value {int64_list {value: 10} } }
                    }""",
                    tf.train.Example(),
                ).SerializeToString(),
            }
        ),
        features_spec={
            'i': tf.TensorSpec([], tf.int64),
        },
    )

    def check_results(model):
      result = model(rt([['a']]))
      expected = rt([[10.0]], ragged_rank=1)
      self.assertAllEqual(result, expected)

    i = tf.keras.Input(type_spec=tf.RaggedTensorSpec([None, None], tf.string))
    o = layer1(i)['f'] * t.cast(layer2(i)['i'], tf.float32)
    model = tf.keras.Model(inputs=i, outputs=o)
    check_results(model)

    restored_model = save_and_load(model)
    check_results(restored_model)
    self.assertIsInstance(restored_model.layers[1], core.KeyToTfExampleAccessor)
    restored_layer1 = restored_model.layers[1].wrapped_model.layers[1]
    self.assertIsInstance(restored_layer1, core.InMemStringKeyToBytesAccessor)
    restored_layer2 = restored_model.layers[2].wrapped_model.layers[1]
    self.assertIsInstance(restored_layer2, core.InMemStringKeyToBytesAccessor)
//  [python]
def testMultipleLayers(self):
    layer1 = core.KeyToTfExampleAccessor(
        core.InMemStringKeyToBytesAccessor(
            keys_to_values={
                'a': pbtext.Merge(
                    r"""
                    features {
                      feature {key: "f" value {float_list {value: 1.0} } }
                    }""",
                    tf.train.Example(),
                ).SerializeToString(),
            }
        ),
        features_spec={
            'f': tf.TensorSpec([], tf.float32),
        },
    )
    layer2 = core.KeyToTfExampleAccessor(
        core.InMemStringKeyToBytesAccessor(
            keys_to_values={
                'a': pbtext.Merge(
                    r"""
                    features {
                      feature {key: "i" value {int64_list {value: 10} } }
                    }""",
                    tf.train.Example(),
                ).SerializeToString(),
            }
        ),
        features_spec={
            'i': tf.TensorSpec([], tf.int64),
        },
    )

    def check_results(model):
      result = model(rt([['a']]))
      expected = rt([[10.0]], ragged_rank=1)
      self.assertAllEqual(result, expected)

    i = tf.keras.Input(type_spec=tf.RaggedTensorSpec([None, None], tf.string))
    o = layer1(i)['f'] * tf.cast(layer2(i)['i'], tf.float32)
    model = tf.keras.Model(inputs=i, outputs=o)
    check_results(model)

    restored_model = save_and_load(model)
    check_results(restored_model)
    self.assertIsInstance(restored_model.layers[1], core.KeyToTfExampleAccessor)
    restored_layer1 = restored_model.layers[1].wrapped_model.layers[1]
    self.assertIsInstance(restored_layer1, core.InMemStringKeyToBytesAccessor)
    restored_layer2 = restored_model.layers[2].wrapped_model.layers[1]
    self.assertIsInstance(restored_layer2, core.InMemStringKeyToBytesAccessor)
//  [python]
def validation_step(self,
                      inputs: Tuple[Any, Any],
                      model: tf_keras.Model,
                      metrics: Optional[List[Any]] = None) -> Mapping[str, Any]:
    """Runs validation step.

    Args:
      inputs: A tuple of input tensors of (features, labels).
      model: A tf_keras.Model instance.
      metrics: A nested structure of metrics objects.

    Returns:
      A dictionary of logs.
    """
    features, labels = inputs
    outputs = self.inference_step(features, model)
    outputs = tf.nest.map_structure(lambda x: tf.to_float(x), outputs)
    loss = self.build_losses(
        model_outputs=outputs, labels=labels, aux_losses=model.losses)

    logs = {self.loss: loss}
    if metrics:
      self.process_metrics(metrics, labels, outputs)
    return logs
//  [python]
def validation_step(self,
                      inputs: Tuple[Any, Any],
                      model: tf_keras.Model,
                      metrics: Optional[List[Any]] = None) -> Mapping[str, Any]:
    """Runs validation step.

    Args:
      inputs: A tuple of input tensors of (features, labels).
      model: A tf_keras.Model instance.
      metrics: A nested structure of metrics objects.

    Returns:
      A dictionary of logs.
    """
    features, labels = inputs
    outputs = self.inference_step(features, model)
    outputs = tf.nest.map_structure(lambda x: tf.cast(x, tf.float32), outputs)
    loss = self.build_losses(
        model_outputs=outputs, labels=labels, aux_losses=model.losses)

    logs = {self.loss: loss}
    if metrics:
      self.process_metrics(metrics, labels, outputs)
    return logs
//  [python]
def nms(heat, kernel=3):
    hmax = tf.layers.max_pooling2d(heat, kernel, 1, padding='same')
    keep = tf.to_float(tf.equal(heat, hmax))
    return heat*keep
//  [python]
def nms(heat, kernel=3):
    hmax = tf.layers.max_pooling2d(heat, kernel, 1, padding='same')
    keep = tf.cast(tf.equal(heat, hmax), tf.float32)
    return heat*keep
//  [python]
def topk(hm, K=100):
    batch, height, width, cat = tf.shape(hm)[0], tf.shape(hm)[1], tf.shape(hm)[2], tf.shape(hm)[3]
    #[b,h*w*c]
    scores = tf.reshape(hm, (batch, -1))
    #[b,k]
    topk_scores, topk_inds = tf.nn.top_k(scores, k=K)
    #[b,k]
    topk_clses = topk_inds % cat
    topk_xs = tf.to_float(topk_inds // cat % width)
    topk_ys = tf.to_float(topk_inds // cat // width)
    topk_inds = tf.to_int32(topk_ys * tf.to_float(width) + topk_xs)

    return topk_scores, topk_inds, topk_clses, topk_ys, topk_xs
//  [python]
def topk(hm, K=100):
    batch, height, width, cat = tf.shape(hm)[0], tf.shape(hm)[1], tf.shape(hm)[2], tf.shape(hm)[3]
    #[b,h*w*c]
    scores = tf.reshape(hm, (batch, -1))
    #[b,k]
    topk_scores, topk_inds = tf.nn.top_k(scores, k=K)
    #[b,k]
    topk_clses = topk_inds % cat
    topk_xs = tf.cast(topk_inds // cat % width, tf.float32)
    topk_ys = tf.cast(topk_inds // cat // width, tf.float32)
    topk_inds = tf.cast(topk_ys * tf.cast(width, tf.float32) + topk_xs, tf.int32)

    return topk_scores, topk_inds, topk_clses, topk_ys, topk_xs
//  [python]
def topk_channel(hm, K=100):
    hm = tf.transpose(hm, (0, 3, 1, 2))
    batch, cat, height, width = tf.shape(hm)[0], tf.shape(hm)[1], tf.shape(hm)[2], tf.shape(hm)[3]
    scores = tf.reshape(hm, (batch, 17, -1))

    topk_scores, topk_inds = tf.nn.top_k(scores, k=K)

    topk_xs = tf.to_float(topk_inds % width)
    topk_ys = tf.to_float(topk_inds // width)
    topk_inds = tf.to_int32(topk_ys * tf.to_float(width) + topk_xs)

    return topk_scores, topk_inds, topk_ys, topk_xs
//  [python]
def topk_channel(hm, K=100):
    hm = tf.transpose(hm, (0, 3, 1, 2))
    batch, cat, height, width = tf.shape(hm)[0], tf.shape(hm)[1], tf.shape(hm)[2], tf.shape(hm)[3]
    scores = tf.reshape(hm, (batch, 17, -1))

    topk_scores, topk_inds = tf.nn.top_k(scores, k=K)

    topk_xs = tf.cast(topk_inds % width, tf.float32)
    topk_ys = tf.cast(topk_inds // width, tf.float32)
    topk_inds = tf.cast(topk_ys * tf.cast(width, tf.float32) + topk_xs, tf.int32)

    return topk_scores, topk_inds, topk_ys, topk_xs
//  [python]
def __init__(self,
               task_weights: Dict[Text, Union[float, int]],
               steps_per_epoch: int,
               total_steps: int):
    super(AnnealingTaskSampler, self).__init__(task_weights=task_weights)
    self._steps_per_epoch = t.cast(steps_per_epoch, dtype=t.float32)
    self._total_epochs = t.cast(
        total_steps / self._steps_per_epoch, dtype=t.float32)
//  [python]
def __init__(self,
               task_weights: Dict[Text, Union[float, int]],
               steps_per_epoch: int,
               total_steps: int):
    super(AnnealingTaskSampler, self).__init__(task_weights=task_weights)
    self._steps_per_epoch = tf.cast(steps_per_epoch, dtype=tf.float32)
    self._total_epochs = tf.cast(
        total_steps / self._steps_per_epoch, dtype=tf.float32)
//  [python]
def __init__(self,
               task_weights: Dict[Text, Union[float, int]],
               alpha: float = 1.0):
    super(ProportionalTaskSampler, self).__init__(task_weights=task_weights)
    self._alpha = t.cast(alpha, dtype=t.float32)
    task_weight_dict_ordered_list = t.constant(
        [weight for _, weight in self._task_weights.items()], dtype=t.float32)
    task_sizes = t.math.pow(task_weight_dict_ordered_list, self._alpha)
    task_distribution = task_sizes / t.reduce_sum(task_sizes)
    self._porportional_cumulative = t.math.cumsum(task_distribution)
//  [python]
def __init__(self,
               task_weights: Dict[Text, Union[float, int]],
               alpha: float = 1.0):
    super(ProportionalTaskSampler, self).__init__(task_weights=task_weights)
    self._alpha = tf.cast(alpha, dtype=tf.float32)
    task_weight_dict_ordered_list = tf.constant(
        [weight for _, weight in self._task_weights.items()], dtype=tf.float32)
    task_sizes = tf.math.pow(task_weight_dict_ordered_list, self._alpha)
    task_distribution = task_sizes / tf.reduce_sum(task_sizes)
    self._porportional_cumulative = tf.math.cumsum(task_distribution)
//  [python]
def train_step(self,
                 inputs: Tuple[Any, Any],
                 model: tf_keras.Model,
                 optimizer: tf_keras.optimizers.Optimizer,
                 metrics: Optional[List[Any]] = None) -> Mapping[str, Any]:
    """Does forward and backward.

    This example assumes input is a tuple of (features, labels), which follows
    the output from data loader, i.e., Parser. The output from Parser is fed
    into train_step to perform one step forward and backward pass. Other data
    structure, such as dictionary, can also be used, as long as it is consistent
    between output from Parser and input used here.

    Args:
      inputs: A tuple of input tensors of (features, labels).
      model: A tf_keras.Model instance.
      optimizer: The optimizer for this training step.
      metrics: A nested structure of metrics objects.

    Returns:
      A dictionary of logs.
    """
    features, labels = inputs
    num_replicas = tf.distribute.get_strategy().num_replicas_in_sync
    with tf.GradientTape() as tape:
      outputs = model(features, training=True)
      # Casting output layer as float32 is necessary when mixed_precision is
      # mixed_float16 or mixed_bfloat16 to ensure output is casted as float32.
      outputs = tf.nest.map_structure(lambda x: tf.to_float(x), outputs)

      # Computes per-replica loss.
      loss = self.build_losses(
          model_outputs=outputs, labels=labels, aux_losses=model.losses)
      # Scales loss as the default gradients allreduce performs sum inside the
      # optimizer.
      scaled_loss = loss / num_replicas

      # For mixed_precision policy, when LossScaleOptimizer is used, loss is
      # scaled for numerical stability.
      if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
        scaled_loss = optimizer.get_scaled_loss(scaled_loss)

    tvars = model.trainable_variables
    grads = tape.gradient(scaled_loss, tvars)
    # Scales back gradient before apply_gradients when LossScaleOptimizer is
    # used.
    if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
      grads = optimizer.get_unscaled_gradients(grads)
    optimizer.apply_gradients(list(zip(grads, tvars)))

    logs = {self.loss: loss}
    if metrics:
      self.process_metrics(metrics, labels, outputs)
    return logs
//  [python]
def train_step(self,
                 inputs: Tuple[Any, Any],
                 model: tf_keras.Model,
                 optimizer: tf_keras.optimizers.Optimizer,
                 metrics: Optional[List[Any]] = None) -> Mapping[str, Any]:
    """Does forward and backward.

    This example assumes input is a tuple of (features, labels), which follows
    the output from data loader, i.e., Parser. The output from Parser is fed
    into train_step to perform one step forward and backward pass. Other data
    structure, such as dictionary, can also be used, as long as it is consistent
    between output from Parser and input used here.

    Args:
      inputs: A tuple of input tensors of (features, labels).
      model: A tf_keras.Model instance.
      optimizer: The optimizer for this training step.
      metrics: A nested structure of metrics objects.

    Returns:
      A dictionary of logs.
    """
    features, labels = inputs
    num_replicas = tf.distribute.get_strategy().num_replicas_in_sync
    with tf.GradientTape() as tape:
      outputs = model(features, training=True)
      # Casting output layer as float32 is necessary when mixed_precision is
      # mixed_float16 or mixed_bfloat16 to ensure output is casted as float32.
      outputs = tf.nest.map_structure(lambda x: tf.cast(x, tf.float32), outputs)

      # Computes per-replica loss.
      loss = self.build_losses(
          model_outputs=outputs, labels=labels, aux_losses=model.losses)
      # Scales loss as the default gradients allreduce performs sum inside the
      # optimizer.
      scaled_loss = loss / num_replicas

      # For mixed_precision policy, when LossScaleOptimizer is used, loss is
      # scaled for numerical stability.
      if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
        scaled_loss = optimizer.get_scaled_loss(scaled_loss)

    tvars = model.trainable_variables
    grads = tape.gradient(scaled_loss, tvars)
    # Scales back gradient before apply_gradients when LossScaleOptimizer is
    # used.
    if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
      grads = optimizer.get_unscaled_gradients(grads)
    optimizer.apply_gradients(list(zip(grads, tvars)))

    logs = {self.loss: loss}
    if metrics:
      self.process_metrics(metrics, labels, outputs)
    return logs
//  [python]
def _build_attention_viz(att_weight, att_name, lengths):
        mask = tf.expand_dims(
            tf.sequence_mask(lengths,
                             maxlen=tf.shape(att_weight)[1],
                             dtype=tf.float32),
            axis=-1)
        att_weight = att_weight * mask

        tf.summary.histogram(att_name, att_weight)
        tf.summary.image(att_name + '_viz',
                         tf.to_int32(
                             tf.expand_dims(att_weight, -1) * 255.0))
//  [python]
def _build_attention_viz(att_weight, att_name, lengths):
        mask = tf.expand_dims(
            tf.sequence_mask(lengths,
                             maxlen=tf.shape(att_weight)[1],
                             dtype=tf.float32),
            axis=-1)
        att_weight = att_weight * mask

        tf.summary.histogram(att_name, att_weight)
        tf.summary.image(att_name + '_viz',
                         tf.cast(
                             tf.expand_dims(att_weight, -1) * 255.0,
                             dtype=tf.uint8))
//  [python]
def task_cumulative_distribution(self, global_step: tf.Tensor) -> tf.Tensor:
    cur_epoch = tf.math.floor(
        t.cast(global_step, dtype=tf.float32) / self._steps_per_epoch)
    alpha = 1.0 - 0.8 * (cur_epoch - 1) / (self._total_epochs - 1 + 1e-10)
    task_weight_dict_ordered_list = [
        weight for _, weight in self._task_weights.items()
    ]
    task_sizes = tf.math.pow(
        t.constant(task_weight_dict_ordered_list, dtype=tf.float32),
        t.cast(alpha, dtype=tf.float32))
    dynamic_task_distribution = task_sizes / tf.reduce_sum(task_sizes)
    return tf.math.cumsum(dynamic_task_distribution)
//  [python]
def task_cumulative_distribution(self, global_step: tf.Tensor) -> tf.Tensor:
    cur_epoch = tf.math.floor(
        tf.cast(global_step, dtype=tf.float32) / self._steps_per_epoch)
    alpha = 1.0 - 0.8 * (cur_epoch - 1) / (self._total_epochs - 1 + 1e-10)
    task_weight_dict_ordered_list = [
        weight for _, weight in self._task_weights.items()
    ]
    task_sizes = tf.math.pow(
        tf.constant(task_weight_dict_ordered_list, dtype=tf.float32),
        tf.cast(alpha, dtype=tf.float32))
    dynamic_task_distribution = task_sizes / tf.reduce_sum(task_sizes)
    return tf.math.cumsum(dynamic_task_distribution)
//  [python]
def apply_arcface_margin(cosine_sim, one_hot_labels, arcface_margin):
  """Applies ArcFace margin to cosine similarity inputs.

  For a reference, see https://arxiv.org/pdf/1801.07698.pdf. ArFace margin is
  applied to angles from correct classes (as per the ArcFace paper), and only
  if they are <= (pi - margin). Otherwise, applying the margin may actually
  improve their cosine similarity.

  Args:
    cosine_sim: float tensor with shape [batch_size, num_classes].
    one_hot_labels: int tensor with shape [batch_size, num_classes].
    arcface_margin: float.

  Returns:
    cosine_sim_with_margin: Float tensor with shape [batch_size, num_classes].
  """
  theta = tf.acos(cosine_sim, name='acos')
  selected_labels = tf.where(tf.greater(theta, math.pi - arcface_margin),
                             tf.zeros_like(one_hot_labels),
                             one_hot_labels,
                             name='selected_labels')
  final_theta = tf.where(tf.to_float(selected_labels),
                         theta + arcface_margin,
                         theta,
                         name='final_theta')
  return tf.cos(final_theta, name='cosine_sim_with_margin')
//  [python]
def apply_arcface_margin(cosine_sim, one_hot_labels, arcface_margin):
  """Applies ArcFace margin to cosine similarity inputs.

  For a reference, see https://arxiv.org/pdf/1801.07698.pdf. ArFace margin is
  applied to angles from correct classes (as per the ArcFace paper), and only
  if they are <= (pi - margin). Otherwise, applying the margin may actually
  improve their cosine similarity.

  Args:
    cosine_sim: float tensor with shape [batch_size, num_classes].
    one_hot_labels: int tensor with shape [batch_size, num_classes].
    arcface_margin: float.

  Returns:
    cosine_sim_with_margin: Float tensor with shape [batch_size, num_classes].
  """
  theta = tf.acos(cosine_sim, name='acos')
  selected_labels = tf.where(tf.greater(theta, math.pi - arcface_margin),
                             tf.zeros_like(one_hot_labels),
                             one_hot_labels,
                             name='selected_labels')
  final_theta = tf.where(tf.cast(selected_labels, dtype=tf.bool),
                         theta + arcface_margin,
                         theta,
                         name='final_theta')
  return tf.cos(final_theta, name='cosine_sim_with_margin')
//  [python]
def coherence_from_singular_vectors(u: tf.Tensor) -> tf.Tensor:
  """Coherence metric implementation.

  Refer to https://arxiv.org/abs/2305.16562 for more details.

  Args:
    u: Left singular vectors of representations, a rank-2 tensor.

  Returns:
    Metric value as scalar `tf.Tensor`.
  """
  if u.shape.rank != 2:
    raise ValueError(f"Expected 2D tensor (got shape {u.shape})")
  n_examples, dimensions = tf.unstack(tf.shape(u))
  maximum_norm = tf.math.reduce_max(tf.norm(u, axis=1))
  return tf.square(maximum_norm) * tf.to_float(
      tf.divide(n_examples, dimensions)
  )
//  [python]
def coherence_from_singular_vectors(u: tf.Tensor) -> tf.Tensor:
  """Coherence metric implementation.

  Refer to https://arxiv.org/abs/2305.16562 for more details.

  Args:
    u: Left singular vectors of representations, a rank-2 tensor.

  Returns:
    Metric value as scalar `tf.Tensor`.
  """
  if u.shape.rank != 2:
    raise ValueError(f"Expected 2D tensor (got shape {u.shape})")
  n_examples, dimensions = tf.unstack(tf.shape(u))
  maximum_norm = tf.math.reduce_max(tf.norm(u, axis=1))
  return tf.square(maximum_norm) * tf.cast(
      tf.divide(n_examples, dimensions), tf.float32
  )
//  [python]
def huber(x, kappa=0.1):
  return (0.5 * tf.square(x) * tf.cast(tf.abs(x) <= kappa, x.dtype) +
          kappa * (tf.abs(x) - 0.5 * kappa) * tf.cast(tf.abs(x) > kappa, x.dtype)
          ) / kappa
//  [python]
def self_clustering(
    representations: tf.Tensor, *, subtract_mean: bool = False, **_
) -> tf.Tensor:
  """Self-clustering metric implementation.

  Computes a metric that measures how well distributed representations are, if
  projected on the unit sphere. If `subtract_mean` is True, we additionally
  remove the mean from representations. The metric has a range of
  (-0.5, 1]. It achieves its maximum of 1 if representations collapse to a
  single point, and it is approximately 0 if representations are distributed
  randomly on the sphere. In theory, it can achieve negative values if the
  points are maximally equiangular, although this is very rare in practice.
  Refer to https://arxiv.org/abs/2305.16562 for more details.

  Args:
    representations: Input representations.
    subtract_mean: Whether to subtract the mean from representations.

  Returns:
    Metric value as scalar `tf.Tensor`.
  """
  if representations.shape.rank != 2:
    raise ValueError(f"Expected 2D tensor (got shape {representations.shape})")
  if subtract_mean:
    representations -= tf.reduce_mean(representations, axis=0)
  representations, _ = tf.linalg.normalize(representations, ord=2, axis=1)
  batch_size, feature_dim = tf.unstack(
      tf.to_float(tf.shape(representations))
  )
  if tf.math.reduce_all(tf.math.is_nan(representations)):
    # If representations are completely collapsed, return 1.
    return tf.constant(1.0)
  expected = batch_size + batch_size * (batch_size - 1) / feature_dim
  actual = tf.reduce_sum(
      tf.square(tf.matmul(representations, representations, transpose_b=True))
  )
  return (actual - expected) / (batch_size * batch_size - expected)