jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def occlusion(flow_fw, flow_bw):
    x_shape = tf.shape(flow_fw)
    H = x_shape[1]
    W = x_shape[2]    
    flow_bw_warped = tf_warp(flow_bw, flow_fw, H, W)
    flow_fw_warped = tf_warp(flow_fw, flow_bw, H, W)
    flow_diff_fw = flow_fw + flow_bw_warped
    flow_diff_bw = flow_bw + flow_fw_warped
    mag_sq_fw = length_sq(flow_fw) + length_sq(flow_bw_warped)
    mag_sq_bw = length_sq(flow_bw) + length_sq(flow_fw_warped)
    occ_thresh_fw =  0.01 * mag_sq_fw + 0.5
    occ_thresh_bw =  0.01 * mag_sq_bw + 0.5
    occ_fw = tf.to_float(tf.cast(length_sq(flow_diff_fw) > occ_thresh_fw, tf.int32))
    occ_bw = tf.to_float(tf.cast(length_sq(flow_diff_bw) > occ_thresh_bw, tf.int32))

    return occ_fw, occ_bw
//  [python]
def occlusion(flow_fw, flow_bw):
    x_shape = tf.shape(flow_fw)
    H = x_shape[1]
    W = x_shape[2]    
    flow_bw_warped = tf_warp(flow_bw, flow_fw, H, W)
    flow_fw_warped = tf_warp(flow_fw, flow_bw, H, W)
    flow_diff_fw = flow_fw + flow_bw_warped
    flow_diff_bw = flow_bw + flow_fw_warped
    mag_sq_fw = length_sq(flow_fw) + length_sq(flow_bw_warped)
    mag_sq_bw = length_sq(flow_bw) + length_sq(flow_fw_warped)
    occ_thresh_fw =  0.01 * mag_sq_fw + 0.5
    occ_thresh_bw =  0.01 * mag_sq_bw + 0.5
    occ_fw = tf.cast(length_sq(flow_diff_fw) > occ_thresh_fw, tf.float32)
    occ_bw = tf.cast(length_sq(flow_diff_bw) > occ_thresh_bw, tf.float32)

    return occ_fw, occ_bw
//  [python]
def _reduce_instance_metrics(
      self, logs: Dict[str, Any], use_masks: bool = False
  ):
    """Updates the per class and mean instance metrics in the logs."""
    if use_masks:
      instance_metrics = self.instance_mask_perclass_metrics
      prefix = 'mask_'
    else:
      instance_metrics = self.instance_box_perclass_metrics
      prefix = ''
    if instance_metrics is None:
      raise ValueError(
          'No instance metrics defined when use_masks is %s' % use_masks
      )
    result = instance_metrics.result()
    iou_thresholds = instance_metrics.get_config()['iou_thresholds']

    for ap_key in instance_metrics.get_average_precision_metrics_keys():
      # (num_iou_thresholds, num_classes)
      per_class_ap = tf.where(
          result['valid_classes'], result[ap_key], tf.zeros_like(result[ap_key])
      )
      # (num_iou_thresholds,)
      mean_ap_by_iou = tf.math.divide_no_nan(
          tf.reduce_sum(per_class_ap, axis=-1),
          tf.reduce_sum(
              tf.to_float(result['valid_classes']),
              axis=-1,
          ),
      )
      logs[f'{prefix}{ap_key}'] = tf.reduce_mean(mean_ap_by_iou)
      for j, iou in enumerate(iou_thresholds):
        if int(iou * 100) in {50, 75}:
          logs[f'{prefix}{ap_key}{int(iou * 100)}'] = mean_ap_by_iou[j]

      if self.task_config.per_category_metrics:
        # (num_classes,)
        per_class_mean_ap = tf.reduce_mean(per_class_ap, axis=0)
        valid_classes = result['valid_classes'].numpy()
        for k in range(self.task_config.model.num_classes):
          if valid_classes[k]:
            logs[f'{prefix}{ap_key} ByCategory/{k}'] = per_class_mean_ap[k]
            for j, iou in enumerate(iou_thresholds):
              if int(iou * 100) in {50, 75}:
                logs[f'{prefix}{ap_key}{int(iou * 100)} ByCategory/{k}'] = (
                    per_class_ap[j][k]
                )
//  [python]
def _reduce_instance_metrics(
      self, logs: Dict[str, Any], use_masks: bool = False
  ):
    """Updates the per class and mean instance metrics in the logs."""
    if use_masks:
      instance_metrics = self.instance_mask_perclass_metrics
      prefix = 'mask_'
    else:
      instance_metrics = self.instance_box_perclass_metrics
      prefix = ''
    if instance_metrics is None:
      raise ValueError(
          'No instance metrics defined when use_masks is %s' % use_masks
      )
    result = instance_metrics.result()
    iou_thresholds = instance_metrics.get_config()['iou_thresholds']

    for ap_key in instance_metrics.get_average_precision_metrics_keys():
      # (num_iou_thresholds, num_classes)
      per_class_ap = tf.where(
          result['valid_classes'], result[ap_key], tf.zeros_like(result[ap_key])
      )
      # (num_iou_thresholds,)
      mean_ap_by_iou = tf.math.divide_no_nan(
          tf.reduce_sum(per_class_ap, axis=-1),
          tf.reduce_sum(
              tf.cast(result['valid_classes'], dtype=per_class_ap.dtype),
              axis=-1,
          ),
      )
      logs[f'{prefix}{ap_key}'] = tf.reduce_mean(mean_ap_by_iou)
      for j, iou in enumerate(iou_thresholds):
        if int(iou * 100) in {50, 75}:
          logs[f'{prefix}{ap_key}{int(iou * 100)}'] = mean_ap_by_iou[j]

      if self.task_config.per_category_metrics:
        # (num_classes,)
        per_class_mean_ap = tf.reduce_mean(per_class_ap, axis=0)
        valid_classes = result['valid_classes'].numpy()
        for k in range(self.task_config.model.num_classes):
          if valid_classes[k]:
            logs[f'{prefix}{ap_key} ByCategory/{k}'] = per_class_mean_ap[k]
            for j, iou in enumerate(iou_thresholds):
              if int(iou * 100) in {50, 75}:
                logs[f'{prefix}{ap_key}{int(iou * 100)} ByCategory/{k}'] = (
                    per_class_ap[j][k]
                )
//  [python]
def call(self, y_true, box_outputs):
    anchor_boxes = tf.tile(
        self.input_anchors.boxes,
        [box_outputs.shape[0] // self.input_anchors.boxes.shape[0], 1])
    num_positives, box_targets = y_true
    mask = tf.to_float(box_targets != 0.0)
    box_outputs = anchors.decode_box_outputs(box_outputs, anchor_boxes) * mask
    box_targets = anchors.decode_box_outputs(box_targets, anchor_boxes) * mask
    normalizer = num_positives * 4.0
    box_iou_loss = iou_utils.iou_loss(box_outputs, box_targets,
                                      self.iou_loss_type)
    box_iou_loss = tf.reduce_sum(box_iou_loss) / normalizer
    return box_iou_loss
//  [python]
def call(self, y_true, box_outputs):
    anchor_boxes = tf.tile(
        self.input_anchors.boxes,
        [box_outputs.shape[0] // self.input_anchors.boxes.shape[0], 1])
    num_positives, box_targets = y_true
    mask = tf.cast(box_targets != 0.0, box_targets.dtype)
    box_outputs = anchors.decode_box_outputs(box_outputs, anchor_boxes) * mask
    box_targets = anchors.decode_box_outputs(box_targets, anchor_boxes) * mask
    normalizer = num_positives * 4.0
    box_iou_loss = iou_utils.iou_loss(box_outputs, box_targets,
                                      self.iou_loss_type)
    box_iou_loss = tf.reduce_sum(box_iou_loss) / normalizer
    return box_iou_loss
//  [python]
def shift_tokens_right(input_ids: tf.Tensor, pad_token_id: int):
    """
    Shift input ids one token to the right, and wrap the last non pad token (the <LID> token) Note that MBart does not
    have a single `decoder_start_token_id` in contrast to other Bart-like models.
    """
    assert pad_token_id is not None, "self.model.config.pad_token_id has to be defined."
    # replace possible -100 values in labels by `pad_token_id`
    input_ids = tf.where(input_ids == -100, tf.fill(shape_list(input_ids), pad_token_id), input_ids)
    language_id_index = (
        tf.reduce_sum(tf.to_int32(tf.math.not_equal(input_ids, pad_token_id)), axis=-1) - 1
    )
    language_id_index = tf.stack([tf.range(shape_list(input_ids)[0]), language_id_index], axis=-1)
    languages_ids = tf.gather_nd(input_ids, language_id_index)

    shifted_input_ids = tf.concat([tf.expand_dims(languages_ids, axis=-1), input_ids[:, :-1]], axis=-1)

    return shifted_input_ids
//  [python]
def shift_tokens_right(input_ids: tf.Tensor, pad_token_id: int):
    """
    Shift input ids one token to the right, and wrap the last non pad token (the <LID> token) Note that MBart does not
    have a single `decoder_start_token_id` in contrast to other Bart-like models.
    """
    assert pad_token_id is not None, "self.model.config.pad_token_id has to be defined."
    # replace possible -100 values in labels by `pad_token_id`
    input_ids = tf.where(input_ids == -100, tf.fill(shape_list(input_ids), pad_token_id), input_ids)
    language_id_index = (
        tf.reduce_sum(tf.cast(tf.math.not_equal(input_ids, pad_token_id), dtype=input_ids.dtype), axis=-1) - 1
    )
    language_id_index = tf.stack([tf.range(shape_list(input_ids)[0]), language_id_index], axis=-1)
    languages_ids = tf.gather_nd(input_ids, language_id_index)

    shifted_input_ids = tf.concat([tf.expand_dims(languages_ids, axis=-1), input_ids[:, :-1]], axis=-1)

    return shifted_input_ids
//  [python]
def decode(self, serialized_example: Any) -> Mapping[str, Any]:
    """Decode the serialized example.

    Args:
      serialized_example: a single serialized tf.Example string.

    Returns:
      decoded_tensors: a dictionary of tensors with the following fields:
        - frame_id: an int64 scalar tensor to identify an example.
        - pillars: a float32 tensor of shape [P, N, D].
        - indices: an int32 tensor of shape [P, 2].
        - gt_classes: an int32 tensor of shape [M].
        - gt_boxes: a float32 tensor of shape [M, 4].
        - gt_attributes: a dict of (name, [M, 1]) float32 pairs.
        - gt_difficulty: an int32 tensor of shape [M].
    """
    parsed_tensors = tf.io.parse_single_example(
        serialized=serialized_example, features=self._feature_description)

    # Convert sparse tensor to dense tensor.
    for k in parsed_tensors:
      if isinstance(parsed_tensors[k], tf.SparseTensor):
        parsed_tensors[k] = tf.sparse.to_dense(
            parsed_tensors[k], default_value=0)

    # Decode features and labels.
    frame_id = parsed_tensors['frame_id']
    pillars, indices = self._decode_pillars(parsed_tensors)
    classes = tf.to_int32(parsed_tensors['bbox/class'])
    boxes = self._decode_boxes(parsed_tensors)
    attr_heading = tf.expand_dims(parsed_tensors['bbox/heading'], axis=1)
    attr_z = tf.expand_dims(parsed_tensors['bbox/z'], axis=1)
    attr_height = tf.expand_dims(parsed_tensors['bbox/height'], axis=1)
    difficulty = tf.to_int32(parsed_tensors['bbox/difficulty'])

    decoded_tensors = {
        'frame_id': frame_id,
        'pillars': pillars,
        'indices': indices,
        'gt_classes': classes,
        'gt_boxes': boxes,
        'gt_attributes': {
            'heading': attr_heading,
            'z': attr_z,
            'height': attr_height,
        },
        'gt_difficulty': difficulty,
    }
    return decoded_tensors
//  [python]
def decode(self, serialized_example: Any) -> Mapping[str, Any]:
    """Decode the serialized example.

    Args:
      serialized_example: a single serialized tf.Example string.

    Returns:
      decoded_tensors: a dictionary of tensors with the following fields:
        - frame_id: an int64 scalar tensor to identify an example.
        - pillars: a float32 tensor of shape [P, N, D].
        - indices: an int32 tensor of shape [P, 2].
        - gt_classes: an int32 tensor of shape [M].
        - gt_boxes: a float32 tensor of shape [M, 4].
        - gt_attributes: a dict of (name, [M, 1]) float32 pairs.
        - gt_difficulty: an int32 tensor of shape [M].
    """
    parsed_tensors = tf.io.parse_single_example(
        serialized=serialized_example, features=self._feature_description)

    # Convert sparse tensor to dense tensor.
    for k in parsed_tensors:
      if isinstance(parsed_tensors[k], tf.SparseTensor):
        parsed_tensors[k] = tf.sparse.to_dense(
            parsed_tensors[k], default_value=0)

    # Decode features and labels.
    frame_id = parsed_tensors['frame_id']
    pillars, indices = self._decode_pillars(parsed_tensors)
    classes = tf.cast(parsed_tensors['bbox/class'], tf.int32)
    boxes = self._decode_boxes(parsed_tensors)
    attr_heading = tf.expand_dims(parsed_tensors['bbox/heading'], axis=1)
    attr_z = tf.expand_dims(parsed_tensors['bbox/z'], axis=1)
    attr_height = tf.expand_dims(parsed_tensors['bbox/height'], axis=1)
    difficulty = tf.cast(parsed_tensors['bbox/difficulty'], tf.int32)

    decoded_tensors = {
        'frame_id': frame_id,
        'pillars': pillars,
        'indices': indices,
        'gt_classes': classes,
        'gt_boxes': boxes,
        'gt_attributes': {
            'heading': attr_heading,
            'z': attr_z,
            'height': attr_height,
        },
        'gt_difficulty': difficulty,
    }
    return decoded_tensors
//  [python]
def _quantile(a, q, axis=None):
    ret_dtype = a.dtype
    if tf.experimental.numpy.ndim(q) > 1:
        raise ValueError("q argument must be a scalar or 1-dimensional!")
    if axis is None:
        axis = 0
        a = tf.reshape(a, [-1])
    elif axis != 0:
        a = tf.experimental.numpy.moveaxis(a, axis, 0)
        axis = 0

    n = a.shape[axis]

    indices = q * (n - 1)

    a = tf.sort(a, axis)

    indices_below = tf.to_int32(tf.math.floor(indices))
    indices_upper = tf.to_int32(tf.math.ceil(indices))

    weights = indices - tf.to_float(indices_below)

    indices_below = tf.clip_by_value(indices_below, 0, n - 1)
    indices_upper = tf.clip_by_value(indices_upper, 0, n - 1)
    tensor_upper = tf.gather(a, indices_upper, axis=axis)
    tensor_below = tf.gather(a, indices_below, axis=axis)

    pred = weights <= 0.5
    out = tf.where(pred, tensor_below, tensor_upper)

    return tf.to_float(out)
//  [python]
def _quantile(a, q, axis=None):
    ret_dtype = a.dtype
    if tf.experimental.numpy.ndim(q) > 1:
        raise ValueError("q argument must be a scalar or 1-dimensional!")
    if axis is None:
        axis = 0
        a = tf.reshape(a, [-1])
    elif axis != 0:
        a = tf.experimental.numpy.moveaxis(a, axis, 0)
        axis = 0

    n = a.shape[axis]

    indices = q * (n - 1)

    a = tf.sort(a, axis)

    indices_below = tf.cast(tf.math.floor(indices), dtype=tf.int32)
    indices_upper = tf.cast(tf.math.ceil(indices), dtype=tf.int32)

    weights = indices - tf.cast(indices_below, dtype=ret_dtype)

    indices_below = tf.clip_by_value(indices_below, 0, n - 1)
    indices_upper = tf.clip_by_value(indices_upper, 0, n - 1)
    tensor_upper = tf.gather(a, indices_upper, axis=axis)
    tensor_below = tf.gather(a, indices_below, axis=axis)

    pred = weights <= 0.5
    out = tf.where(pred, tensor_below, tensor_upper)

    return tf.cast(out, ret_dtype)
//  [python]
def test_recompute_grad(self):
    """Tests that the gradient is computed correctly with recompute_grad."""
    dense = tf.keras.layers.Dense(10, input_shape=(8,))
    dropout = recomputing_dropout.RecomputingDropout(
        0.4, force_recomputation=True)

    @recompute_grad.recompute_grad
    def recompute_dense_dropout(x):
      return dropout(dense(x), training=True)

    # Define the model using dropout.
    def f(x):
      with tf.GradientTape() as tape:
        h1 = recompute_dense_dropout(x)
        h2 = recompute_dense_dropout(x)
        y = tf.math.reduce_sum(h1 + h2)
      return (tf.cast(tf.math.not_equal(h1, 0), tf.float32),
              tf.cast(tf.math.not_equal(h2, 0),
                      tf.float32), tape.gradient(y, dense.trainable_variables))

    x = tf.convert_to_tensor(np.random.normal(size=(4, 8)), tf.float32)
    mask1, mask2, gradients = f(x)
    self.evaluate(tf.compat.v1.initializers.global_variables())

    mask1, mask2, gradients = self.evaluate([mask1, mask2, gradients])
    # Make sure entries were masked and there is randomness.
    self.assertGreaterEqual(np.sum(mask1 == 0), 2)
    self.assertGreaterEqual(np.sum(mask2 == 0), 2)
    self.assertNotAllEqual(mask1, mask2)

    # Use the masks to compute exact gradients.
    def g(x):
      with tf.GradientTape() as tape:
        # Rescale proportional to dropout rate.
        h1 = (dense(x) * mask1) / 0.6
        h2 = (dense(x) * mask2) / 0.6
        y = tf.math.reduce_sum(h1 + h2)
      return tape.gradient(y, dense.trainable_variables)

    expected_gradients = self.evaluate(g(x))
    self.assertAllClose(gradients, expected_gradients)
//  [python]
def cummin(
    x: Union[tf.Tensor, tf.Variable],
    /,
    *,
    axis: int = 0,
    exclusive: bool = False,
    reverse: bool = False,
    dtype: Optional[tf.DType] = None,
    out: Optional[Union[tf.Tensor, tf.Variable]] = None,
) -> Union[tf.Tensor, tf.Variable]:
    dtype = ivy.as_native_dtype(dtype)
    if reverse:
        x = tf.reverse(x, axis=[axis])
    x_unstacked = tf.unstack(x, axis=axis)
    cummin_x_unstacked = []
    cummin_x_unstacked.append(x_unstacked[0])
    for i, x_sub in enumerate(x_unstacked[1:]):
        cummin_x_sub = tf.minimum(cummin_x_unstacked[i], x_sub)
        cummin_x_unstacked.append(cummin_x_sub)
    cummin_x = tf.stack(cummin_x_unstacked, axis=axis)
    if reverse:
        cummin_x = tf.reverse(cummin_x, axis=[axis])
    if dtype is None:
        return cummin_x
    else:
        return tf.cast(cummin_x, dtype)
//  [python]
def test_nested_recompute_grad(self):
    """Tests nested usage of recompute_grad."""
    dense = tf.keras.layers.Dense(
        5, input_shape=(8,), bias_initializer='glorot_normal')
    dropout = recomputing_dropout.RecomputingDropout(
        0.4, force_recomputation=True)

    @recompute_grad.recompute_grad
    def recompute_dense_dropout_tower(x):
      return dropout(dense(x), training=True)

    def make_head():
      inputs = tf.keras.Input(shape=(5,))
      x = tf.keras.layers.Dense(
          3, activation='tanh', name='dense', bias_initializer='glorot_normal')(
              inputs)
      x = recomputing_dropout.RecomputingDropout(0.45)(x)
      outputs = {
          'head_mask': tf.cast(tf.math.not_equal(x, 0), tf.float32),
          'y': tf.reduce_sum(x),
      }
      return tf.keras.Model(inputs, outputs, name='head')

    head = make_head()

    # Nest recompute_grad inside another recompute_grad function.
    @recompute_grad.recompute_grad
    def recompute_model(x):
      y1 = recompute_dense_dropout_tower(x)
      y2 = recompute_dense_dropout_tower(x)
      outputs = head(y1 + y2, training=True)
      outputs.update({
          'tower1_mask': tf.cast(tf.math.not_equal(y1, 0), tf.float32),
          'tower2_mask': tf.cast(tf.math.not_equal(y2, 0), tf.float32),
      })
      return outputs

    def f(x):
      with tf.GradientTape() as tape:
        outputs = recompute_model(x)
      outputs['gradients'] = tape.gradient(
          outputs.pop('y'),
          dense.trainable_variables + head.trainable_variables)
      return outputs

    x = tf.convert_to_tensor(np.random.normal(size=(4, 8)), tf.float32)
    outputs = f(x)
    self.evaluate(tf.compat.v1.initializers.global_variables())
    outputs = self.evaluate(outputs)

    # Verify gradients are correct.
    def g(x):
      with tf.GradientTape() as tape:
        y1 = dense(x) * outputs['tower1_mask'] / 0.6
        y2 = dense(x) * outputs['tower2_mask'] / 0.6
        y = tf.reduce_sum(
            head.get_layer('dense')(y1 + y2) * outputs['head_mask'] / 0.55)
      return tape.gradient(y,
                           dense.trainable_variables + head.trainable_variables)
    # Increase tolerance from default of 1e-6 to reduce flakiness.
    self.assertAllClose(
        outputs['gradients'], self.evaluate(g(x)), rtol=2e-5, atol=2e-5)
//  [python]
def normal_max_for_show(disp):
    disp_max = tf.reduce_max(disp)
    disp_new = disp/disp_max
    disp_new = disp_new*255
    disp_new_uint8 = tf.cast(disp_new, dtype=tf.uint8)
    return disp_new_uint8
//  [python]
def render_lm2d(lm2d_batch_xy, h, w):
    """
    :param lm2d_batch:
    :param h:
    :param w:
    :return:
    """
    # preprocess
    """
    row correspond to y
    column correspond to x
    (row, column) = (y, x)
    """
    x = lm2d_batch_xy[:, :, 0]
    y = lm2d_batch_xy[:, :, 1]
    x = tf.clip_by_value(x, 0, w-1)
    y = tf.clip_by_value(y, 0, h-1)
    lm2d_batch = tf.stack([y, x], axis=-1)

    #
    visual_lm2d = []
    for b_it in range(lm2d_batch.shape[0]):
        lm2d = lm2d_batch[b_it]
        lm2d = tf.to_int64(lm2d)
        r = tf.sparse_tensor_to_dense(
            tf.SparseTensor(indices=lm2d, values=tf.linspace(0.0, 1.0, lm2d.shape[0]),
                            dense_shape=[h, w]),
            validate_indices=False
        )
        g = tf.sparse_tensor_to_dense(
            tf.SparseTensor(indices=lm2d, values=tf.zeros(shape=[lm2d.shape[0]], dtype=tf.float32),
                            dense_shape=[h, w]),
            validate_indices=False
        )
        b = tf.sparse_tensor_to_dense(
            tf.SparseTensor(indices=lm2d, values=tf.zeros(shape=[lm2d.shape[0]], dtype=tf.float32),
                            dense_shape=[h, w]),
            validate_indices=False
        )
        rgb = tf.stack([r, g, b], axis=-1)
        visual_lm2d.append(rgb)

    visual_lm2d = tf.stack(visual_lm2d)
    return visual_lm2d
//  [python]
def render_lm2d(lm2d_batch_xy, h, w):
    """
    :param lm2d_batch:
    :param h:
    :param w:
    :return:
    """
    # preprocess
    """
    row correspond to y
    column correspond to x
    (row, column) = (y, x)
    """
    x = lm2d_batch_xy[:, :, 0]
    y = lm2d_batch_xy[:, :, 1]
    x = tf.clip_by_value(x, 0, w-1)
    y = tf.clip_by_value(y, 0, h-1)
    lm2d_batch = tf.stack([y, x], axis=-1)

    #
    visual_lm2d = []
    for b_it in range(lm2d_batch.shape[0]):
        lm2d = lm2d_batch[b_it]
        lm2d = tf.cast(lm2d, dtype=tf.int64)
        r = tf.sparse_tensor_to_dense(
            tf.SparseTensor(indices=lm2d, values=tf.linspace(0.0, 1.0, lm2d.shape[0]),
                            dense_shape=[h, w]),
            validate_indices=False
        )
        g = tf.sparse_tensor_to_dense(
            tf.SparseTensor(indices=lm2d, values=tf.zeros(shape=[lm2d.shape[0]], dtype=tf.float32),
                            dense_shape=[h, w]),
            validate_indices=False
        )
        b = tf.sparse_tensor_to_dense(
            tf.SparseTensor(indices=lm2d, values=tf.zeros(shape=[lm2d.shape[0]], dtype=tf.float32),
                            dense_shape=[h, w]),
            validate_indices=False
        )
        rgb = tf.stack([r, g, b], axis=-1)
        visual_lm2d.append(rgb)

    visual_lm2d = tf.stack(visual_lm2d)
    return visual_lm2d
//  [python]
def resample(img, flow):
    # img, NCHW
    # flow, N2HW
    B = tf.shape(img)[0]
    c = tf.shape(img)[1]
    h = tf.shape(img)[2]
    w = tf.shape(img)[3]
    img_flat = tf.reshape(tf.transpose(img, [0, 2, 3, 1]), [-1, c])

    dx, dy = tf.unstack(flow, axis=1)
    xf, yf = tf.meshgrid(tf.range(w), tf.range(h))
    xf = xf + dx
    yf = yf + dy

    alpha = tf.expand_dims(xf - tf.floor(xf), axis=-1)
    beta = tf.expand_dims(yf - tf.floor(yf), axis=-1)

    xL = tf.clip_by_value(tf.floor(xf), 0, w - 1)
    xR = tf.clip_by_value(tf.floor(xf) + 1, 0, w - 1)
    yT = tf.clip_by_value(tf.floor(yf), 0, h - 1)
    yB = tf.clip_by_value(tf.floor(yf) + 1, 0, h - 1)

    batch_ids = tf.tile(tf.expand_dims(tf.expand_dims(tf.range(B), axis=-1), axis=-1), [1, h, w])

    def get(y, x):
        idx = tf.reshape(batch_ids * h * w + y * w + x, [-1])
        idx = tf.cast(idx, tf.int32)
        return tf.gather(img_flat, idx)

    val = tf.zeros_like(alpha)
    val += (1 - alpha) * (1 - beta) * tf.reshape(get(yT, xL), [-1, h, w, c])
    val += (0 + alpha) * (1 - beta) * tf.reshape(get(yT, xR), [-1, h, w, c])
    val += (1 - alpha) * (0 + beta) * tf.reshape(get(yB, xL), [-1, h, w, c])
    val += (0 + alpha) * (0 + beta) * tf.reshape(get(yB, xR), [-1, h, w, c])

    # we need to enforce the channel_dim known during compile-time here
    shp = img.shape.as_list()
    return tf.reshape(tf.transpose(val, [0, 3, 1, 2]), [-1, shp[1], h, w])
//  [python]
def resample(img, flow):
    # img, NCHW
    # flow, N2HW
    B = tf.shape(img)[0]
    c = tf.shape(img)[1]
    h = tf.shape(img)[2]
    w = tf.shape(img)[3]
    img_flat = tf.reshape(tf.transpose(img, [0, 2, 3, 1]), [-1, c])

    dx, dy = tf.unstack(flow, axis=1)
    xf, yf = tf.meshgrid(tf.cast(tf.range(w), tf.float32), tf.cast(tf.range(h), tf.float32))
    xf = xf + dx
    yf = yf + dy

    alpha = tf.expand_dims(xf - tf.floor(xf), axis=-1)
    beta = tf.expand_dims(yf - tf.floor(yf), axis=-1)

    xL = tf.clip_by_value(tf.cast(tf.floor(xf), dtype=tf.int32), 0, w - 1)
    xR = tf.clip_by_value(tf.cast(tf.floor(xf) + 1, dtype=tf.int32), 0, w - 1)
    yT = tf.clip_by_value(tf.cast(tf.floor(yf), dtype=tf.int32), 0, h - 1)
    yB = tf.clip_by_value(tf.cast(tf.floor(yf) + 1, dtype=tf.int32), 0, h - 1)

    batch_ids = tf.tile(tf.expand_dims(tf.expand_dims(tf.range(B), axis=-1), axis=-1), [1, h, w])

    def get(y, x):
        idx = tf.reshape(batch_ids * h * w + y * w + x, [-1])
        idx = tf.cast(idx, tf.int32)
        return tf.gather(img_flat, idx)

    val = tf.zeros_like(alpha)
    val += (1 - alpha) * (1 - beta) * tf.reshape(get(yT, xL), [-1, h, w, c])
    val += (0 + alpha) * (1 - beta) * tf.reshape(get(yT, xR), [-1, h, w, c])
    val += (1 - alpha) * (0 + beta) * tf.reshape(get(yB, xL), [-1, h, w, c])
    val += (0 + alpha) * (0 + beta) * tf.reshape(get(yB, xR), [-1, h, w, c])

    # we need to enforce the channel_dim known during compile-time here
    shp = img.shape.as_list()
    return tf.reshape(tf.transpose(val, [0, 3, 1, 2]), [-1, shp[1], h, w])
//  [python]
def normal_depthmap_for_show_bgMax(disp):
    #disp_min = tf.contrib.distributions.percentile(disp, q=0, axis=[1, 2], interpolation='lower')
    disp_min = tf.reduce_min(disp)
    #disp_max = disp_min+255
    #disp = tf.clip_by_value(disp, disp_min, disp_max)
    #disp_max = tf.expand_dims(tf.expand_dims(disp_max, 1), 1)
    disp_min = tf.reshape(disp_min, [1, 1, 1, 1])
    disp_new = disp-disp_min
    disp_new_uint8 = tf.to_int32(disp_new)
    return disp_new_uint8
//  [python]
def normal_depthmap_for_show_bgMax(disp):
    #disp_min = tf.contrib.distributions.percentile(disp, q=0, axis=[1, 2], interpolation='lower')
    disp_min = tf.reduce_min(disp)
    #disp_max = disp_min+255
    #disp = tf.clip_by_value(disp, disp_min, disp_max)
    #disp_max = tf.expand_dims(tf.expand_dims(disp_max, 1), 1)
    disp_min = tf.reshape(disp_min, [1, 1, 1, 1])
    disp_new = disp-disp_min
    disp_new_uint8 = tf.cast(disp_new, dtype=tf.uint8)
    return disp_new_uint8
//  [python]
def _router_z_loss(router_logits: tf.Tensor) -> float:
  """Computes router z-loss.

   The router z-loss was introduced in Designing Effective Sparse Expert Models
   (https://arxiv.org/abs/2202.08906). It encourages router logits to remain
   small in an effort to improve stability.

  Args:
    router_logits: <float32>[num_groups, tokens_per_group, num_experts] router
      logits.

  Returns:
    Scalar router z-loss <float32>.
  """
  num_groups = tf.shape(router_logits)[0]
  tokens_per_group = router_logits.shape[1]

  log_z = tf.math.reduce_logsumexp(router_logits, axis=-1)
  z_loss = log_z**2
  return tf.math.reduce_sum(z_loss) / tf.cast(
      num_groups * tokens_per_group, tf.float32)
//  [python]
def normal_depthmap_for_show(disp):
    disp_max = tf.contrib.distributions.percentile(disp, q=100, axis=[1, 2], interpolation='lower')
    disp_max = tf.reduce_max(disp_max)
    disp_min = disp_max-255*2
    disp = tf.clip_by_value(disp, disp_min, disp_max)
    #disp_max = tf.expand_dims(tf.expand_dims(disp_max, 1), 1)
    disp_min = tf.reshape(disp_min, [1, 1, 1, 1])

    disp_new = disp-disp_min
    #disp_new = (disp_new - disp_min) / (disp_max - disp_min)

    # disp_new = []
    # for i in range(disp.shape[0]):
    #     #disp_i = tf_render.clip_by_value(disp[i], disp_min[i], disp_max[i])
    #     dn = (disp[i] - disp_min[i]) / (disp_max[i] - disp_min[i])
    #     disp_new.append(dn)
    # disp_new = tf_render.stack(disp_new)
    disp_new = tf.to_int64(disp_new)
    return disp_new
//  [python]
def normal_depthmap_for_show(disp):
    disp_max = tf.contrib.distributions.percentile(disp, q=100, axis=[1, 2], interpolation='lower')
    disp_max = tf.reduce_max(disp_max)
    disp_min = disp_max-255*2
    disp = tf.clip_by_value(disp, disp_min, disp_max)
    #disp_max = tf.expand_dims(tf.expand_dims(disp_max, 1), 1)
    disp_min = tf.reshape(disp_min, [1, 1, 1, 1])

    disp_new = disp-disp_min
    #disp_new = (disp_new - disp_min) / (disp_max - disp_min)

    # disp_new = []
    # for i in range(disp.shape[0]):
    #     #disp_i = tf_render.clip_by_value(disp[i], disp_min[i], disp_max[i])
    #     dn = (disp[i] - disp_min[i]) / (disp_max[i] - disp_min[i])
    #     disp_new.append(dn)
    # disp_new = tf_render.stack(disp_new)
    disp_new = tf.cast(disp_new, dtype=tf.uint8)
    return disp_new
//  [python]
def read_and_decode(filename):
    # 根据文件名生成一个队列
    with tf.name_scope("input"):
        filename_queue = tf.train.string_input_producer([filename])
        reader = tf.TFRecordReader()
        _, serialized_example = reader.read(filename_queue)   # 返回文件名和文件
        features = tf.parse_single_example(serialized_example,
                                           features={
                                               'label': tf.FixedLenFeature([], tf.int64),
                                               'image_byte': tf.FixedLenFeature([], tf.string),
                                           })
        image = tf.decode_raw(features['image_byte'], tf.uint8)
        image = tf.reshape(image, [400, 400, 3])
        image = tf.cast(image, tf.float32)
        image = tf.image.per_image_standardization(image)  # 将图像标准化，有利于加速训练
        label = tf.cast(features['label'], tf.int32)

    return image, label
//  [python]
def sanitize_coordinates(x1, x2, img_size, padding=0):
    x1 = tf.minimum(x1, x2)
    x2 = tf.maximum(x1, x2)
    x1 = tf.clip_by_value(x1 - padding, clip_value_min=0., clip_value_max=1000000.)
    x2 = tf.clip_by_value(x2 + padding, clip_value_min=0., clip_value_max=tf.cast(img_size, tf.float32))
    return x1, x2
//  [python]
def load_img(path_to_img):
    """Define a function to load an image and limit its maximum dimension to 512 pixels."""
    max_dim = 512
    img = tf.io.read_file(path_to_img)
    img = tf.image.decode_image(img, channels=3)
    img = tf.image.convert_image_dtype(img, tf.float32)

    shape = tf.to_int32(tf.shape(img)[:-1])
    long_dim = max(shape)
    scale = max_dim / long_dim

    new_shape = tf.to_int32(shape * scale)

    img = tf.image.resize(img, new_shape)
    img = img[tf.newaxis, :]
    return img
//  [python]
def load_img(path_to_img):
    """Define a function to load an image and limit its maximum dimension to 512 pixels."""
    max_dim = 512
    img = tf.io.read_file(path_to_img)
    img = tf.image.decode_image(img, channels=3)
    img = tf.image.convert_image_dtype(img, tf.float32)

    shape = tf.cast(tf.shape(img)[:-1], tf.float32)
    long_dim = max(shape)
    scale = max_dim / long_dim

    new_shape = tf.cast(shape * scale, tf.int32)

    img = tf.image.resize(img, new_shape)
    img = img[tf.newaxis, :]
    return img
//  [python]
def output_list(x,seq_length):
    training = tf.constant(False,dtype = tf.bool,name = 'Training')
    config_path = os.path.join(FLAGS.model,'model.json')
    model_configure = read_config(config_path)
    logits, ratio = inference(x, 
                              seq_length, 
                              training=training,
                              full_sequence_len = FLAGS.segment_len,
                              configure = model_configure)
    ratio = tf.constant(ratio,dtype = tf.float32,shape = [])
    seq_length_r = tf.to_int32(tf.round(tf.to_float(seq_length)/ratio))
    prob_logits = path_prob(logits)
    predict,log_prob = tf.nn.ctc_beam_search_decoder(tf.transpose(logits, perm=[1, 0, 2]), 
                                            seq_length_r, 
                                            merge_repeated=False,
                                            beam_width = FLAGS.beam_width)
    return predict[0],logits,prob_logits,log_prob
//  [python]
def output_list(x,seq_length):
    training = tf.constant(False,dtype = tf.bool,name = 'Training')
    config_path = os.path.join(FLAGS.model,'model.json')
    model_configure = read_config(config_path)
    logits, ratio = inference(x, 
                              seq_length, 
                              training=training,
                              full_sequence_len = FLAGS.segment_len,
                              configure = model_configure)
    ratio = tf.constant(ratio,dtype = tf.float32,shape = [])
    seq_length_r = tf.cast(tf.round(tf.cast(seq_length,dtype = tf.float32)/ratio),tf.int32)
    prob_logits = path_prob(logits)
    predict,log_prob = tf.nn.ctc_beam_search_decoder(tf.transpose(logits, perm=[1, 0, 2]), 
                                            seq_length_r, 
                                            merge_repeated=False,
                                            beam_width = FLAGS.beam_width)
    return predict[0],logits,prob_logits,log_prob
//  [python]
def chebyshev5(self, x, L, Fout, K, training=False):
        fstart = time.time()
        N, M, Fin = x.get_shape()
        N, M, Fin = int(N), int(M), int(Fin)
        # Rescale Laplacian and store as a TF sparse tensor. Copy to not modify the shared L.
        indices = np.column_stack((L.row, L.col))
        L = tf.SparseTensor(indices, L.data, L.shape)
        L = tf.sparse_reorder(L)
        L = tf.to_float(L)
        # Transform to Chebyshev basis
        x0 = tf.transpose(x, perm=[1, 2, 0])  # M x Fin x N
        x0 = tf.reshape(x0, [M, Fin*N])  # M x Fin*N
        x = tf.expand_dims(x0, 0)  # 1 x M x Fin*N
        def concat(x, x_):
            x_ = tf.expand_dims(x_, 0)  # 1 x M x Fin*N
            return tf.concat([x, x_], axis=0)  # K x M x Fin*N
        if K > 1:
            x1 = tf.sparse_tensor_dense_matmul(L, x0)
            x = concat(x, x1)
        for k in range(2, K):
            x2 = 2 * tf.sparse_tensor_dense_matmul(L, x1) - x0  # M x Fin*N
            x = concat(x, x2)
            x0, x1 = x1, x2
        x = tf.reshape(x, [K, M, Fin, N])  # K x M x Fin x N
        x = tf.transpose(x, perm=[3, 1, 2, 0])  # N x M x Fin x K
        x = tf.reshape(x, [N*M, Fin*K])  # N*M x Fin*K
        # Filter: Fin*Fout filters of order K, i.e. one filterbank per output feature.
        W = self._weight_variable_cheby(K, Fin, Fout, regularization=True)
        # Drop filters of the convolutional layer
        if self.dropFilt<1:
            W = tf.reshape(W, [Fin, K, Fout])
            W = tf.transpose(W, perm=[1,0,2])
            mask = tf.ones((Fin, Fout))
            dropout = tf.cond(training, lambda: float(self.dropFilt), lambda: 1.0)
            mask = tf.nn.dropout(mask, dropout) * (dropout)
            W = tf.multiply(W, mask)
            W = tf.transpose(W, perm=[1,0,2])
            W = tf.reshape(W, [Fin*K, Fout])
        x = tf.matmul(x, W)  # N*M x Fout
        return tf.reshape(x, [N, M, Fout])
//  [python]
def chebyshev5(self, x, L, Fout, K, training=False):
        fstart = time.time()
        N, M, Fin = x.get_shape()
        N, M, Fin = int(N), int(M), int(Fin)
        # Rescale Laplacian and store as a TF sparse tensor. Copy to not modify the shared L.
        indices = np.column_stack((L.row, L.col))
        L = tf.SparseTensor(indices, L.data, L.shape)
        L = tf.sparse_reorder(L)
        L = tf.cast(L, self.dtype)
        # Transform to Chebyshev basis
        x0 = tf.transpose(x, perm=[1, 2, 0])  # M x Fin x N
        x0 = tf.reshape(x0, [M, Fin*N])  # M x Fin*N
        x = tf.expand_dims(x0, 0)  # 1 x M x Fin*N
        def concat(x, x_):
            x_ = tf.expand_dims(x_, 0)  # 1 x M x Fin*N
            return tf.concat([x, x_], axis=0)  # K x M x Fin*N
        if K > 1:
            x1 = tf.sparse_tensor_dense_matmul(L, x0)
            x = concat(x, x1)
        for k in range(2, K):
            x2 = 2 * tf.sparse_tensor_dense_matmul(L, x1) - x0  # M x Fin*N
            x = concat(x, x2)
            x0, x1 = x1, x2
        x = tf.reshape(x, [K, M, Fin, N])  # K x M x Fin x N
        x = tf.transpose(x, perm=[3, 1, 2, 0])  # N x M x Fin x K
        x = tf.reshape(x, [N*M, Fin*K])  # N*M x Fin*K
        # Filter: Fin*Fout filters of order K, i.e. one filterbank per output feature.
        W = self._weight_variable_cheby(K, Fin, Fout, regularization=True)
        # Drop filters of the convolutional layer
        if self.dropFilt<1:
            W = tf.reshape(W, [Fin, K, Fout])
            W = tf.transpose(W, perm=[1,0,2])
            mask = tf.ones((Fin, Fout))
            dropout = tf.cond(training, lambda: float(self.dropFilt), lambda: 1.0)
            mask = tf.nn.dropout(mask, dropout) * (dropout)
            W = tf.multiply(W, mask)
            W = tf.transpose(W, perm=[1,0,2])
            W = tf.reshape(W, [Fin*K, Fout])
        x = tf.matmul(x, W)  # N*M x Fout
        return tf.reshape(x, [N, M, Fout])
//  [python]
def crop(pred, boxes):
    # pred [num_obj, 138, 138], gt [num_bboxes, 4]
    # sanitize coordination (make sure the bboxes are in range 0 <= x, y <= image size)
    shape_pred = tf.shape(pred)
    pred_w = shape_pred[0]
    pred_h = shape_pred[1]

    xmin, xmax = sanitize_coordinates(boxes[:, 0], boxes[:, 2], pred_w, padding=1)
    ymin, ymax = sanitize_coordinates(boxes[:, 1], boxes[:, 3], pred_h, padding=1)

    rows = tf.broadcast_to(tf.range(pred_w)[None, :, None], shape_pred)
    cols = tf.broadcast_to(tf.range(pred_h)[:, None, None], shape_pred)

    xmin = xmin[None, None, :]
    ymin = ymin[None, None, :]
    xmax = xmax[None, None, :]
    ymax = ymax[None, None, :]

    mask_left = (rows >= t.cast(xmin, cols.dtype))
    mask_right = (rows <= t.cast(xmax, cols.dtype))
    mask_bottom = (cols >= t.cast(ymin, rows.dtype))
    mask_top = (cols <= t.cast(ymax, rows.dtype))

    crop_mask = tf.math.logical_and(tf.math.logical_and(mask_left, mask_right),
                                    tf.math.logical_and(mask_bottom, mask_top))
    return pred * t.cast(crop_mask, tf.float32)
//  [python]
def crop(pred, boxes):
    # pred [num_obj, 138, 138], gt [num_bboxes, 4]
    # sanitize coordination (make sure the bboxes are in range 0 <= x, y <= image size)
    shape_pred = tf.shape(pred)
    pred_w = shape_pred[0]
    pred_h = shape_pred[1]

    xmin, xmax = sanitize_coordinates(boxes[:, 0], boxes[:, 2], pred_w, padding=1)
    ymin, ymax = sanitize_coordinates(boxes[:, 1], boxes[:, 3], pred_h, padding=1)

    rows = tf.broadcast_to(tf.range(pred_w)[None, :, None], shape_pred)
    cols = tf.broadcast_to(tf.range(pred_h)[:, None, None], shape_pred)

    xmin = xmin[None, None, :]
    ymin = ymin[None, None, :]
    xmax = xmax[None, None, :]
    ymax = ymax[None, None, :]

    mask_left = (rows >= tf.cast(xmin, cols.dtype))
    mask_right = (rows <= tf.cast(xmax, cols.dtype))
    mask_bottom = (cols >= tf.cast(ymin, rows.dtype))
    mask_top = (cols <= tf.cast(ymax, rows.dtype))

    crop_mask = tf.math.logical_and(tf.math.logical_and(mask_left, mask_right),
                                    tf.math.logical_and(mask_bottom, mask_top))
    return pred * tf.cast(crop_mask, tf.float32)
//  [python]
def check_model_generate(
        self, config, input_ids, attention_mask, decoder_input_ids, decoder_attention_mask, **kwargs
    ):
        self.assertIsNotNone(config.question_encoder)
        self.assertIsNotNone(config.generator)

        for model_class in self.all_generative_model_classes:
            model = model_class(config, retriever=self.get_retriever(config))

            self.assertTrue(model.config.is_encoder_decoder)

            input_ids = t.cast(input_ids, t.int32)
            outputs = model.generate(
                input_ids=input_ids,
                num_beams=2,
                num_return_sequences=2,
                decoder_start_token_id=config.generator.eos_token_id,
                max_new_tokens=5,
            )

            self.assertIsNotNone(outputs)
//  [python]
def check_model_generate(
        self, config, input_ids, attention_mask, decoder_input_ids, decoder_attention_mask, **kwargs
    ):
        self.assertIsNotNone(config.question_encoder)
        self.assertIsNotNone(config.generator)

        for model_class in self.all_generative_model_classes:
            model = model_class(config, retriever=self.get_retriever(config))

            self.assertTrue(model.config.is_encoder_decoder)

            input_ids = tf.cast(input_ids, tf.int32)
            outputs = model.generate(
                input_ids=input_ids,
                num_beams=2,
                num_return_sequences=2,
                decoder_start_token_id=config.generator.eos_token_id,
                max_new_tokens=5,
            )

            self.assertIsNotNone(outputs)
//  [python]
def _get_batch_shape_tensor(
    field: Field, batch_rank: int, dtype: tf.dtypes.DType
) -> tf.Tensor:
  """Returns the batch shape tensor of the field."""
  if batch_rank == 0:
    return tf.constant([], dtype=dtype)

  assert field.shape.rank >= batch_rank

  if isinstance(field, tf.Tensor):
    return tf.shape(field, dtype)[:batch_rank]

  assert isinstance(field, tf.RaggedTensor)

  dims = [tf.size(field.row_splits, dtype) - tf.constant(1, dtype=dtype)]

  for d in range(1, batch_rank):
    if isinstance(field, tf.Tensor):
      dims.extend(tf.unstack(tf.shape(field, dtype)[1: (batch_rank + 1 - d)]))
      break

    if field.uniform_row_length is None:
      # NOTE: this exception should never be raised if the graph piece is
      # constructed using  `GraphPieceBase._from_data()` with static shapes
      # validatition being enabled (true by default). The `_from_data()`
      # checks that all input fields have fully defined batch dimensions, except
      # the outermost. This means that the ragged fields must have uniform
      # batch dimensions.
      raise ValueError(
          f'The dimension={d} of the ragged field must be uniform as it is part'
          f' of the graph shape (graph rank={batch_rank}), got ragged.'
      )

    dims.append(tf.to_int32(field.uniform_row_length))
    field = field.values

  return tf.stack(dims)
//  [python]
def _get_batch_shape_tensor(
    field: Field, batch_rank: int, dtype: tf.dtypes.DType
) -> tf.Tensor:
  """Returns the batch shape tensor of the field."""
  if batch_rank == 0:
    return tf.constant([], dtype=dtype)

  assert field.shape.rank >= batch_rank

  if isinstance(field, tf.Tensor):
    return tf.shape(field, dtype)[:batch_rank]

  assert isinstance(field, tf.RaggedTensor)

  dims = [tf.size(field.row_splits, dtype) - tf.constant(1, dtype=dtype)]

  for d in range(1, batch_rank):
    if isinstance(field, tf.Tensor):
      dims.extend(tf.unstack(tf.shape(field, dtype)[1: (batch_rank + 1 - d)]))
      break

    if field.uniform_row_length is None:
      # NOTE: this exception should never be raised if the graph piece is
      # constructed using  `GraphPieceBase._from_data()` with static shapes
      # validatition being enabled (true by default). The `_from_data()`
      # checks that all input fields have fully defined batch dimensions, except
      # the outermost. This means that the ragged fields must have uniform
      # batch dimensions.
      raise ValueError(
          f'The dimension={d} of the ragged field must be uniform as it is part'
          f' of the graph shape (graph rank={batch_rank}), got ragged.'
      )

    dims.append(tf.cast(field.uniform_row_length, dtype))
    field = field.values

  return tf.stack(dims)
//  [python]
def check_model_generate_from_context_input_ids(
        self, config, input_ids, attention_mask, decoder_input_ids, decoder_attention_mask, **kwargs
    ):
        self.assertIsNotNone(config.question_encoder)
        self.assertIsNotNone(config.generator)

        retriever = self.get_retriever(config)

        for i, model_class in enumerate(self.all_generative_model_classes):
            model = model_class(config)
            self.assertTrue(model.config.is_encoder_decoder)

            question_hidden_states = model.question_encoder(input_ids, attention_mask=attention_mask)[0]

            out = retriever(
                input_ids,
                question_hidden_states.numpy(),
                prefix=config.generator.prefix,
                return_tensors="tf",
            )

            context_input_ids, context_attention_mask, retrieved_doc_embeds = (
                out["context_input_ids"],
                out["context_attention_mask"],
                out["retrieved_doc_embeds"],
            )
            retrieved_doc_embeds = t.cast(retrieved_doc_embeds, t.float32)

            # compute doc_scores
            doc_scores = tf.squeeze(
                tf.matmul(tf.expand_dims(question_hidden_states, axis=[1]), retrieved_doc_embeds, transpose_b=True),
                axis=[1],
            )

            outputs = model.generate(
                context_input_ids=context_input_ids,
                context_attention_mask=context_attention_mask,
                doc_scores=doc_scores,
            )

            self.assertIsNotNone(outputs)
//  [python]
def check_model_generate_from_context_input_ids(
        self, config, input_ids, attention_mask, decoder_input_ids, decoder_attention_mask, **kwargs
    ):
        self.assertIsNotNone(config.question_encoder)
        self.assertIsNotNone(config.generator)

        retriever = self.get_retriever(config)

        for i, model_class in enumerate(self.all_generative_model_classes):
            model = model_class(config)
            self.assertTrue(model.config.is_encoder_decoder)

            question_hidden_states = model.question_encoder(input_ids, attention_mask=attention_mask)[0]

            out = retriever(
                input_ids,
                question_hidden_states.numpy(),
                prefix=config.generator.prefix,
                return_tensors="tf",
            )

            context_input_ids, context_attention_mask, retrieved_doc_embeds = (
                out["context_input_ids"],
                out["context_attention_mask"],
                out["retrieved_doc_embeds"],
            )
            retrieved_doc_embeds = tf.cast(retrieved_doc_embeds, tf.float32)

            # compute doc_scores
            doc_scores = tf.squeeze(
                tf.matmul(tf.expand_dims(question_hidden_states, axis=[1]), retrieved_doc_embeds, transpose_b=True),
                axis=[1],
            )

            outputs = model.generate(
                context_input_ids=context_input_ids,
                context_attention_mask=context_attention_mask,
                doc_scores=doc_scores,
            )

            self.assertIsNotNone(outputs)
//  [python]
def check_model_without_retriever(
        self, config, input_ids, attention_mask, decoder_input_ids, decoder_attention_mask, **kwargs
    ):
        self.assertIsNotNone(config.question_encoder)
        self.assertIsNotNone(config.generator)

        retriever = self.get_retriever(config)

        for model_class in self.all_model_classes:
            model = model_class(config)
            self.assertTrue(model.config.is_encoder_decoder)

            question_hidden_states = model.question_encoder(input_ids, attention_mask=attention_mask)[0]

            out = retriever(
                input_ids,
                question_hidden_states.numpy(),
                prefix=config.generator.prefix,
                return_tensors="tf",
            )

            context_input_ids, context_attention_mask, retrieved_doc_embeds = (
                out["context_input_ids"],
                out["context_attention_mask"],
                out["retrieved_doc_embeds"],
            )

            retrieved_doc_embeds = tf.cast(retrieved_doc_embeds, tf.float32)

            # compute doc_scores
            doc_scores = tf.squeeze(
                tf.matmul(tf.expand_dims(question_hidden_states, axis=[1]), retrieved_doc_embeds, transpose_b=True),
                axis=[1],
            )

            outputs = model(
                input_ids=None,
                context_input_ids=context_input_ids,
                context_attention_mask=context_attention_mask,
                doc_scores=doc_scores,
                decoder_input_ids=decoder_input_ids,
                decoder_attention_mask=decoder_attention_mask,
            )

            # logits
            self.assertEqual(
                outputs.logits.shape,
                (self.n_docs * decoder_input_ids.shape[0], decoder_input_ids.shape[1], config.generator.vocab_size),
            )

            # generator encoder last hidden states
            self.assertEqual(
                outputs.generator_enc_last_hidden_state.shape,
                (self.n_docs * decoder_input_ids.shape[0], self.max_combined_length, config.generator.hidden_size),
            )
            # doc scores
            self.assertEqual(outputs.doc_scores.shape, (input_ids.shape[0], self.n_docs))