jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def serve(self, images: tf.Tensor):
    """Cast image to float and run inference.

    Args:
      images: uint8 Tensor of shape [batch_size, None, None, 3]

    Returns:
      Tensor holding detection output logits.
    """
    if self._input_type != 'tflite':
      with tf.device('cpu:0'):
        images = tf.cast(images, dtype=tf.float32)
        images_spec = tf.TensorSpec(
            shape=self._input_image_size + [3], dtype=tf.float32)
        image_info_spec = tf.TensorSpec(shape=[4, 2], dtype=tf.float32)

        images, image_info = tf.nest.map_structure(
            tf.identity,
            tf.map_fn(
                self._build_inputs,
                elems=images,
                fn_output_signature=(images_spec, image_info_spec),
                parallel_iterations=32))

    outputs = self.model.call(
        inputs=images, image_info=image_info, training=False)

    masks = outputs['segmentation_outputs']
    masks = tf.image.resize(masks, self._input_image_size, method='bilinear')
    classes = tf.math.argmax(masks, axis=-1)
    scores = tf.nn.softmax(masks, axis=-1)
    final_outputs = {
        'semantic_logits': masks,
        'semantic_scores': scores,
        'semantic_classes': classes,
        'image_info': image_info,
        'panoptic_category_mask': outputs['category_mask'],
        'panoptic_instance_mask': outputs['instance_mask'],
    }

    return final_outputs
//  [python]
def create_model_gan(input_shape, generator, discriminator):
    input = Input(input_shape)

    gen_out = generator(input)
    # dis_out = discriminator(concatenate([gen_out, input], axis=3))
    gen_argmax = tf.argmax(gen_out, axis=-1)
    gen_argmax = tf.expand_dims(gen_argmax, axis=-1)
    gen_argmax = tf.cast(gen_argmax, tf.float32)
    half = 19. / 2.
    gen_argmax = (gen_argmax / half) - 1.
    dis_out = discriminator(gen_argmax)

    model = tf.keras.Model(inputs=[input], outputs=[dis_out, gen_out], name='dcgan')
    return model
//  [python]
def create_model_gan(input_shape, generator, discriminator):
    input = Input(input_shape)

    gen_out = generator(input)
    # dis_out = discriminator(concatenate([gen_out, input], axis=3))
    gen_argmax = tf.math.argmax(gen_out, axis=-1)
    gen_argmax = tf.expand_dims(gen_argmax, axis=-1)
    gen_argmax = tf.cast(gen_argmax, tf.float32)
    half = 19. / 2.
    gen_argmax = (gen_argmax / half) - 1.
    dis_out = discriminator(gen_argmax)

    model = tf.keras.Model(inputs=[input], outputs=[dis_out, gen_out], name='dcgan')
    return model
//  [python]
def accuracy(logits, labels):
  indices = tf.argmax(logits, 1, output_type=tf.int32)
  correct = tf.reduce_sum(tf.cast(tf.math.equal(indices, labels), tf.float32))
  return correct / tf.cast(tf.shape(labels)[0], tf.float32)
//  [python]
def accuracy(logits, labels):
  indices = tf.math.argmax(logits, 1, output_type=tf.int32)
  correct = tf.reduce_sum(tf.cast(tf.math.equal(indices, labels), tf.float32))
  return correct / tf.cast(tf.shape(labels)[0], tf.float32)
//  [python]
def validation_step(self, inputs, model: tf_keras.Model, metrics=None):
    features, labels = inputs, inputs
    outputs = self.inference_step(features, model)
    loss = self.build_losses(
        labels=labels, model_outputs=outputs, aux_losses=model.losses
    )
    logs = {self.loss: loss}
    if metrics:
      self.process_metrics(metrics, labels, outputs)
    if model.compiled_metrics:
      self.process_compiled_metrics(model.compiled_metrics, labels, outputs)
      logs.update({m.name: m.result() for m in metrics or []})
      logs.update({m.name: m.result() for m in model.metrics})
    if self.metric_type == 'matthews_corrcoef':
      logs.update({
          'sentence_prediction': (
              tf.expand_dims(  # Ensure one prediction along batch dimension.
                  tf.argmax(outputs, axis=1), axis=1
              )
          ),
          'labels': labels[self.label_field],
      })
    else:
      logs.update({
          'sentence_prediction': outputs,
          'labels': labels[self.label_field],
      })
    return logs
//  [python]
def validation_step(self, inputs, model: tf_keras.Model, metrics=None):
    features, labels = inputs, inputs
    outputs = self.inference_step(features, model)
    loss = self.build_losses(
        labels=labels, model_outputs=outputs, aux_losses=model.losses
    )
    logs = {self.loss: loss}
    if metrics:
      self.process_metrics(metrics, labels, outputs)
    if model.compiled_metrics:
      self.process_compiled_metrics(model.compiled_metrics, labels, outputs)
      logs.update({m.name: m.result() for m in metrics or []})
      logs.update({m.name: m.result() for m in model.metrics})
    if self.metric_type == 'matthews_corrcoef':
      logs.update({
          'sentence_prediction': (
              tf.expand_dims(  # Ensure one prediction along batch dimension.
                  tf.math.argmax(outputs, axis=1), axis=1
              )
          ),
          'labels': labels[self.label_field],
      })
    else:
      logs.update({
          'sentence_prediction': outputs,
          'labels': labels[self.label_field],
      })
    return logs
//  [python]
def argmax(in_tensor: Tensor, axis: Union[int, Iterable, tuple[int], None] = None):
    if BACKEND == BackEnd.NUMPY:
        return numpy.argmax(a=in_tensor, axis=axis)
    elif BACKEND == BackEnd.TENSORFLOW:
        return tf.argmax(input=in_tensor, axis=axis)
    else:
        raise NotImplementedError
//  [python]
def argmax(in_tensor: Tensor, axis: Union[int, Iterable, tuple[int], None] = None):
    if BACKEND == BackEnd.NUMPY:
        return numpy.argmax(a=in_tensor, axis=axis)
    elif BACKEND == BackEnd.TENSORFLOW:
        return tensorflow.math.argmax(input_tensor=in_tensor, axis=axis)
    else:
        raise NotImplementedError
//  [python]
def argmax(x, axis=None, keepdim=False, dtype='int64'):
    """
    Returns the index with the largest value across axes of a tensor.

    Parameters
    ----------
    x : tensor
        A Tensor
    axis : int
        An integer, the axis to reduce across. Default to 0.
    dtype : tensor or str
        An optional output dtype (nt32 or int64). Defaults to int64.

    Returns
    -------
        A Tensor of type output_type.

    Examples
    ---------
    >>> import tensorlayerx as tlx
    >>> x = tlx.ops.constant(value=[10, 20, 5, 6, 15])
    >>> y = tlx.ops.argmax(x)

    """

    return t.argmax(x, axis=axis, output_type=dtype_str(dtype))
//  [python]
def argmax(x, axis=None, keepdim=False, dtype='int64'):
    """
    Returns the index with the largest value across axes of a tensor.

    Parameters
    ----------
    x : tensor
        A Tensor
    axis : int
        An integer, the axis to reduce across. Default to 0.
    dtype : tensor or str
        An optional output dtype (nt32 or int64). Defaults to int64.

    Returns
    -------
        A Tensor of type output_type.

    Examples
    ---------
    >>> import tensorlayerx as tlx
    >>> x = tlx.ops.constant(value=[10, 20, 5, 6, 15])
    >>> y = tlx.ops.argmax(x)

    """

    return tf.math.argmax(x, axis=axis, output_type=dtype_str(dtype))
//  [python]
def mode(self, *, parameters, independent):
        if self.temperature_mode is None:
            probabilities, action_values = parameters.get(('probabilities', 'action_values'))
        else:
            probabilities, temperature, action_values = parameters.get(
                ('probabilities', 'temperature', 'action_values')
            )

        # Distribution parameter summaries
        dependencies = list()
        if not independent:
            def fn_summary():
                axis = range(self.action_spec.rank + 1)
                probs = tf.math.reduce_mean(input_tensor=probabilities, axis=axis)
                probs = [probs[n] for n in range(self.action_spec.num_values)]
                if self.temperature_mode is not None:
                    probs.append(tf.math.reduce_mean(input_tensor=temperature, axis=axis))
                return probs

            prefix = 'distributions/' + self.name + '-probability'
            names = [prefix + str(n) for n in range(self.action_spec.num_values)]
            if self.temperature_mode is not None:
                names.append('distributions/' + self.name + '-temperature')
            dependencies.extend(self.summary(
                label='distribution', name=names, data=fn_summary, step='timesteps'
            ))

        # Distribution parameter tracking
        def fn_tracking():
            return tf.math.reduce_mean(input_tensor=probabilities, axis=0)

        dependencies.extend(
            self.track(label='distribution', name='probabilities', data=fn_tracking)
        )

        if self.temperature_mode is not None:

            def fn_tracking():
                return tf.math.reduce_mean(input_tensor=temperature, axis=0)

            dependencies.extend(
                self.track(label='distribution', name='temperature', data=fn_tracking)
            )

        with tf.control_dependencies(control_inputs=dependencies):
            action = tf.argmax(input=action_values, axis=-1)
            return tf_util.cast(x=action, dtype='int')
//  [python]
def mode(self, *, parameters, independent):
        if self.temperature_mode is None:
            probabilities, action_values = parameters.get(('probabilities', 'action_values'))
        else:
            probabilities, temperature, action_values = parameters.get(
                ('probabilities', 'temperature', 'action_values')
            )

        # Distribution parameter summaries
        dependencies = list()
        if not independent:
            def fn_summary():
                axis = range(self.action_spec.rank + 1)
                probs = tf.math.reduce_mean(input_tensor=probabilities, axis=axis)
                probs = [probs[n] for n in range(self.action_spec.num_values)]
                if self.temperature_mode is not None:
                    probs.append(tf.math.reduce_mean(input_tensor=temperature, axis=axis))
                return probs

            prefix = 'distributions/' + self.name + '-probability'
            names = [prefix + str(n) for n in range(self.action_spec.num_values)]
            if self.temperature_mode is not None:
                names.append('distributions/' + self.name + '-temperature')
            dependencies.extend(self.summary(
                label='distribution', name=names, data=fn_summary, step='timesteps'
            ))

        # Distribution parameter tracking
        def fn_tracking():
            return tf.math.reduce_mean(input_tensor=probabilities, axis=0)

        dependencies.extend(
            self.track(label='distribution', name='probabilities', data=fn_tracking)
        )

        if self.temperature_mode is not None:

            def fn_tracking():
                return tf.math.reduce_mean(input_tensor=temperature, axis=0)

            dependencies.extend(
                self.track(label='distribution', name='temperature', data=fn_tracking)
            )

        with tf.control_dependencies(control_inputs=dependencies):
            action = tf.math.argmax(input=action_values, axis=-1)
            return tf_util.cast(x=action, dtype='int')
//  [python]
def __call__(
        self,
        logits: tf.Tensor,
    ) -> List[Tuple[str, float]]:
        # compute pred with argmax for attention models
        out_idxs = tf.argmax(logits, axis=2)
        preds_prob = tf.math.reduce_max(tf.nn.softmax(logits, axis=-1), axis=-1)

        # decode raw output of the model with tf_label_to_idx
        out_idxs = tf.cast(out_idxs, dtype="int32")
        embedding = tf.constant(self._embedding, dtype=tf.string)
        decoded_strings_pred = tf.strings.reduce_join(inputs=tf.nn.embedding_lookup(embedding, out_idxs), axis=-1)
        decoded_strings_pred = tf.strings.split(decoded_strings_pred, "<eos>")
        decoded_strings_pred = tf.sparse.to_dense(decoded_strings_pred.to_sparse(), default_value="not valid")[:, 0]
        word_values = [word.decode() for word in decoded_strings_pred.numpy().tolist()]

        # compute probabilties for each word up to the EOS token
        probs = [
            preds_prob[i, : len(word)].numpy().clip(0, 1).mean().item() if word else 0.0
            for i, word in enumerate(word_values)
        ]

        return list(zip(word_values, probs))
//  [python]
def __call__(
        self,
        logits: tf.Tensor,
    ) -> List[Tuple[str, float]]:
        # compute pred with argmax for attention models
        out_idxs = tf.math.argmax(logits, axis=2)
        preds_prob = tf.math.reduce_max(tf.nn.softmax(logits, axis=-1), axis=-1)

        # decode raw output of the model with tf_label_to_idx
        out_idxs = tf.cast(out_idxs, dtype="int32")
        embedding = tf.constant(self._embedding, dtype=tf.string)
        decoded_strings_pred = tf.strings.reduce_join(inputs=tf.nn.embedding_lookup(embedding, out_idxs), axis=-1)
        decoded_strings_pred = tf.strings.split(decoded_strings_pred, "<eos>")
        decoded_strings_pred = tf.sparse.to_dense(decoded_strings_pred.to_sparse(), default_value="not valid")[:, 0]
        word_values = [word.decode() for word in decoded_strings_pred.numpy().tolist()]

        # compute probabilties for each word up to the EOS token
        probs = [
            preds_prob[i, : len(word)].numpy().clip(0, 1).mean().item() if word else 0.0
            for i, word in enumerate(word_values)
        ]

        return list(zip(word_values, probs))
//  [python]
def topk_class_boxes(params, cls_outputs: T,
                     box_outputs: T) -> Tuple[T, T, T, T]:
  """Pick the topk class and box outputs."""
  batch_size = tf.shape(cls_outputs)[0]
  num_classes = params['num_classes']

  max_nms_inputs = params['nms_configs'].get('max_nms_inputs', 0)
  if max_nms_inputs > 0:
    # Prune anchors and detections to only keep max_nms_inputs.
    # Due to some issues, top_k is currently slow in graph model.
    logging.info('use max_nms_inputs for pre-nms topk.')
    cls_outputs_reshape = tf.reshape(cls_outputs, [batch_size, -1])
    _, cls_topk_indices = tf.math.top_k(
        cls_outputs_reshape, k=max_nms_inputs, sorted=False)
    indices = cls_topk_indices // num_classes
    classes = cls_topk_indices % num_classes
    cls_indices = tf.stack([indices, classes], axis=2)

    cls_outputs_topk = tf.gather_nd(cls_outputs, cls_indices, batch_dims=1)
    box_outputs_topk = tf.gather_nd(
        box_outputs, tf.expand_dims(indices, 2), batch_dims=1)
  else:
    logging.info('use max_reduce for pre-nms topk.')
    # Keep all anchors, but for each anchor, just keep the max probablity for
    # each class.
    cls_outputs_idx = tf.argmax(cls_outputs, axis=-1, output_type=tf.int32)
    num_anchors = tf.shape(cls_outputs)[1]

    classes = cls_outputs_idx
    indices = tf.tile(
        tf.expand_dims(tf.range(num_anchors), axis=0), [batch_size, 1])
    cls_outputs_topk = tf.reduce_max(cls_outputs, -1)
    box_outputs_topk = box_outputs

  return cls_outputs_topk, box_outputs_topk, classes, indices
//  [python]
def topk_class_boxes(params, cls_outputs: T,
                     box_outputs: T) -> Tuple[T, T, T, T]:
  """Pick the topk class and box outputs."""
  batch_size = tf.shape(cls_outputs)[0]
  num_classes = params['num_classes']

  max_nms_inputs = params['nms_configs'].get('max_nms_inputs', 0)
  if max_nms_inputs > 0:
    # Prune anchors and detections to only keep max_nms_inputs.
    # Due to some issues, top_k is currently slow in graph model.
    logging.info('use max_nms_inputs for pre-nms topk.')
    cls_outputs_reshape = tf.reshape(cls_outputs, [batch_size, -1])
    _, cls_topk_indices = tf.math.top_k(
        cls_outputs_reshape, k=max_nms_inputs, sorted=False)
    indices = cls_topk_indices // num_classes
    classes = cls_topk_indices % num_classes
    cls_indices = tf.stack([indices, classes], axis=2)

    cls_outputs_topk = tf.gather_nd(cls_outputs, cls_indices, batch_dims=1)
    box_outputs_topk = tf.gather_nd(
        box_outputs, tf.expand_dims(indices, 2), batch_dims=1)
  else:
    logging.info('use max_reduce for pre-nms topk.')
    # Keep all anchors, but for each anchor, just keep the max probablity for
    # each class.
    cls_outputs_idx = tf.math.argmax(cls_outputs, axis=-1, output_type=tf.int32)
    num_anchors = tf.shape(cls_outputs)[1]

    classes = cls_outputs_idx
    indices = tf.tile(
        tf.expand_dims(tf.range(num_anchors), axis=0), [batch_size, 1])
    cls_outputs_topk = tf.reduce_max(cls_outputs, -1)
    box_outputs_topk = box_outputs

  return cls_outputs_topk, box_outputs_topk, classes, indices
//  [python]
def postprocess(outputs: dict[str, tf.Tensor]) -> dict[str, tf.Tensor]:
  """Performs post-processing on model output.

  Args:
    outputs: The raw model output.

  Returns:
    Postprocessed model output.
  """
  predictions = {
      "detection_boxes":  # Box coordinates are relative values here.
          box_ops.cycxhw_to_yxyx(outputs["box_outputs"]),
      "detection_scores":
          tf.math.reduce_max(
              tf.nn.softmax(outputs["cls_outputs"])[:, :, 1:], axis=-1),
      "detection_classes":
          tf.argmax(outputs["cls_outputs"][:, :, 1:], axis=-1) + 1,
      # Fix this. It's not being used at the moment.
      "num_detections":
          tf.reduce_sum(
              tf.cast(
                  tf.math.greater(
                      tf.math.reduce_max(outputs["cls_outputs"], axis=-1), 0),
                  tf.int32),
              axis=-1)
  }
  return predictions
//  [python]
def postprocess(outputs: dict[str, tf.Tensor]) -> dict[str, tf.Tensor]:
  """Performs post-processing on model output.

  Args:
    outputs: The raw model output.

  Returns:
    Postprocessed model output.
  """
  predictions = {
      "detection_boxes":  # Box coordinates are relative values here.
          box_ops.cycxhw_to_yxyx(outputs["box_outputs"]),
      "detection_scores":
          tf.math.reduce_max(
              tf.nn.softmax(outputs["cls_outputs"])[:, :, 1:], axis=-1),
      "detection_classes":
          tf.math.argmax(outputs["cls_outputs"][:, :, 1:], axis=-1) + 1,
      # Fix this. It's not being used at the moment.
      "num_detections":
          tf.reduce_sum(
              tf.cast(
                  tf.math.greater(
                      tf.math.reduce_max(outputs["cls_outputs"], axis=-1), 0),
                  tf.int32),
              axis=-1)
  }
  return predictions
//  [python]
def evaluate_cset(infiller, cset, char2idx, max_length):
  """Evaluates given cognate set."""
  tgt_index = 0
  inputs = []
  input_mask = []
  # Find possible target positions
  for i, p in enumerate(cset):
    if p.strip():
      if p == '<TARGET>':
        tgt_index = i
        inputs.append([char2idx['<TARGET>']] * max_length)
        input_mask.append([0.0] * max_length)
      else:
        seq = [char2idx['<BOS>']] + [
            char2idx[c] if c in char2idx else char2idx['<UNK>']
            for c in p.split()
        ] + [char2idx['<EOS>']]
        template = [char2idx['<BLANK>']] * max_length
        for j in range(min(len(seq), max_length)):
          template[j] = seq[j]
        inputs.append(template)
        input_mask.append([1.0] * max_length)
    else:
      inputs.append([char2idx['<BLANK>']] * max_length)
      input_mask.append([0.0] * max_length)

  inputs = tf.constant([inputs], dtype='int32')
  input_mask = tf.constant([input_mask], dtype='float32')

  logits = infiller(inputs, input_mask, training=False)
  trow = tf.argmax(logits[0, tgt_index, :, :], axis=-1)
  return trow.numpy()
//  [python]
def evaluate_cset(infiller, cset, char2idx, max_length):
  """Evaluates given cognate set."""
  tgt_index = 0
  inputs = []
  input_mask = []
  # Find possible target positions
  for i, p in enumerate(cset):
    if p.strip():
      if p == '<TARGET>':
        tgt_index = i
        inputs.append([char2idx['<TARGET>']] * max_length)
        input_mask.append([0.0] * max_length)
      else:
        seq = [char2idx['<BOS>']] + [
            char2idx[c] if c in char2idx else char2idx['<UNK>']
            for c in p.split()
        ] + [char2idx['<EOS>']]
        template = [char2idx['<BLANK>']] * max_length
        for j in range(min(len(seq), max_length)):
          template[j] = seq[j]
        inputs.append(template)
        input_mask.append([1.0] * max_length)
    else:
      inputs.append([char2idx['<BLANK>']] * max_length)
      input_mask.append([0.0] * max_length)

  inputs = tf.constant([inputs], dtype='int32')
  input_mask = tf.constant([input_mask], dtype='float32')

  logits = infiller(inputs, input_mask, training=False)
  trow = tf.math.argmax(logits[0, tgt_index, :, :], axis=-1)
  return trow.numpy()
//  [python]
def build_mock_model(path):
  """Build and save the mock model with the given signature."""
  module = tf.Module()
  # We have to set this useless variable in order for the TF C API to correctly
  # intake it
  module.var = tf.Variable(0, dtype=tf.int64)

  def action(*inputs):
    result = tf.argmax(
        tf.cast(inputs[0]['mask'], tf.int32), axis=-1) + module.var
    return {POLICY_DECISION_LABEL: result}
  module.action = tf.function()(action)
  action = {
      'action': module.action.get_concrete_function(get_input_signature())
  }
  tf.saved_model.save(module, path, signatures=action)
  output_spec_path = get_output_spec_path(path)
  with open(output_spec_path, 'w') as f:
    print(f'Writing output spec to {output_spec_path}.')
    f.write(POLICY_OUTPUT_SPEC)
//  [python]
def build_mock_model(path):
  """Build and save the mock model with the given signature."""
  module = tf.Module()
  # We have to set this useless variable in order for the TF C API to correctly
  # intake it
  module.var = tf.Variable(0, dtype=tf.int64)

  def action(*inputs):
    result = tf.math.argmax(
        tf.cast(inputs[0]['mask'], tf.int32), axis=-1) + module.var
    return {POLICY_DECISION_LABEL: result}
  module.action = tf.function()(action)
  action = {
      'action': module.action.get_concrete_function(get_input_signature())
  }
  tf.saved_model.save(module, path, signatures=action)
  output_spec_path = get_output_spec_path(path)
  with open(output_spec_path, 'w') as f:
    print(f'Writing output spec to {output_spec_path}.')
    f.write(POLICY_OUTPUT_SPEC)
//  [python]
def build_losses(self, outputs, labels, aux_losses=None):
    """Builds DETR losses."""
    cls_outputs = outputs['cls_outputs']
    box_outputs = outputs['box_outputs']
    cls_targets = labels['classes']
    box_targets = labels['boxes']

    cost = self._compute_cost(
        cls_outputs, box_outputs, cls_targets, box_targets)

    _, indices = matchers.hungarian_matching(cost)
    indices = tf.stop_gradient(indices)

    target_index = tf.argmax(indices, axis=1)
    cls_assigned = tf.gather(cls_outputs, target_index, batch_dims=1, axis=1)
    box_assigned = tf.gather(box_outputs, target_index, batch_dims=1, axis=1)

    background = tf.equal(cls_targets, 0)
    num_boxes = tf.reduce_sum(
        tf.cast(tf.logical_not(background), tf.float32), axis=-1)

    # Down-weight background to account for class imbalance.
    xentropy = tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=cls_targets, logits=cls_assigned)
    cls_loss = self._task_config.losses.lambda_cls * tf.where(
        background, self._task_config.losses.background_cls_weight * xentropy,
        xentropy)
    cls_weights = tf.where(
        background,
        self._task_config.losses.background_cls_weight * tf.ones_like(cls_loss),
        tf.ones_like(cls_loss))

    # Box loss is only calculated on non-background class.
    l_1 = tf.reduce_sum(tf.abs(box_assigned - box_targets), axis=-1)
    box_loss = self._task_config.losses.lambda_box * tf.where(
        background, tf.zeros_like(l_1), l_1)

    # Giou loss is only calculated on non-background class.
    giou = tf.linalg.diag_part(1.0 - box_ops.bbox_generalized_overlap(
        box_ops.cycxhw_to_yxyx(box_assigned),
        box_ops.cycxhw_to_yxyx(box_targets)
        ))
    giou_loss = self._task_config.losses.lambda_giou * tf.where(
        background, tf.zeros_like(giou), giou)

    # Consider doing all reduce once in train_step to speed up.
    num_boxes_per_replica = tf.reduce_sum(num_boxes)
    cls_weights_per_replica = tf.reduce_sum(cls_weights)
    replica_context = tf.distribute.get_replica_context()
    num_boxes_sum, cls_weights_sum = replica_context.all_reduce(
        tf.distribute.ReduceOp.SUM,
        [num_boxes_per_replica, cls_weights_per_replica])
    cls_loss = tf.math.divide_no_nan(
        tf.reduce_sum(cls_loss), cls_weights_sum)
    box_loss = tf.math.divide_no_nan(
        tf.reduce_sum(box_loss), num_boxes_sum)
    giou_loss = tf.math.divide_no_nan(
        tf.reduce_sum(giou_loss), num_boxes_sum)

    aux_losses = tf.add_n(aux_losses) if aux_losses else 0.0

    total_loss = cls_loss + box_loss + giou_loss + aux_losses
    return total_loss, cls_loss, box_loss, giou_loss
//  [python]
def build_losses(self, outputs, labels, aux_losses=None):
    """Builds DETR losses."""
    cls_outputs = outputs['cls_outputs']
    box_outputs = outputs['box_outputs']
    cls_targets = labels['classes']
    box_targets = labels['boxes']

    cost = self._compute_cost(
        cls_outputs, box_outputs, cls_targets, box_targets)

    _, indices = matchers.hungarian_matching(cost)
    indices = tf.stop_gradient(indices)

    target_index = tf.math.argmax(indices, axis=1)
    cls_assigned = tf.gather(cls_outputs, target_index, batch_dims=1, axis=1)
    box_assigned = tf.gather(box_outputs, target_index, batch_dims=1, axis=1)

    background = tf.equal(cls_targets, 0)
    num_boxes = tf.reduce_sum(
        tf.cast(tf.logical_not(background), tf.float32), axis=-1)

    # Down-weight background to account for class imbalance.
    xentropy = tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=cls_targets, logits=cls_assigned)
    cls_loss = self._task_config.losses.lambda_cls * tf.where(
        background, self._task_config.losses.background_cls_weight * xentropy,
        xentropy)
    cls_weights = tf.where(
        background,
        self._task_config.losses.background_cls_weight * tf.ones_like(cls_loss),
        tf.ones_like(cls_loss))

    # Box loss is only calculated on non-background class.
    l_1 = tf.reduce_sum(tf.abs(box_assigned - box_targets), axis=-1)
    box_loss = self._task_config.losses.lambda_box * tf.where(
        background, tf.zeros_like(l_1), l_1)

    # Giou loss is only calculated on non-background class.
    giou = tf.linalg.diag_part(1.0 - box_ops.bbox_generalized_overlap(
        box_ops.cycxhw_to_yxyx(box_assigned),
        box_ops.cycxhw_to_yxyx(box_targets)
        ))
    giou_loss = self._task_config.losses.lambda_giou * tf.where(
        background, tf.zeros_like(giou), giou)

    # Consider doing all reduce once in train_step to speed up.
    num_boxes_per_replica = tf.reduce_sum(num_boxes)
    cls_weights_per_replica = tf.reduce_sum(cls_weights)
    replica_context = tf.distribute.get_replica_context()
    num_boxes_sum, cls_weights_sum = replica_context.all_reduce(
        tf.distribute.ReduceOp.SUM,
        [num_boxes_per_replica, cls_weights_per_replica])
    cls_loss = tf.math.divide_no_nan(
        tf.reduce_sum(cls_loss), cls_weights_sum)
    box_loss = tf.math.divide_no_nan(
        tf.reduce_sum(box_loss), num_boxes_sum)
    giou_loss = tf.math.divide_no_nan(
        tf.reduce_sum(giou_loss), num_boxes_sum)

    aux_losses = tf.add_n(aux_losses) if aux_losses else 0.0

    total_loss = cls_loss + box_loss + giou_loss + aux_losses
    return total_loss, cls_loss, box_loss, giou_loss
//  [python]
def _get_box_class_per_point(
    box: tf.Tensor, box_class: tf.Tensor, point_xyz: tf.Tensor
) -> tf.Tensor:
    """Extracts point labels.

    Args:
      box: [M, 7] box tensor.
      box_class: [M] class of each box.
      point_xyz: [N, 3] points.

    Returns:
      point_box_class: [N] box class of each point.
    """
    n = point_xyz.shape[0]
    m = box.shape[0]
    if m == 0:
        return tf.zeros([n], dtype=tf.int32)

    # [N, M]
    point_in_box = box_utils.is_within_box_3d(point_xyz, box)
    # [N]
    point_in_any_box = tf.math.reduce_any(point_in_box, axis=-1)
    # [N]
    point_box_idx = tf.argmax(point_in_box, axis=-1, output_type=tf.int32)
    # [N]
    point_box_class = tf.where(
        point_in_any_box, tf.gather(box_class, point_box_idx), 0
    )

    return point_box_class
//  [python]
def _get_box_class_per_point(
    box: tf.Tensor, box_class: tf.Tensor, point_xyz: tf.Tensor
) -> tf.Tensor:
    """Extracts point labels.

    Args:
      box: [M, 7] box tensor.
      box_class: [M] class of each box.
      point_xyz: [N, 3] points.

    Returns:
      point_box_class: [N] box class of each point.
    """
    n = point_xyz.shape[0]
    m = box.shape[0]
    if m == 0:
        return tf.zeros([n], dtype=tf.int32)

    # [N, M]
    point_in_box = box_utils.is_within_box_3d(point_xyz, box)
    # [N]
    point_in_any_box = tf.math.reduce_any(point_in_box, axis=-1)
    # [N]
    point_box_idx = tf.math.argmax(point_in_box, axis=-1, output_type=tf.int32)
    # [N]
    point_box_class = tf.where(
        point_in_any_box, tf.gather(box_class, point_box_idx), 0
    )

    return point_box_class
//  [python]
def validation_step(self, inputs, model: tf_keras.Model, metrics=None):
    features, labels = inputs, inputs
    outputs = self.inference_step(features, model)
    loss = self.build_losses(
        labels=labels, model_outputs=outputs, aux_losses=model.losses)
    logs = {self.loss: loss}
    if metrics:
      self.process_metrics(metrics, labels, outputs)
    if model.compiled_metrics:
      self.process_compiled_metrics(model.compiled_metrics, labels, outputs)
      logs.update({m.name: m.result() for m in metrics or []})
      logs.update({m.name: m.result() for m in model.metrics})
    if self.metric_type == 'matthews_corrcoef':
      logs.update({
          'sentence_prediction':  # Ensure one prediction along batch dimension.
              tf.expand_dims(tf.argmax(outputs, axis=1), axis=1),
          'labels':
              labels[self.label_field],
      })
    else:
      logs.update({
          'sentence_prediction': outputs,
          'labels': labels[self.label_field],
      })
    return logs
//  [python]
def validation_step(self, inputs, model: tf_keras.Model, metrics=None):
    features, labels = inputs, inputs
    outputs = self.inference_step(features, model)
    loss = self.build_losses(
        labels=labels, model_outputs=outputs, aux_losses=model.losses)
    logs = {self.loss: loss}
    if metrics:
      self.process_metrics(metrics, labels, outputs)
    if model.compiled_metrics:
      self.process_compiled_metrics(model.compiled_metrics, labels, outputs)
      logs.update({m.name: m.result() for m in metrics or []})
      logs.update({m.name: m.result() for m in model.metrics})
    if self.metric_type == 'matthews_corrcoef':
      logs.update({
          'sentence_prediction':  # Ensure one prediction along batch dimension.
              tf.expand_dims(tf.math.argmax(outputs, axis=1), axis=1),
          'labels':
              labels[self.label_field],
      })
    else:
      logs.update({
          'sentence_prediction': outputs,
          'labels': labels[self.label_field],
      })
    return logs
//  [python]
def balanced_sample(features,
                    instance_ids,
                    num_samples,
                    valid_mask=None,
                    max_instance_id=None):
  """Samples features by encouraging a balanced selections based on instance id.

  Args:
    features: A tf.float32 tensor of size [batch_size, n, f].
    instance_ids: A tf.int32 tensor of size [batch_size, n].
    num_samples: An int determining the number of samples.
    valid_mask: A tf.bool tensor of size [batch_size, n] that is True when an
      element is valid and False if it needs to be ignored. By default the value
      is None which means it is not applied.
    max_instance_id: If set, instance ids larger than that value will be
      ignored. If not set, it will be computed from instance_ids tensor.

  Returns:
    A tf.float32 tensor of size [batch_size, num_samples, f] containing sampled
      features.
    A tf.int32 tensor of size [batch_size, num_samples] containing sampled
      instance ids
    A tf.int32 tensor of size [batch_size, num_samples] containing sampled
      indices.
  """
  sampling_probs = get_balanced_sampling_probability(
      instance_ids=instance_ids,
      valid_mask=valid_mask,
      max_instance_id=max_instance_id)
  dist = tfp.distributions.OneHotCategorical(probs=sampling_probs)
  sample_obj = tfp.distributions.Sample(dist, sample_shape=num_samples)
  samples = sample_obj.sample()
  sampled_indices = tf.argmax(samples, axis=2)
  sampled_features = tf.gather(features, sampled_indices, batch_dims=1)
  sampled_instance_ids = tf.gather(instance_ids, sampled_indices, batch_dims=1)
  return sampled_features, sampled_instance_ids, sampled_indices
//  [python]
def balanced_sample(features,
                    instance_ids,
                    num_samples,
                    valid_mask=None,
                    max_instance_id=None):
  """Samples features by encouraging a balanced selections based on instance id.

  Args:
    features: A tf.float32 tensor of size [batch_size, n, f].
    instance_ids: A tf.int32 tensor of size [batch_size, n].
    num_samples: An int determining the number of samples.
    valid_mask: A tf.bool tensor of size [batch_size, n] that is True when an
      element is valid and False if it needs to be ignored. By default the value
      is None which means it is not applied.
    max_instance_id: If set, instance ids larger than that value will be
      ignored. If not set, it will be computed from instance_ids tensor.

  Returns:
    A tf.float32 tensor of size [batch_size, num_samples, f] containing sampled
      features.
    A tf.int32 tensor of size [batch_size, num_samples] containing sampled
      instance ids
    A tf.int32 tensor of size [batch_size, num_samples] containing sampled
      indices.
  """
  sampling_probs = get_balanced_sampling_probability(
      instance_ids=instance_ids,
      valid_mask=valid_mask,
      max_instance_id=max_instance_id)
  dist = tfp.distributions.OneHotCategorical(probs=sampling_probs)
  sample_obj = tfp.distributions.Sample(dist, sample_shape=num_samples)
  samples = sample_obj.sample()
  sampled_indices = tf.math.argmax(samples, axis=2)
  sampled_features = tf.gather(features, sampled_indices, batch_dims=1)
  sampled_instance_ids = tf.gather(instance_ids, sampled_indices, batch_dims=1)
  return sampled_features, sampled_instance_ids, sampled_indices
//  [python]
def call(self, dataset_iter, training=True, generate=False):
    observations, actions, maze_maps, value_maps = next(dataset_iter)

    states = self.process_states(observations, maze_maps, training=training)

    with tf.GradientTape() as tape:
      embed = self.embed_states(states)
      logit = self._policy(embed)
      pred_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
          labels=actions, logits=logit)
      pred = tf.argmax(logit, axis=-1, output_type=tf.dtypes.int32)
      acc = tf.reduce_sum(tf.cast(pred == actions, tf.float32)
                          ) / tf.cast(tf.shape(pred)[0], tf.float32)

      value_maps = tf.reshape(value_maps, [-1, self._num_actions])
      valid_indices = tf.where(tf.reduce_sum(value_maps, axis=-1) == 1)
      value_maps = tf.gather(value_maps, valid_indices)
      logit = self._action_network(embed)
      logit = tf.reshape(logit, [-1, self._num_actions])
      logit = tf.gather(logit, valid_indices)
      aux_loss = tf.nn.softmax_cross_entropy_with_logits(
          labels=value_maps, logits=logit)

      loss = tf.reduce_mean(pred_loss) + self._aux_weight * tf.reduce_mean(aux_loss)
    if training:
      grads = tape.gradient(loss, self.trainable_variables)
      self._optimizer.apply_gradients(zip(grads, self.trainable_variables))

    return {'loss': loss,
            'pred_loss': tf.reduce_mean(pred_loss),
            'aux_loss': tf.reduce_mean(aux_loss),
            'acc': acc
            }
//  [python]
def call(self, dataset_iter, training=True, generate=False):
    observations, actions, maze_maps, value_maps = next(dataset_iter)

    states = self.process_states(observations, maze_maps, training=training)

    with tf.GradientTape() as tape:
      embed = self.embed_states(states)
      logit = self._policy(embed)
      pred_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
          labels=actions, logits=logit)
      pred = tf.math.argmax(logit, axis=-1, output_type=tf.dtypes.int32)
      acc = tf.reduce_sum(tf.cast(pred == actions, tf.float32)
                          ) / tf.cast(tf.shape(pred)[0], tf.float32)

      value_maps = tf.reshape(value_maps, [-1, self._num_actions])
      valid_indices = tf.where(tf.reduce_sum(value_maps, axis=-1) == 1)
      value_maps = tf.gather(value_maps, valid_indices)
      logit = self._action_network(embed)
      logit = tf.reshape(logit, [-1, self._num_actions])
      logit = tf.gather(logit, valid_indices)
      aux_loss = tf.nn.softmax_cross_entropy_with_logits(
          labels=value_maps, logits=logit)

      loss = tf.reduce_mean(pred_loss) + self._aux_weight * tf.reduce_mean(aux_loss)
    if training:
      grads = tape.gradient(loss, self.trainable_variables)
      self._optimizer.apply_gradients(zip(grads, self.trainable_variables))

    return {'loss': loss,
            'pred_loss': tf.reduce_mean(pred_loss),
            'aux_loss': tf.reduce_mean(aux_loss),
            'acc': acc
            }
//  [python]
def add_embeddings(self):
        """
        embedding层
        """
        with tf.variable_scope('embedding'):
            if self.config.embeddings is not None:
                embeddings = tf.Variable(self.config.embeddings, 
                    name="embeddings", trainable=False)
            else:
                embeddings = tf.get_variable('embeddings', 
                    shape=[self.config.vocab_size, self.config.embedding_size], 
                    initializer=tf.uniform_unit_scaling_initializer())
            q_embed = tf.nn.embedding_lookup(embeddings, self.q)
            a_embed = tf.nn.embedding_lookup(embeddings, self.a)
            return q_embed, a_embed
//  [python]
def add_embeddings(self):
        """
        embedding层
        """
        with tf.variable_scope('embedding'):
            if self.config.embeddings is not None:
                embeddings = tf.Variable(self.config.embeddings, 
                    name="embeddings", trainable=False)
            else:
                embeddings = tf.get_variable('embeddings', 
                    shape=[self.config.vocab_size, self.config.embedding_size], 
                    initializer=tf.keras.initializers.VarianceScaling())
            q_embed = tf.nn.embedding_lookup(embeddings, self.q)
            a_embed = tf.nn.embedding_lookup(embeddings, self.a)
            return q_embed, a_embed
//  [python]
def build_mock_model(path):
    """Build and save the mock model with the given signature."""
    module = tf.Module()
    # We have to set this useless variable in order for the TF C API to correctly
    # intake it
    module.var = tf.Variable(0, dtype=tf.int64)

    def action(*inputs):
        result = (
            tf.argmax(tf.cast(inputs[0]["mask"], tf.int32), axis=-1) + module.var
        )
        return {POLICY_DECISION_LABEL: result}

    module.action = tf.function()(action)
    action = {"action": module.action.get_concrete_function(get_input_signature())}
    tf.saved_model.save(module, path, signatures=action)
    output_spec_path = get_output_spec_path(path)
    with open(output_spec_path, "w") as f:
        print(f"Writing output spec to {output_spec_path}.")
        f.write(POLICY_OUTPUT_SPEC)
//  [python]
def build_mock_model(path):
    """Build and save the mock model with the given signature."""
    module = tf.Module()
    # We have to set this useless variable in order for the TF C API to correctly
    # intake it
    module.var = tf.Variable(0, dtype=tf.int64)

    def action(*inputs):
        result = (
            tf.math.argmax(tf.cast(inputs[0]["mask"], tf.int32), axis=-1) + module.var
        )
        return {POLICY_DECISION_LABEL: result}

    module.action = tf.function()(action)
    action = {"action": module.action.get_concrete_function(get_input_signature())}
    tf.saved_model.save(module, path, signatures=action)
    output_spec_path = get_output_spec_path(path)
    with open(output_spec_path, "w") as f:
        print(f"Writing output spec to {output_spec_path}.")
        f.write(POLICY_OUTPUT_SPEC)
//  [python]
def _get_ground_truth(self):
        """
        Returns:
        """
        res = []

        for i,outputs in enumerate(self.head_outputs):
            shape = wmlt.combined_static_and_dynamic_shape(outputs['heatmaps_ct'])[1:3]
            t_res = self.box2box_transform.get_deltas(self.gt_boxes,
                                                tf.ones_like(self.gt_labels),
                                                self.gt_length,
                                                output_size=shape)
            g_heatmaps_c, _, _= tfop.center2_boxes_encode(self.gt_boxes,
                                                         self.gt_labels,
                                                         self.gt_length,
                                                         shape,
                                                         global_cfg.MODEL.MOT.FAIR_MOT_NUM_CLASSES,
                                                         0.7)
            B,H,W,_ = btf.combined_static_and_dynamic_shape(g_heatmaps_c)
            bg_prob = tf.ones([B,H,W,1])*0.5
            heatmaps = tf.concat([bg_prob,g_heatmaps_c],axis=-1)
            classes = tf.argmax(heatmaps,axis=-1)
            t_res['id_classes'] = classes
            res.append(t_res)
        return res
//  [python]
def _get_ground_truth(self):
        """
        Returns:
        """
        res = []

        for i,outputs in enumerate(self.head_outputs):
            shape = wmlt.combined_static_and_dynamic_shape(outputs['heatmaps_ct'])[1:3]
            t_res = self.box2box_transform.get_deltas(self.gt_boxes,
                                                tf.ones_like(self.gt_labels),
                                                self.gt_length,
                                                output_size=shape)
            g_heatmaps_c, _, _= tfop.center2_boxes_encode(self.gt_boxes,
                                                         self.gt_labels,
                                                         self.gt_length,
                                                         shape,
                                                         global_cfg.MODEL.MOT.FAIR_MOT_NUM_CLASSES,
                                                         0.7)
            B,H,W,_ = btf.combined_static_and_dynamic_shape(g_heatmaps_c)
            bg_prob = tf.ones([B,H,W,1])*0.5
            heatmaps = tf.concat([bg_prob,g_heatmaps_c],axis=-1)
            classes = tf.math.argmax(heatmaps,axis=-1)
            t_res['id_classes'] = classes
            res.append(t_res)
        return res
//  [python]
def main(unused_argv):
  g = load_graph()
  g.init()

  # Define Model
  dimensions = [FLAGS.features_num] + [FLAGS.hidden_dim] * (FLAGS.hops_num - 1) + [FLAGS.class_num]
  model = EgoGraphSAGE(dimensions, act_func=tf.nn.relu, dropout=FLAGS.in_drop_rate)

  # prepare train dataset
  train_data = EgoSAGESupervisedDataLoader(
    g, gl.Mask.TRAIN,
    node_type=FLAGS.node_type, edge_type=FLAGS.edge_type,
    nbrs_num=nbrs_num, hops_num=FLAGS.hops_num,
  )
  train_embedding = model.forward(train_data.src_ego)
  train_labels = train_data.src_ego.src.labels
  loss = tf.reduce_mean(
    tf.nn.sparse_softmax_cross_entropy_with_logits(
      labels=train_labels, logits=train_embedding,
    )
  )
  optimizer = tf.train.AdamOptimizer(learning_rate=FLAGS.learning_rate)

  # prepare test dataset
  test_data = EgoSAGESupervisedDataLoader(
    g, gl.Mask.TEST,
    node_type=FLAGS.node_type, edge_type=FLAGS.edge_type,
    nbrs_num=nbrs_num, hops_num=FLAGS.hops_num,
  )
  test_embedding = model.forward(test_data.src_ego)
  test_labels = test_data.src_ego.src.labels
  test_indices = tf.argmax(test_embedding, 1, output_type=tf.int32)
  test_acc = tf.div(
    tf.reduce_sum(tf.cast(tf.equal(test_indices, test_labels), tf.float32)),
    tf.cast(tf.shape(test_labels)[0], tf.float32),
  )

  # train and test
  trainer = LocalTrainer()
  trainer.train(train_data.iterator, loss, optimizer, epochs=FLAGS.epochs)
  trainer.test(test_data.iterator, test_acc)

  # finish
  g.close()
//  [python]
def main(unused_argv):
  g = load_graph()
  g.init()

  # Define Model
  dimensions = [FLAGS.features_num] + [FLAGS.hidden_dim] * (FLAGS.hops_num - 1) + [FLAGS.class_num]
  model = EgoGraphSAGE(dimensions, act_func=tf.nn.relu, dropout=FLAGS.in_drop_rate)

  # prepare train dataset
  train_data = EgoSAGESupervisedDataLoader(
    g, gl.Mask.TRAIN,
    node_type=FLAGS.node_type, edge_type=FLAGS.edge_type,
    nbrs_num=nbrs_num, hops_num=FLAGS.hops_num,
  )
  train_embedding = model.forward(train_data.src_ego)
  train_labels = train_data.src_ego.src.labels
  loss = tf.reduce_mean(
    tf.nn.sparse_softmax_cross_entropy_with_logits(
      labels=train_labels, logits=train_embedding,
    )
  )
  optimizer = tf.train.AdamOptimizer(learning_rate=FLAGS.learning_rate)

  # prepare test dataset
  test_data = EgoSAGESupervisedDataLoader(
    g, gl.Mask.TEST,
    node_type=FLAGS.node_type, edge_type=FLAGS.edge_type,
    nbrs_num=nbrs_num, hops_num=FLAGS.hops_num,
  )
  test_embedding = model.forward(test_data.src_ego)
  test_labels = test_data.src_ego.src.labels
  test_indices = tf.math.argmax(test_embedding, 1, output_type=tf.int32)
  test_acc = tf.div(
    tf.reduce_sum(tf.cast(tf.math.equal(test_indices, test_labels), tf.float32)),
    tf.cast(tf.shape(test_labels)[0], tf.float32),
  )

  # train and test
  trainer = LocalTrainer()
  trainer.train(train_data.iterator, loss, optimizer, epochs=FLAGS.epochs)
  trainer.test(test_data.iterator, test_acc)

  # finish
  g.close()
//  [python]
def init_vq_bottleneck(bottleneck_size, hidden_size):
  """Get lookup table for VQ bottleneck."""
  means = tf.get_variable(
      name="means",
      shape=[bottleneck_size, hidden_size],
      initializer=tf.uniform_unit_scaling_initializer())
  ema_count = tf.get_variable(
      name="ema_count",
      shape=[bottleneck_size],
      initializer=tf.constant_initializer(0),
      trainable=False)
  with tf.colocate_with(means):
    ema_means = tf.get_variable(
        name="ema_means",
        initializer=means.initialized_value(),
        trainable=False)

  return means, ema_means, ema_count
//  [python]
def init_vq_bottleneck(bottleneck_size, hidden_size):
  """Get lookup table for VQ bottleneck."""
  means = tf.get_variable(
      name="means",
      shape=[bottleneck_size, hidden_size],
      initializer=tf.initializers.variance_scaling())
  ema_count = tf.get_variable(
      name="ema_count",
      shape=[bottleneck_size],
      initializer=tf.constant_initializer(0),
      trainable=False)
  with tf.colocate_with(means):
    ema_means = tf.get_variable(
        name="ema_means",
        initializer=means.initialized_value(),
        trainable=False)

  return means, ema_means, ema_count
//  [python]
def add_embeddings(self):
        with tf.variable_scope('embedding'):
            if self.config.embeddings is not None:
                embeddings = tf.Variable(self.config.embeddings, name="embeddings", trainable=False)
            else:
                embeddings = tf.get_variable('embeddings', shape=[self.config.vocab_size, self.config.embedding_size], initializer=tf.uniform_unit_scaling_initializer())
            q_embed = tf.nn.embedding_lookup(embeddings, self.q)
            a_embed = tf.nn.embedding_lookup(embeddings, self.a)
            return q_embed, a_embed
//  [python]
def add_embeddings(self):
        with tf.variable_scope('embedding'):
            if self.config.embeddings is not None:
                embeddings = tf.Variable(self.config.embeddings, name="embeddings", trainable=False)
            else:
                embeddings = tf.get_variable('embeddings', shape=[self.config.vocab_size, self.config.embedding_size], initializer=tf.keras.initializers.VarianceScaling())
            q_embed = tf.nn.embedding_lookup(embeddings, self.q)
            a_embed = tf.nn.embedding_lookup(embeddings, self.a)
            return q_embed, a_embed