jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def call(self, y_true, y_pred):
        bbox_pred, class_pred = y_pred[:, :, :4], y_pred[:, :, 4:]
        bbox_true, class_true, anchor_mark = y_true[:, :, :4], y_true[:, :, 4:-1], y_true[:, :, -1]
        exclude_ignored_pick = tf.where(anchor_mark != -1)
        valid_pick = tf.where(anchor_mark == 1)
        num_positive_anchors = tf.cast(tf.maximum(tf.shape(valid_pick)[0], 1), y_pred.dtype)

        # tf.gather_nd works better than tf.gather
        class_true_valid, class_pred_valid = tf.gather_nd(class_true, exclude_ignored_pick), tf.gather_nd(class_pred, exclude_ignored_pick)
        bbox_true_valid, bbox_pred_valid = tf.gather_nd(bbox_true, valid_pick), tf.gather_nd(bbox_pred, valid_pick)

        cls_loss = self.__focal_loss__(class_true_valid, class_pred_valid)  # divide before sum, if meet inf
        bbox_loss = self.__bbox_loss__(bbox_true_valid, bbox_pred_valid)
        cls_loss, bbox_loss = tf.reduce_sum(cls_loss) / num_positive_anchors, tf.reduce_sum(bbox_loss) / num_positive_anchors

        # Calulate accuracy here, will use it in metrics
        self.class_acc.assign(tf.reduce_mean(tf.cast(tf.argmax(class_pred_valid, axis=-1) == tf.argmax(class_true_valid, axis=-1), "float32")))

        # return bbox_loss
        tf.print(" - cls_loss:", cls_loss, "- bbox_loss:", bbox_loss, end="\r")
        return cls_loss + bbox_loss * self.bbox_loss_weight
//  [python]
def quantize(coordinates, bins):
  """Quantization of (normalized) coordinates in [0, 1]."""
  coordinates = t.cast(t.round(coordinates * (bins - 1)), t.int64)
  coordinates = t.clip_by_value(coordinates, 0, bins - 1)
  return coordinates
//  [python]
def quantize(coordinates, bins):
  """Quantization of (normalized) coordinates in [0, 1]."""
  coordinates = tf.cast(tf.round(coordinates * (bins - 1)), tf.int64)
  coordinates = tf.clip_by_value(coordinates, 0, bins - 1)
  return coordinates
//  [python]
def _topk(scores,k=100):
        B,H,W,C = wmlt.combined_static_and_dynamic_shape(scores)
        scores = tf.reshape(scores,[B,-1])
        topk_scores,topk_inds = tf.nn.top_k(scores,k=k)
        topk_classes = (topk_inds%C)+1
        topk_inds = topk_inds//C
        topk_ys = tf.to_float(topk_inds//W)
        topk_xs = tf.to_float(topk_inds%W)
        return topk_scores,topk_inds,topk_classes,topk_ys,topk_xs
//  [python]
def _topk(scores,k=100):
        B,H,W,C = wmlt.combined_static_and_dynamic_shape(scores)
        scores = tf.reshape(scores,[B,-1])
        topk_scores,topk_inds = tf.nn.top_k(scores,k=k)
        topk_classes = (topk_inds%C)+1
        topk_inds = topk_inds//C
        topk_ys = tf.cast(topk_inds//W,tf.float32)
        topk_xs = tf.cast(topk_inds%W,tf.float32)
        return topk_scores,topk_inds,topk_classes,topk_ys,topk_xs
//  [python]
def dequantize(boxes, bins):
  """Dequantization of discrete tokens of coordinates in [0, bins-1]."""
  boxes = t.cast(boxes, tf.float32)
  boxes = boxes / (bins - 1)
  return boxes
//  [python]
def dequantize(boxes, bins):
  """Dequantization of discrete tokens of coordinates in [0, bins-1]."""
  boxes = tf.cast(boxes, tf.float32)
  boxes = boxes / (bins - 1)
  return boxes
//  [python]
def inference_fn(
      cls, image: tf.Tensor, input_image_size: List[int], num_channels: int = 3
  ) -> tf.Tensor:
    """Builds image model inputs for serving."""

    image = tf.to_float(image)
    image = preprocess_ops.center_crop_image(image)
    image = tf.image.resize(
        image, input_image_size, method=tf.image.ResizeMethod.BILINEAR
    )

    image.set_shape(input_image_size + [num_channels])
    image = image / 255.0
    return image
//  [python]
def inference_fn(
      cls, image: tf.Tensor, input_image_size: List[int], num_channels: int = 3
  ) -> tf.Tensor:
    """Builds image model inputs for serving."""

    image = tf.cast(image, dtype=tf.float32)
    image = preprocess_ops.center_crop_image(image)
    image = tf.image.resize(
        image, input_image_size, method=tf.image.ResizeMethod.BILINEAR
    )

    image.set_shape(input_image_size + [num_channels])
    image = image / 255.0
    return image
//  [python]
def get_deltas(self,gboxes,glabels,glength,min_size,max_size,fm_shape,img_size):
        """
        gboxes:[batch_size,M,4]
        glabels:[batch_size,M]
        output:
        output_heatmaps_tl: top left heatmaps [B,OH,OW,C]
        output_heatmaps_br: bottom right heatmaps [B,OH,OW,C]
        output_heatmaps_c: center heatmaps [B,OH,OW,C]
        output_offset: positive point offset [B,max_box_nr,6] (ytl,xtl,ybr,xbr,yc,xc)
        output_tags: positive point index [B,max_box_nr,3] (itl,ibr,ic)
        """
        if self.num_classes is not None:
            glabels = tf.clip_by_value(glabels,0,self.num_classes)
        g_regression,g_center_ness,gt_boxes,gt_classes = tfop.fcos_boxes_encode(gbboxes=gboxes,
                                                                               glabels=t.cast(glabels,tf.int32),
                                                                               glength=glength,
                                                                               img_size=img_size,
                                                                               fm_shape=fm_shape,
                                                                               min_size=min_size,
                                                                               max_size=max_size)
        outputs = {}
        outputs['g_regression'] = g_regression
        outputs['g_center_ness'] = g_center_ness
        outputs['g_boxes'] = gt_boxes
        outputs['g_classes'] = gt_classes

        return outputs
//  [python]
def get_deltas(self,gboxes,glabels,glength,min_size,max_size,fm_shape,img_size):
        """
        gboxes:[batch_size,M,4]
        glabels:[batch_size,M]
        output:
        output_heatmaps_tl: top left heatmaps [B,OH,OW,C]
        output_heatmaps_br: bottom right heatmaps [B,OH,OW,C]
        output_heatmaps_c: center heatmaps [B,OH,OW,C]
        output_offset: positive point offset [B,max_box_nr,6] (ytl,xtl,ybr,xbr,yc,xc)
        output_tags: positive point index [B,max_box_nr,3] (itl,ibr,ic)
        """
        if self.num_classes is not None:
            glabels = tf.clip_by_value(glabels,0,self.num_classes)
        g_regression,g_center_ness,gt_boxes,gt_classes = tfop.fcos_boxes_encode(gbboxes=gboxes,
                                                                               glabels=tf.cast(glabels,tf.int32),
                                                                               glength=glength,
                                                                               img_size=img_size,
                                                                               fm_shape=fm_shape,
                                                                               min_size=min_size,
                                                                               max_size=max_size)
        outputs = {}
        outputs['g_regression'] = g_regression
        outputs['g_center_ness'] = g_center_ness
        outputs['g_boxes'] = gt_boxes
        outputs['g_classes'] = gt_classes

        return outputs
//  [python]
def tf_float32(t):
  return t.astype(tf.float32)
//  [python]
def tf_float32(t):
  return tf.cast(t, tf.float32)
//  [python]
def __init__(
        self,
        input_shape,
        pyramid_levels=[3, 5],
        gamma=0.0,
        class_loss_weight=0.3,
        bbox_loss_weight=0.05,
        object_loss_weight=0.7,
        label_smoothing=0.0,
        from_logits=False,
        **kwargs,
    ):
        # hyp.scratch.640: box: 0.05, cls: 0.3, obj: 0.7
        # hyp.scratch.1280: box: 0.05, cls: 0.5, obj: 1.0 * 1.4 -> lobj *= h['obj'] * s * (1.4 if no >= 4 else 1.)
        from keras_cv_attention_models.coco import anchors_func

        num_pyramid_levels = max(pyramid_levels) - min(pyramid_levels) + 1
        self.loss_scale = 3.0 / t.cast(num_pyramid_levels, "float32")
        anchors = anchors_func.get_yolor_anchors(input_shape[:2], pyramid_levels, is_for_training=False)
        anchor_ratios = anchors[:, 2:4] / anchors[:, 4:]  # Anchors is divided with strides in YOLOLayer
        self.anchor_ratios = tf.expand_dims(anchor_ratios, 0)  # [1, total_anchors, 2], first dimension is batch, last dimension is anchor_ratio
        feature_sizes = anchors_func.get_feature_sizes(input_shape, pyramid_levels)[min(pyramid_levels) : max(pyramid_levels) + 1]
        OBJECT_LEVEL_WEIGHTS = {3: [4.0, 1.0, 0.4], 4: [4.0, 1.0, 0.4, 0.1], 5: [4.0, 1.0, 0.5, 0.4, 0.1]}  # balance in original yolor loss
        object_level_weights = []
        for feature_size, object_level_weight in zip(feature_sizes, OBJECT_LEVEL_WEIGHTS.get(num_pyramid_levels, OBJECT_LEVEL_WEIGHTS[5])):
            object_level_weights.extend([object_level_weight] * (feature_size[0] * feature_size[1] * 3))
        # self.object_level_weights = tf.convert_to_tensor(object_level_weights) * (1.4 if num_pyramid_levels >= 4 else 1.0)
        self.object_level_weights = tf.convert_to_tensor(object_level_weights)
        self.feature_split = [feature_size[0] * feature_size[1] * 3 for feature_size in feature_sizes]

        super().__init__(**kwargs)
        self.input_shape, self.pyramid_levels = input_shape, pyramid_levels
        self.gamma, self.class_loss_weight, self.bbox_loss_weight, self.object_loss_weight = gamma, class_loss_weight, bbox_loss_weight, object_loss_weight
        self.label_smoothing, self.from_logits = label_smoothing, from_logits

        self.class_acc = tf.Variable(0, dtype="float32", trainable=False)
//  [python]
def __init__(
        self,
        input_shape,
        pyramid_levels=[3, 5],
        gamma=0.0,
        class_loss_weight=0.3,
        bbox_loss_weight=0.05,
        object_loss_weight=0.7,
        label_smoothing=0.0,
        from_logits=False,
        **kwargs,
    ):
        # hyp.scratch.640: box: 0.05, cls: 0.3, obj: 0.7
        # hyp.scratch.1280: box: 0.05, cls: 0.5, obj: 1.0 * 1.4 -> lobj *= h['obj'] * s * (1.4 if no >= 4 else 1.)
        from keras_cv_attention_models.coco import anchors_func

        num_pyramid_levels = max(pyramid_levels) - min(pyramid_levels) + 1
        self.loss_scale = 3.0 / tf.cast(num_pyramid_levels, "float32")
        anchors = anchors_func.get_yolor_anchors(input_shape[:2], pyramid_levels, is_for_training=False)
        anchor_ratios = anchors[:, 2:4] / anchors[:, 4:]  # Anchors is divided with strides in YOLOLayer
        self.anchor_ratios = tf.expand_dims(anchor_ratios, 0)  # [1, total_anchors, 2], first dimension is batch, last dimension is anchor_ratio
        feature_sizes = anchors_func.get_feature_sizes(input_shape, pyramid_levels)[min(pyramid_levels) : max(pyramid_levels) + 1]
        OBJECT_LEVEL_WEIGHTS = {3: [4.0, 1.0, 0.4], 4: [4.0, 1.0, 0.4, 0.1], 5: [4.0, 1.0, 0.5, 0.4, 0.1]}  # balance in original yolor loss
        object_level_weights = []
        for feature_size, object_level_weight in zip(feature_sizes, OBJECT_LEVEL_WEIGHTS.get(num_pyramid_levels, OBJECT_LEVEL_WEIGHTS[5])):
            object_level_weights.extend([object_level_weight] * (feature_size[0] * feature_size[1] * 3))
        # self.object_level_weights = tf.convert_to_tensor(object_level_weights) * (1.4 if num_pyramid_levels >= 4 else 1.0)
        self.object_level_weights = tf.convert_to_tensor(object_level_weights)
        self.feature_split = [feature_size[0] * feature_size[1] * 3 for feature_size in feature_sizes]

        super().__init__(**kwargs)
        self.input_shape, self.pyramid_levels = input_shape, pyramid_levels
        self.gamma, self.class_loss_weight, self.bbox_loss_weight, self.object_loss_weight = gamma, class_loss_weight, bbox_loss_weight, object_loss_weight
        self.label_smoothing, self.from_logits = label_smoothing, from_logits

        self.class_acc = tf.Variable(0, dtype="float32", trainable=False)
//  [python]
def decode_offsets(ofs: tf.Tensor):
    _, h, w, _ = ofs.get_shape()
    x,y = tf.meshgrid(tf.range(w), tf.range(h))
    yx = tf.stack([y,x], axis=-1)
    yx = tf.to_float(yx)
    return yx - ofs
//  [python]
def decode_offsets(ofs: tf.Tensor):
    _, h, w, _ = ofs.get_shape()
    x,y = tf.meshgrid(tf.range(w), tf.range(h))
    yx = tf.stack([y,x], axis=-1)
    yx = tf.cast(yx, tf.float32)
    return yx - ofs
//  [python]
def __valid_call_single__(self, bbox_labels_true, bbox_labels_pred):
        bbox_labels_true_assined = tf.stop_gradient(self.anchor_assign(bbox_labels_true, bbox_labels_pred))
        if self.with_encoded_bboxes:
            bboxes_true, bboxes_true_encoded, labels_true, object_true_idx_nd = tf.split(bbox_labels_true_assined, [4, 4, -1, 1], axis=-1)
        else:
            bboxes_true, labels_true, object_true_idx_nd = tf.split(bbox_labels_true_assined, [4, -1, 1], axis=-1)
        object_true_idx_nd = tf.cast(object_true_idx_nd, tf.int32)
        object_true = tf.tensor_scatter_nd_update(tf.zeros_like(bbox_labels_pred[:, -1]), object_true_idx_nd, tf.ones_like(bboxes_true[:, -1]))

        bbox_labels_pred_valid = tf.gather_nd(bbox_labels_pred, object_true_idx_nd)
        if self.use_object_scores:
            bboxes_pred, labels_pred, _ = tf.split(bbox_labels_pred_valid, [self.regression_len, -1, 1], axis=-1)
            object_pred = bbox_labels_pred[:, -1]
        else:
            bboxes_pred, labels_pred = tf.split(bbox_labels_pred_valid, [self.regression_len, -1], axis=-1)
        # bboxes_true.set_shape(bboxes_pred.shape)

        anchors_centers = tf.gather_nd(self.anchor_assign.anchors_centers, object_true_idx_nd)
        anchors_hws = tf.gather_nd(self.anchor_assign.anchors_hws, object_true_idx_nd)
        bboxes_pred_top_left, bboxes_pred_bottom_right, bboxes_pred_center, bboxes_pred_hw = self.anchor_assign.__decode_bboxes__(
            bboxes_pred, anchors_centers, anchors_hws
        )

        if self.label_smoothing > 0:
            labels_true = labels_true * (1.0 - self.label_smoothing) + 0.5 * self.label_smoothing
        # class_loss = tf.reduce_sum(K.binary_crossentropy(labels_true, labels_pred))
        class_loss = K.binary_crossentropy(labels_true, labels_pred)
        bbox_loss = self.__iou_loss__(bboxes_true, bboxes_pred_top_left, bboxes_pred_bottom_right, bboxes_pred_hw)

        l1_loss = self.__l1_loss__(bboxes_true_encoded, bboxes_pred) if self.use_l1_loss else 0.0  # mean absolute error
        dfl_loss = self.__dfl_loss__(bboxes_true_encoded, bboxes_pred, labels_true) if self.use_dfl_loss else 0.0
        # class_loss, bbox_loss, l1_loss, dfl_loss = tf.reduce_mean(class_loss), tf.reduce_mean(bbox_loss), tf.reduce_mean(l1_loss), tf.reduce_mean(dfl_loss)
        class_loss, bbox_loss, l1_loss, dfl_loss = self.__reduce_with_class_weight__(labels_true, class_loss, bbox_loss, l1_loss, dfl_loss)

        object_loss = tf.reduce_sum(K.binary_crossentropy(tf.cast(object_true, object_pred.dtype), object_pred)) if self.use_object_scores else 0.0

        num_valid_anchors = tf.cast(tf.shape(bboxes_pred)[0], bboxes_pred.dtype)
        class_acc = tf.reduce_mean(tf.cast(tf.argmax(labels_true, axis=-1) == tf.argmax(labels_pred, axis=-1), "float32"))
        return class_loss, bbox_loss, object_loss, l1_loss, dfl_loss, num_valid_anchors, class_acc
//  [python]
def build_prompt_seq_from_task_id(
    task_vocab_id: int, response_seq=None, prompt_shape=None
):
  """Build prompt seq just using task id.

  Args:
    task_vocab_id: Vocab id for the task.
    response_seq: an (optional) discerte target sequen with shape (bsz, ..., k).
    prompt_shape: an (optional) tuple for prompt shape. One and only one of
      `response_seq` and `prompt_shape` should be specified.

  Returns:
    discrete input sequence of task id with shape (bsz, ..., 1).
  """
  task_id = tf.constant(task_vocab_id)
  if response_seq is not None:
    prompt_seq = tf.zeros_like(response_seq[..., :1]) + tf.to_int64(
        task_id
    )
  if prompt_shape is not None:
    assert response_seq is None, 'double specification'
    prompt_seq = tf.zeros(prompt_shape, dtype=tf.int64) + tf.to_int64(
        task_id
    )
  return prompt_seq
//  [python]
def build_prompt_seq_from_task_id(
    task_vocab_id: int, response_seq=None, prompt_shape=None
):
  """Build prompt seq just using task id.

  Args:
    task_vocab_id: Vocab id for the task.
    response_seq: an (optional) discerte target sequen with shape (bsz, ..., k).
    prompt_shape: an (optional) tuple for prompt shape. One and only one of
      `response_seq` and `prompt_shape` should be specified.

  Returns:
    discrete input sequence of task id with shape (bsz, ..., 1).
  """
  task_id = tf.constant(task_vocab_id)
  if response_seq is not None:
    prompt_seq = tf.zeros_like(response_seq[..., :1]) + tf.cast(
        task_id, response_seq.dtype
    )
  if prompt_shape is not None:
    assert response_seq is None, 'double specification'
    prompt_seq = tf.zeros(prompt_shape, dtype=tf.int64) + tf.cast(
        task_id, dtype=tf.int64
    )
  return prompt_seq
//  [python]
def _random_int_up_to(maxval, random_fn):
  # Need to cast because the int kernel for uniform doesn't support bcast.
  # We add one because maxval is exclusive, and this will get rounded down
  # when we cast back to int.
  float_maxval = t.cast(maxval, t.float32)
  return t.cast(
      random_fn(
          shape=t.shape(maxval),
          minval=t.zeros_like(float_maxval),
          maxval=float_maxval),
      dtype=maxval.dtype)
//  [python]
def _random_int_up_to(maxval, random_fn):
  # Need to cast because the int kernel for uniform doesn't support bcast.
  # We add one because maxval is exclusive, and this will get rounded down
  # when we cast back to int.
  float_maxval = tf.cast(maxval, tf.float32)
  return tf.cast(
      random_fn(
          shape=tf.shape(maxval),
          minval=tf.zeros_like(float_maxval),
          maxval=float_maxval),
      dtype=maxval.dtype)
//  [python]
def _random_int_from_range(minval, maxval, random_fn):
  # Need to cast because the int kernel for uniform doesn't support bcast.
  # We add one because maxval is exclusive, and this will get rounded down
  # when we cast back to int.
  float_minval = t.cast(minval, t.float32)
  float_maxval = t.cast(maxval, t.float32)
  return t.cast(
      random_fn(tf.shape(maxval), minval=float_minval, maxval=float_maxval),
      maxval.dtype)
//  [python]
def _random_int_from_range(minval, maxval, random_fn):
  # Need to cast because the int kernel for uniform doesn't support bcast.
  # We add one because maxval is exclusive, and this will get rounded down
  # when we cast back to int.
  float_minval = tf.cast(minval, tf.float32)
  float_maxval = tf.cast(maxval, tf.float32)
  return tf.cast(
      random_fn(tf.shape(maxval), minval=float_minval, maxval=float_maxval),
      maxval.dtype)
//  [python]
def jitter_bbox(bbox, min_range=0.0, max_range=0.05, truncation=True):
  """Jitter the bbox.

  Args:
    bbox: `float` tensor of shape (n, 4), ranged between 0 and 1.
    min_range: min jitter range in ratio to bbox size.
    max_range: max jitter range in ratio to bbox size.
    truncation: whether to truncate resulting bbox to remain [0, 1].
  Note: To create noisy positives, set min_range=0, which enables truncated
    normal distribution. max_range <=0.05: noisy duplicates, <=0.02: near
    duplicate. To create negatives: set min_range >= 0.1 to avoid false
    negatives; suggested max_range <=0.4 to avoid too much randomness.

  Returns:
    jittered bbox.
  """
  n = tf.shape(bbox)[0]
  h = bbox[:, 2] - bbox[:, 0]
  w = bbox[:, 3] - bbox[:, 1]
  noise = tf.stack([h, w, h, w], -1)
  if min_range == 0:
    noise_rate = tf.random.truncated_normal(
        [n, 4], mean=0, stddev=max_range / 2.0, dtype=bbox.dtype
    )
  else:
    noise_rate1 = tf.random.uniform([n, 4], min_range, max_range)
    noise_rate2 = tf.random.uniform([n, 4], -max_range, -min_range)
    selector = t.cast(tf.random.uniform([n, 4], 0, 1) < 0.5, tf.float32)
    noise_rate = noise_rate1 * selector + noise_rate2 * (1.0 - selector)
  bbox = bbox + noise * noise_rate
  return truncation_bbox(bbox) if truncation else bbox
//  [python]
def jitter_bbox(bbox, min_range=0.0, max_range=0.05, truncation=True):
  """Jitter the bbox.

  Args:
    bbox: `float` tensor of shape (n, 4), ranged between 0 and 1.
    min_range: min jitter range in ratio to bbox size.
    max_range: max jitter range in ratio to bbox size.
    truncation: whether to truncate resulting bbox to remain [0, 1].
  Note: To create noisy positives, set min_range=0, which enables truncated
    normal distribution. max_range <=0.05: noisy duplicates, <=0.02: near
    duplicate. To create negatives: set min_range >= 0.1 to avoid false
    negatives; suggested max_range <=0.4 to avoid too much randomness.

  Returns:
    jittered bbox.
  """
  n = tf.shape(bbox)[0]
  h = bbox[:, 2] - bbox[:, 0]
  w = bbox[:, 3] - bbox[:, 1]
  noise = tf.stack([h, w, h, w], -1)
  if min_range == 0:
    noise_rate = tf.random.truncated_normal(
        [n, 4], mean=0, stddev=max_range / 2.0, dtype=bbox.dtype
    )
  else:
    noise_rate1 = tf.random.uniform([n, 4], min_range, max_range)
    noise_rate2 = tf.random.uniform([n, 4], -max_range, -min_range)
    selector = tf.cast(tf.random.uniform([n, 4], 0, 1) < 0.5, tf.float32)
    noise_rate = noise_rate1 * selector + noise_rate2 * (1.0 - selector)
  bbox = bbox + noise * noise_rate
  return truncation_bbox(bbox) if truncation else bbox
//  [python]
def _sample_from_other_batch(sentences, random_fn):
  """Samples sentences from other batches."""
  # other_batch: <int64>[num_sentences]: The batch to sample from for each
  # sentence.
  other_batch = random_fn(
      shape=[tf.size(sentences)],
      minval=0,
      maxval=sentences.nrows() - 1,
      dtype=tf.int64)

  other_batch += t.cast(other_batch >= sentences.value_rowids(), tf.int64)

  # other_sentence: <int64>[num_sentences]: The sentence within each batch
  # that we sampled.
  other_sentence = _random_int_up_to(
      tf.gather(sentences.row_lengths(), other_batch), random_fn)
  return sentences.with_values(tf.stack([other_batch, other_sentence], axis=1))
//  [python]
def _sample_from_other_batch(sentences, random_fn):
  """Samples sentences from other batches."""
  # other_batch: <int64>[num_sentences]: The batch to sample from for each
  # sentence.
  other_batch = random_fn(
      shape=[tf.size(sentences)],
      minval=0,
      maxval=sentences.nrows() - 1,
      dtype=tf.int64)

  other_batch += tf.cast(other_batch >= sentences.value_rowids(), tf.int64)

  # other_sentence: <int64>[num_sentences]: The sentence within each batch
  # that we sampled.
  other_sentence = _random_int_up_to(
      tf.gather(sentences.row_lengths(), other_batch), random_fn)
  return sentences.with_values(tf.stack([other_batch, other_sentence], axis=1))
//  [python]
def sampling_points_to_voxel_index(sampling_points, voxel_size):
  """Transforms the sampling points from [-1, 1] to [0, voxel_size]."""
  voxel_size = tf.convert_to_tensor(value=voxel_size)
  max_size = tf.to_int32(voxel_size - 1)
  return 0.5 * ((sampling_points + 1) * max_size)
//  [python]
def sampling_points_to_voxel_index(sampling_points, voxel_size):
  """Transforms the sampling points from [-1, 1] to [0, voxel_size]."""
  voxel_size = tf.convert_to_tensor(value=voxel_size)
  max_size = tf.cast(voxel_size - 1, sampling_points.dtype)
  return 0.5 * ((sampling_points + 1) * max_size)
//  [python]
def ious_of_masks(mi_a, mi_b, h=544, w=704):
    n_mask_a = mi_a[-1,0] + 1
    n_mask_b = mi_b[-1,0] + 1
    mask_stack_a = tf.scatter_nd(mi_a, tf.ones((tf.shape(mi_a)[0],), tf.bool), [n_mask_a, h, w])
    mask_stack_b = tf.scatter_nd(mi_b, tf.ones((tf.shape(mi_b)[0],), tf.bool), [n_mask_b, h, w])
    #mask_areas_a = tf.math.count_nonzero(mask_stack_a, axis=(1,2))
    mask_areas_b = tf.math.count_nonzero(mask_stack_b, axis=(1,2))

    def iou_one_row(one_mask):
        intersects = tf.math.count_nonzero(one_mask & mask_stack_b, axis=(1,2))
        unions = tf.math.count_nonzero(one_mask) + mask_areas_b - intersects
        return tf.to_float(intersects) / (tf.to_float(unions) + 1.0e-7)
    ious = tf.map_fn(
        iou_one_row,
        mask_stack_a,
        fn_output_signature=tf.TensorSpec((None,), tf.float32),
    )
    return ious
//  [python]
def ious_of_masks(mi_a, mi_b, h=544, w=704):
    n_mask_a = mi_a[-1,0] + 1
    n_mask_b = mi_b[-1,0] + 1
    mask_stack_a = tf.scatter_nd(mi_a, tf.ones((tf.shape(mi_a)[0],), tf.bool), [n_mask_a, h, w])
    mask_stack_b = tf.scatter_nd(mi_b, tf.ones((tf.shape(mi_b)[0],), tf.bool), [n_mask_b, h, w])
    #mask_areas_a = tf.math.count_nonzero(mask_stack_a, axis=(1,2))
    mask_areas_b = tf.math.count_nonzero(mask_stack_b, axis=(1,2))

    def iou_one_row(one_mask):
        intersects = tf.math.count_nonzero(one_mask & mask_stack_b, axis=(1,2))
        unions = tf.math.count_nonzero(one_mask) + mask_areas_b - intersects
        return tf.cast(intersects, tf.float32) / (tf.cast(unions, tf.float32) + 1.0e-7)
    ious = tf.map_fn(
        iou_one_row,
        mask_stack_a,
        fn_output_signature=tf.TensorSpec((None,), tf.float32),
    )
    return ious
//  [python]
def sampling_points_from_3d_grid(grid_size, dtype=tf.float32):
  """Returns a tensor of shape `[M, 3]`, with M the number of sampling points."""
  sampling_points = grid.generate((-1.0, -1.0, -1.0), (1.0, 1.0, 1.0),
                                  grid_size)
  sampling_points = tf.cast(sampling_points, dtype)
  return tf.reshape(sampling_points, [-1, 3])
//  [python]
def transform_volume(voxels, transformation_matrix, voxel_size=(128, 128, 128)):
  """Apply a transformation to the input voxels."""
  voxels = tf.convert_to_tensor(voxels)
  volume_sampling = sampling_points_from_3d_grid(voxel_size)
  volume_sampling = tf.matmul(transformation_matrix,
                              tf.transpose(a=volume_sampling))
  volume_sampling = tf.to_float(tf.linalg.matrix_transpose(volume_sampling))
  volume_sampling = sampling_points_to_voxel_index(volume_sampling, voxel_size)
  interpolated_voxels = trilinear.interpolate(voxels, volume_sampling)
  return tf.reshape(interpolated_voxels, tf.shape(voxels))
//  [python]
def transform_volume(voxels, transformation_matrix, voxel_size=(128, 128, 128)):
  """Apply a transformation to the input voxels."""
  voxels = tf.convert_to_tensor(voxels)
  volume_sampling = sampling_points_from_3d_grid(voxel_size)
  volume_sampling = tf.matmul(transformation_matrix,
                              tf.transpose(a=volume_sampling))
  volume_sampling = tf.cast(tf.linalg.matrix_transpose(volume_sampling),
                            tf.float32)
  volume_sampling = sampling_points_to_voxel_index(volume_sampling, voxel_size)
  interpolated_voxels = trilinear.interpolate(voxels, volume_sampling)
  return tf.reshape(interpolated_voxels, tf.shape(voxels))
//  [python]
def _build_metrics_op(self):
        with tf.name_scope('metrics'):
            # Accuracy
            predicted = tf.argmax(self._pred, axis=-1)
            equal_first = tf.cast(tf.equal(predicted[:, :-1], self._placeholders.context[:, 1:]), tf.float32)
            self._accuracy = tf.reduce_mean(equal_first)
            self._elapsed_time = tf.compat.v1.placeholder(dtype=tf.float32, shape=[], name='elapsed_time')
//  [python]
def update_state(self, y_true, y_pred, sample_weight=None):
        y_true = tf.squeeze(y_true, axis=-1)
        y_pred += 1

        zeros_y_pred = tf.zeros(tf.shape(y_pred), tf.int64)
        zeros_y_pred += y_pred
        indices = tf.to_int64(tf.where(tf.equal(y_true, 0), 0, 1))

        y_true *= indices
        zeros_y_pred *= indices

        return super().update_state(y_true, zeros_y_pred, sample_weight)
//  [python]
def update_state(self, y_true, y_pred, sample_weight=None):
        y_true = tf.squeeze(y_true, axis=-1)
        y_pred += 1

        zeros_y_pred = tf.zeros(tf.shape(y_pred), tf.int64)
        zeros_y_pred += y_pred
        indices = tf.cast(tf.where(tf.equal(y_true, 0), 0, 1), tf.int64)

        y_true *= indices
        zeros_y_pred *= indices

        return super().update_state(y_true, zeros_y_pred, sample_weight)
//  [python]
def __call__(self, global_step: int):
    global_step = tf.to_float(global_step)
    warmup_steps = self._warmup_steps
    init_lr = self._init_learning_rate
    total_steps = self._total_steps

    linear_warmup = global_step / warmup_steps * init_lr

    cosine_learning_rate = init_lr * (tf.cos(np.pi *
                                             (global_step - warmup_steps) /
                                             (total_steps - warmup_steps)) +
                                      1.0) / 2.0

    learning_rate = tf.where(global_step < warmup_steps, linear_warmup,
                             cosine_learning_rate)
    return learning_rate
//  [python]
def __call__(self, global_step: int):
    global_step = tf.cast(global_step, dtype=tf.float32)
    warmup_steps = self._warmup_steps
    init_lr = self._init_learning_rate
    total_steps = self._total_steps

    linear_warmup = global_step / warmup_steps * init_lr

    cosine_learning_rate = init_lr * (tf.cos(np.pi *
                                             (global_step - warmup_steps) /
                                             (total_steps - warmup_steps)) +
                                      1.0) / 2.0

    learning_rate = tf.where(global_step < warmup_steps, linear_warmup,
                             cosine_learning_rate)
    return learning_rate
//  [python]
def read_variable(self,
                    variable: tf.Variable,
                    as_dtype: Optional[tf.DType] = None):
    if as_dtype is not None:
      variable = t.cast(variable, dtype=as_dtype)
    return variable
//  [python]
def read_variable(self,
                    variable: tf.Variable,
                    as_dtype: Optional[tf.DType] = None):
    if as_dtype is not None:
      variable = tf.cast(variable, dtype=as_dtype)
    return variable
//  [python]
def scatter_update_3d(tensor, idxs, updates, name=''):
    """
    tensor: [batch, rows, 1]
    idxs:   [batch, row_idx (N)]
    returns [batch, rows, 1] with updates applied according to row_idx above
    """

    batch = shape_list(idxs)[0]
    num_rows = shape_list(idxs)[1]

    # [batch]
    rows = t.cast(tf.range(batch), idxs.dtype)

    # [batch] -> [batch, 1, 1]
    for dim in range(2):
        rows = tf.expand_dims(rows, axis=1)

    # [batch, num_rows, 1]
    rows = tf.tile(rows, [1, num_rows, 1])

    # [batch, num_rows, 1] + [batch, num_rows, 1] to create
    # an index (row, col) into tensor on the first dimension
    row_cols = tf.concat([rows, idxs], axis=2)

    return tf.compat.v1.tensor_scatter_update(
        tensor, row_cols, updates, name=name)
//  [python]
def scatter_update_3d(tensor, idxs, updates, name=''):
    """
    tensor: [batch, rows, 1]
    idxs:   [batch, row_idx (N)]
    returns [batch, rows, 1] with updates applied according to row_idx above
    """

    batch = shape_list(idxs)[0]
    num_rows = shape_list(idxs)[1]

    # [batch]
    rows = tf.cast(tf.range(batch), idxs.dtype)

    # [batch] -> [batch, 1, 1]
    for dim in range(2):
        rows = tf.expand_dims(rows, axis=1)

    # [batch, num_rows, 1]
    rows = tf.tile(rows, [1, num_rows, 1])

    # [batch, num_rows, 1] + [batch, num_rows, 1] to create
    # an index (row, col) into tensor on the first dimension
    row_cols = tf.concat([rows, idxs], axis=2)

    return tf.compat.v1.tensor_scatter_update(
        tensor, row_cols, updates, name=name)
//  [python]
def __call__(self, step: int):
    lr = self._lr_schedule(step)
    if self._warmup_steps:
      if self._warmup_lr is not None:
        initial_learning_rate = t.convert_to_tensor(
            self._warmup_lr, name="initial_learning_rate")
      else:
        initial_learning_rate = t.convert_to_tensor(
            self._lr_schedule.initial_learning_rate,
            name="initial_learning_rate")
      dtype = initial_learning_rate.dtype
      global_step_recomp = t.cast(step, dtype)
      warmup_steps = t.cast(self._warmup_steps, dtype)
      warmup_lr = initial_learning_rate * global_step_recomp / warmup_steps
      lr = t.cond(global_step_recomp < warmup_steps, lambda: warmup_lr,
                   lambda: lr)
    return lr
//  [python]
def __call__(self, step: int):
    lr = self._lr_schedule(step)
    if self._warmup_steps:
      if self._warmup_lr is not None:
        initial_learning_rate = tf.convert_to_tensor(
            self._warmup_lr, name="initial_learning_rate")
      else:
        initial_learning_rate = tf.convert_to_tensor(
            self._lr_schedule.initial_learning_rate,
            name="initial_learning_rate")
      dtype = initial_learning_rate.dtype
      global_step_recomp = tf.cast(step, dtype)
      warmup_steps = tf.cast(self._warmup_steps, dtype)
      warmup_lr = initial_learning_rate * global_step_recomp / warmup_steps
      lr = tf.cond(global_step_recomp < warmup_steps, lambda: warmup_lr,
                   lambda: lr)
    return lr
//  [python]
def get_piecewise_lr(global_step, boundaries, lr_values, show_summary=True):
    ''' Piesewise learning rate'''
    # args example: (from https://github.com/tensorflow/models/blob/master/official/resnet/imagenet_main.py)
    # initial_learning_rate = 0.1 * params['batch_size'] / 128
    # batches_per_epoch = _NUM_IMAGES['train'] / params['batch_size']
    # boundaries = [int(batches_per_epoch * epoch) for epoch in [100, 150, 200]]
    # lr_values = [initial_learning_rate * decay for decay in [1, 0.1, 0.01, 0.001]]
    learning_rate = tf.train.piecewise_constant(
        tf.to_int32(global_step), boundaries, lr_values)
    learning_rate = tf.identity(learning_rate, name='learning_rate')
    if show_summary:
        tf.summary.scalar('learning_rate', learning_rate)
    return learning_rate
//  [python]
def get_piecewise_lr(global_step, boundaries, lr_values, show_summary=True):
    ''' Piesewise learning rate'''
    # args example: (from https://github.com/tensorflow/models/blob/master/official/resnet/imagenet_main.py)
    # initial_learning_rate = 0.1 * params['batch_size'] / 128
    # batches_per_epoch = _NUM_IMAGES['train'] / params['batch_size']
    # boundaries = [int(batches_per_epoch * epoch) for epoch in [100, 150, 200]]
    # lr_values = [initial_learning_rate * decay for decay in [1, 0.1, 0.01, 0.001]]
    learning_rate = tf.train.piecewise_constant(
        tf.cast(global_step, tf.int32), boundaries, lr_values)
    learning_rate = tf.identity(learning_rate, name='learning_rate')
    if show_summary:
        tf.summary.scalar('learning_rate', learning_rate)
    return learning_rate
//  [python]
def object_to_world(voxels,
                    euler_angles_x,
                    euler_angles_y,
                    translation_vector,
                    target_volume_size=(128, 128, 128)):
  """Apply the transformations to the voxels and place them in world coords."""
  scale_factor = 1.82  # object to world voxel space scale factor

  translation_vector = tf.expand_dims(translation_vector, axis=-1)

  sampling_points = t.cast(sampling_points_from_3d_grid(target_volume_size),
                            tf.float32)  # 128^3 X 3
  transf_matrix_x = rotation_matrix_3d.from_euler(euler_angles_x)  # [B, 3, 3]
  transf_matrix_y = rotation_matrix_3d.from_euler(euler_angles_y)  # [B, 3, 3]
  transf_matrix = tf.matmul(transf_matrix_x, transf_matrix_y)  # [B, 3, 3]
  transf_matrix = transf_matrix*scale_factor  # [B, 3, 3]
  sampling_points = tf.matmul(transf_matrix,
                              tf.transpose(sampling_points))  # [B, 3, N]
  translation_vector = tf.matmul(transf_matrix, translation_vector)  # [B, 3, 1]
  sampling_points = sampling_points - translation_vector
  sampling_points = tf.linalg.matrix_transpose(sampling_points)
  sampling_points = sampling_points_to_voxel_index(sampling_points,
                                                   target_volume_size)
  sampling_points = t.cast(sampling_points, tf.float32)
  interpolated_points = trilinear.interpolate(voxels, sampling_points)
  interpolated_voxels = tf.reshape(interpolated_points,
                                   [-1] + list(target_volume_size)+[4])

  return interpolated_voxels
//  [python]
def object_to_world(voxels,
                    euler_angles_x,
                    euler_angles_y,
                    translation_vector,
                    target_volume_size=(128, 128, 128)):
  """Apply the transformations to the voxels and place them in world coords."""
  scale_factor = 1.82  # object to world voxel space scale factor

  translation_vector = tf.expand_dims(translation_vector, axis=-1)

  sampling_points = tf.cast(sampling_points_from_3d_grid(target_volume_size),
                            tf.float32)  # 128^3 X 3
  transf_matrix_x = rotation_matrix_3d.from_euler(euler_angles_x)  # [B, 3, 3]
  transf_matrix_y = rotation_matrix_3d.from_euler(euler_angles_y)  # [B, 3, 3]
  transf_matrix = tf.matmul(transf_matrix_x, transf_matrix_y)  # [B, 3, 3]
  transf_matrix = transf_matrix*scale_factor  # [B, 3, 3]
  sampling_points = tf.matmul(transf_matrix,
                              tf.transpose(sampling_points))  # [B, 3, N]
  translation_vector = tf.matmul(transf_matrix, translation_vector)  # [B, 3, 1]
  sampling_points = sampling_points - translation_vector
  sampling_points = tf.linalg.matrix_transpose(sampling_points)
  sampling_points = sampling_points_to_voxel_index(sampling_points,
                                                   target_volume_size)
  sampling_points = tf.cast(sampling_points, tf.float32)
  interpolated_points = trilinear.interpolate(voxels, sampling_points)
  interpolated_voxels = tf.reshape(interpolated_points,
                                   [-1] + list(target_volume_size)+[4])

  return interpolated_voxels
//  [python]
def __call__(self, x):
    # Keeps the computation inside the layer norm to be float32.
    compute_dtype = x.dtype
    x = t.cast(x, dtype=t.float32)
    variance = tf.math.reduce_mean(tf.math.square(x), axis=-1, keepdims=True)
    x = x * tf.math.rsqrt(variance + self.variance_epsilon)
    x = t.cast(x, dtype=compute_dtype)
    return self.read_variable(self.weight, as_dtype=compute_dtype) * x
//  [python]
def __call__(self, x):
    # Keeps the computation inside the layer norm to be float32.
    compute_dtype = x.dtype
    x = tf.cast(x, dtype=tf.float32)
    variance = tf.math.reduce_mean(tf.math.square(x), axis=-1, keepdims=True)
    x = x * tf.math.rsqrt(variance + self.variance_epsilon)
    x = tf.cast(x, dtype=compute_dtype)
    return self.read_variable(self.weight, as_dtype=compute_dtype) * x
//  [python]
def make_attention_mask(query_input,
                        key_input,
                        pairwise_fn=tf.multiply,
                        dtype=tf.float32):
  """Mask-making helper for attention weights.

  In case of 1d inputs (i.e., `[batch..., len_q]`, `[batch..., len_kv]`, the
  attention weights will be `[batch..., heads, len_q, len_kv]` and this
  function will produce `[batch..., 1, len_q, len_kv]`.

  Args:
    query_input: a batched, flat input of query_length size
    key_input: a batched, flat input of key_length size
    pairwise_fn: broadcasting elementwise comparison function
    dtype: mask return dtype

  Returns:
    A `[batch..., 1, len_q, len_kv]` shaped mask for 1d attention.
  """
  mask = pairwise_fn(
      tf.expand_dims(query_input, axis=-1), tf.expand_dims(key_input, axis=-2))
  mask = tf.expand_dims(mask, axis=-3)
  return tf.to_float(mask)
//  [python]
def make_attention_mask(query_input,
                        key_input,
                        pairwise_fn=tf.multiply,
                        dtype=tf.float32):
  """Mask-making helper for attention weights.

  In case of 1d inputs (i.e., `[batch..., len_q]`, `[batch..., len_kv]`, the
  attention weights will be `[batch..., heads, len_q, len_kv]` and this
  function will produce `[batch..., 1, len_q, len_kv]`.

  Args:
    query_input: a batched, flat input of query_length size
    key_input: a batched, flat input of key_length size
    pairwise_fn: broadcasting elementwise comparison function
    dtype: mask return dtype

  Returns:
    A `[batch..., 1, len_q, len_kv]` shaped mask for 1d attention.
  """
  mask = pairwise_fn(
      tf.expand_dims(query_input, axis=-1), tf.expand_dims(key_input, axis=-2))
  mask = tf.expand_dims(mask, axis=-3)
  return tf.cast(mask, dtype=dtype)
//  [python]
def _relative_position_bucket(relative_position,
                                bidirectional=True,
                                num_buckets=32,
                                max_distance=128):
    """Translate relative position to a bucket number for relative attention.

    The relative position is defined as memory_position - query_position, i.e.
    the distance in tokens from the attending position to the attended-to
    position.

    If bidirectional=False, then positive relative positions are invalid.

    We use smaller buckets for small absolute relative_position and larger
    buckets for larger absolute relative_positions.

    All relative positions >=max_distance map to the same bucket.

    All relative positions <=-max_distance map to the same bucket.

    This should allow for more graceful generalization to longer sequences
    than the model has been trained on.

    Args:
      relative_position: an int32 Tensor
      bidirectional: a boolean - whether the attention is bidirectional
      num_buckets: an integer
      max_distance: an integer

    Returns:
      a Tensor with the same shape as relative_position, containing int32
      values in the range [0, num_buckets)
    """
    ret = 0
    n = -relative_position
    if bidirectional:
      num_buckets //= 2
      ret += t.cast(t.less(n, 0), t.int32) * num_buckets
      n = t.abs(n)
    else:
      n = t.maximum(n, 0)
    # now n is in the range [0, inf)
    max_exact = num_buckets // 2
    is_small = t.less(n, max_exact)
    val_if_large = max_exact + t.cast(
        t.log(
            t.cast(n, t.float32) / max_exact + np.finfo(np.float32).eps) /
        math.log(max_distance / max_exact) * (num_buckets - max_exact),
        t.int32,
    )
    val_if_large = t.minimum(val_if_large, num_buckets - 1)
    ret += t.where(is_small, n, val_if_large)
    return ret
//  [python]
def _relative_position_bucket(relative_position,
                                bidirectional=True,
                                num_buckets=32,
                                max_distance=128):
    """Translate relative position to a bucket number for relative attention.

    The relative position is defined as memory_position - query_position, i.e.
    the distance in tokens from the attending position to the attended-to
    position.

    If bidirectional=False, then positive relative positions are invalid.

    We use smaller buckets for small absolute relative_position and larger
    buckets for larger absolute relative_positions.

    All relative positions >=max_distance map to the same bucket.

    All relative positions <=-max_distance map to the same bucket.

    This should allow for more graceful generalization to longer sequences
    than the model has been trained on.

    Args:
      relative_position: an int32 Tensor
      bidirectional: a boolean - whether the attention is bidirectional
      num_buckets: an integer
      max_distance: an integer

    Returns:
      a Tensor with the same shape as relative_position, containing int32
      values in the range [0, num_buckets)
    """
    ret = 0
    n = -relative_position
    if bidirectional:
      num_buckets //= 2
      ret += tf.cast(tf.math.less(n, 0), tf.int32) * num_buckets
      n = tf.math.abs(n)
    else:
      n = tf.math.maximum(n, 0)
    # now n is in the range [0, inf)
    max_exact = num_buckets // 2
    is_small = tf.math.less(n, max_exact)
    val_if_large = max_exact + tf.dtypes.cast(
        tf.math.log(
            tf.cast(n, tf.float32) / max_exact + np.finfo(np.float32).eps) /
        math.log(max_distance / max_exact) * (num_buckets - max_exact),
        tf.int32,
    )
    val_if_large = tf.math.minimum(val_if_large, num_buckets - 1)
    ret += tf.where(is_small, n, val_if_large)
    return ret
//  [python]
def __call__(self, sigmoids, labels):
    levels = sorted(sigmoids.keys())

    labels_bce = tf.squeeze(labels, axis=-1)
    labels = tf.cast(labels, tf.float32)

    bce_losses = []
    ssim_losses = []
    iou_losses = []

    for level in levels:
      bce_losses.append(
          self._binary_crossentropy(labels_bce, sigmoids[level]))
      ssim_losses.append(
          1 - self._ssim(sigmoids[level], labels, max_val=1.0))
      iou_losses.append(
          self._iou_loss(sigmoids[level], labels))

    total_bce_loss = tf.math.add_n(bce_losses)
    total_ssim_loss = tf.math.add_n(ssim_losses)
    total_iou_loss = tf.math.add_n(iou_losses)

    total_loss = total_bce_loss + total_ssim_loss + total_iou_loss
    total_loss = total_loss / len(levels)

    return total_loss
//  [python]
def __call__(self,
               inputs=None,
               encoder_mask=None,
               dense_inputs=None,
               training=False):
    """Applies Transformer model on the inputs.

    Args:
      inputs: input word ids. Optional if dense data are provided.
      encoder_mask: the encoder self-attention mask.
      dense_inputs: dense input data. Concat after the embedding if word ids are
        provided.
      training: whether it is training pass, affecting dropouts.

    Returns:
      output of a transformer encoder.
    """
    # Casts inputs to the dtype.
    if encoder_mask is not None:
      encoder_mask = t.cast(encoder_mask, self.compute_dtype)
    cfg = self.config
    inputs_array = []
    if inputs is not None:
      inputs_array.append(
          self.input_embed(inputs, one_hot=cfg.one_hot_embedding))
    if dense_inputs is not None:
      inputs_array.append(dense_inputs)
    if not inputs_array:
      raise ValueError("At least one of inputs and dense_inputs must not be "
                       "None.")
    x = tf.concat(inputs_array, axis=1)
    tensor_shape = tf_utils.get_shape_list(x)
    tensor_shape[-2] = 1
    x = self.input_dropout(x, noise_shape=tensor_shape, training=training)
    if inputs is not None:
      input_length = tf_utils.get_shape_list(inputs)[1]
    else:
      input_length = 0

    attention_outputs = []
    for i in range(cfg.num_layers):
      position_bias = self.get_relpos_bias(input_length, dense_inputs, i)
      x = self.encoder_layers[i](
          x,
          attention_mask=encoder_mask,
          position_bias=position_bias,
          training=training)
      if self.config.return_attention_scores:
        x, attention_scores = x
        attention_outputs.append(attention_scores)

    encoded = self.output_norm(x)
    encoded = self.output_dropout(encoded, training=training)
    if self.config.return_attention_scores:
      return encoded, attention_outputs
    else:
      return encoded