jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _noam_learning_rate_decay(self, init_lr, global_step, warmup_steps=4000.0):
		# Noam scheme from tensor2tensor:
		step = t.cast(global_step + 1, dtype=t.float32)
		return tf.maximum(init_lr * warmup_steps**0.5 * tf.minimum(step * warmup_steps**-1.5, step**-0.5), 1e-4)
//  [python]
def _noam_learning_rate_decay(self, init_lr, global_step, warmup_steps=4000.0):
		# Noam scheme from tensor2tensor:
		step = tf.cast(global_step + 1, dtype=tf.float32)
		return tf.maximum(init_lr * warmup_steps**0.5 * tf.minimum(step * warmup_steps**-1.5, step**-0.5), 1e-4)
//  [python]
def _convert_head_mask_to_5d(self, head_mask, num_hidden_layers):
        """-> [num_hidden_layers x batch x num_heads x seq_length x seq_length]"""
        if head_mask.shape.rank == 1:
            head_mask = head_mask[None, None, :, None, None]
            head_mask = tf.repeat(head_mask, repeats=num_hidden_layers, axis=0)
        elif head_mask.shape.rank == 2:
            head_mask = head_mask[:, None, :, None, None]
        assert head_mask.shape.rank == 5, f"head_mask.dim != 5, instead {head_mask.dim()}"
        head_mask = tf.to_float(head_mask)  # switch to float if need + fp16 compatibility
        return head_mask
//  [python]
def _convert_head_mask_to_5d(self, head_mask, num_hidden_layers):
        """-> [num_hidden_layers x batch x num_heads x seq_length x seq_length]"""
        if head_mask.shape.rank == 1:
            head_mask = head_mask[None, None, :, None, None]
            head_mask = tf.repeat(head_mask, repeats=num_hidden_layers, axis=0)
        elif head_mask.shape.rank == 2:
            head_mask = head_mask[:, None, :, None, None]
        assert head_mask.shape.rank == 5, f"head_mask.dim != 5, instead {head_mask.dim()}"
        head_mask = tf.cast(head_mask, tf.float32)  # switch to float if need + fp16 compatibility
        return head_mask
//  [python]
def build_losses(self,
                   labels,
                   model_outputs,
                   metrics,
                   aux_losses=None) -> tf.Tensor:
    metrics = dict([(metric.name, metric) for metric in metrics])

    # generator lm and (optional) nsp loss.
    lm_prediction_losses = tf_keras.losses.sparse_categorical_crossentropy(
        labels['masked_lm_ids'],
        tf.cast(model_outputs['lm_outputs'], tf.float32),
        from_logits=True)
    lm_label_weights = labels['masked_lm_weights']
    lm_numerator_loss = tf.reduce_sum(lm_prediction_losses * lm_label_weights)
    lm_denominator_loss = tf.reduce_sum(lm_label_weights)
    mlm_loss = tf.math.divide_no_nan(lm_numerator_loss, lm_denominator_loss)
    metrics['lm_example_loss'].update_state(mlm_loss)
    if 'next_sentence_labels' in labels:
      sentence_labels = labels['next_sentence_labels']
      sentence_outputs = tf.cast(
          model_outputs['sentence_outputs'], dtype=tf.float32)
      sentence_loss = tf_keras.losses.sparse_categorical_crossentropy(
          sentence_labels, sentence_outputs, from_logits=True)
      metrics['next_sentence_loss'].update_state(sentence_loss)
      total_loss = mlm_loss + sentence_loss
    else:
      total_loss = mlm_loss

    # discriminator replaced token detection (rtd) loss.
    rtd_logits = model_outputs['disc_logits']
    rtd_labels = tf.cast(model_outputs['disc_label'], tf.float32)
    input_mask = tf.cast(labels['input_mask'], tf.float32)
    rtd_ind_loss = tf.nn.sigmoid_cross_entropy_with_logits(
        logits=rtd_logits, labels=rtd_labels)
    rtd_numerator = tf.reduce_sum(input_mask * rtd_ind_loss)
    rtd_denominator = tf.reduce_sum(input_mask)
    rtd_loss = tf.math.divide_no_nan(rtd_numerator, rtd_denominator)
    metrics['discriminator_loss'].update_state(rtd_loss)
    total_loss = total_loss + \
        self.task_config.model.discriminator_loss_weight * rtd_loss

    if aux_losses:
      total_loss += tf.add_n(aux_losses)

    metrics['total_loss'].update_state(total_loss)
    return total_loss
//  [python]
def hf_compute_loss(self, labels, logits):
        loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction=keras.losses.Reduction.NONE)
        if tf.executing_eagerly():  # Data-dependent conditionals are forbidden in XLA
            if tf.math.reduce_any(labels == -1):
                tf.print("Using `-1` to mask the loss for the token is deprecated. Please use `-100` instead.")

        if self.config.tf_legacy_loss:
            # make sure only labels that are not equal to -100
            # are taken into account as loss
            if tf.math.reduce_any(labels == -1):
                tf.print("Using `-1` to mask the loss for the token is deprecated. Please use `-100` instead.")
                active_loss = tf.reshape(labels, (-1,)) != -1
            else:
                active_loss = tf.reshape(labels, (-1,)) != -100
            reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, shape_list(logits)[2])), active_loss)
            labels = tf.boolean_mask(tf.reshape(labels, (-1,)), active_loss)

            return loss_fn(labels, reduced_logits)

        # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway
        unmasked_loss = loss_fn(tf.nn.relu(labels), logits)
        # make sure only labels that are not equal to -100 or -1
        # are taken into account as loss
        loss_mask = tf.to_float(labels >= 0)
        # Avoid possible division by zero later
        # Masked positions will have a loss of NaN because -100 and -1 are not valid labels
        masked_loss = unmasked_loss * loss_mask
        reduced_masked_loss = tf.reduce_sum(masked_loss) / tf.reduce_sum(loss_mask)
        return tf.reshape(reduced_masked_loss, (1,))
//  [python]
def hf_compute_loss(self, labels, logits):
        loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction=keras.losses.Reduction.NONE)
        if tf.executing_eagerly():  # Data-dependent conditionals are forbidden in XLA
            if tf.math.reduce_any(labels == -1):
                tf.print("Using `-1` to mask the loss for the token is deprecated. Please use `-100` instead.")

        if self.config.tf_legacy_loss:
            # make sure only labels that are not equal to -100
            # are taken into account as loss
            if tf.math.reduce_any(labels == -1):
                tf.print("Using `-1` to mask the loss for the token is deprecated. Please use `-100` instead.")
                active_loss = tf.reshape(labels, (-1,)) != -1
            else:
                active_loss = tf.reshape(labels, (-1,)) != -100
            reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, shape_list(logits)[2])), active_loss)
            labels = tf.boolean_mask(tf.reshape(labels, (-1,)), active_loss)

            return loss_fn(labels, reduced_logits)

        # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway
        unmasked_loss = loss_fn(tf.nn.relu(labels), logits)
        # make sure only labels that are not equal to -100 or -1
        # are taken into account as loss
        loss_mask = tf.cast(labels >= 0, dtype=unmasked_loss.dtype)
        # Avoid possible division by zero later
        # Masked positions will have a loss of NaN because -100 and -1 are not valid labels
        masked_loss = unmasked_loss * loss_mask
        reduced_masked_loss = tf.reduce_sum(masked_loss) / tf.reduce_sum(loss_mask)
        return tf.reshape(reduced_masked_loss, (1,))
//  [python]
def hf_compute_loss(self, labels, logits):
        loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction=keras.losses.Reduction.NONE)
        if self.config.tf_legacy_loss:
            # make sure only labels that are not equal to -100
            # are taken into account as loss
            next_sentence_active_loss = tf.not_equal(tf.reshape(labels, (-1,)), -100)
            next_sentence_reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, 2)), next_sentence_active_loss)
            next_sentence_label = tf.boolean_mask(tf.reshape(labels, (-1,)), next_sentence_active_loss)

            return loss_fn(next_sentence_label, next_sentence_reduced_logits)

        # make sure only labels that are not equal to -100
        # are taken into account as loss

        # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway
        unmasked_ns_loss = loss_fn(y_true=tf.nn.relu(labels), y_pred=logits)
        ns_loss_mask = tf.to_float(tf.cast(labels != -100, tf.bool))
        # Just zero out samples where label is -100, no reduction
        masked_ns_loss = unmasked_ns_loss * ns_loss_mask

        return masked_ns_loss
//  [python]
def hf_compute_loss(self, labels, logits):
        loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction=keras.losses.Reduction.NONE)
        if self.config.tf_legacy_loss:
            # make sure only labels that are not equal to -100
            # are taken into account as loss
            next_sentence_active_loss = tf.not_equal(tf.reshape(labels, (-1,)), -100)
            next_sentence_reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, 2)), next_sentence_active_loss)
            next_sentence_label = tf.boolean_mask(tf.reshape(labels, (-1,)), next_sentence_active_loss)

            return loss_fn(next_sentence_label, next_sentence_reduced_logits)

        # make sure only labels that are not equal to -100
        # are taken into account as loss

        # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway
        unmasked_ns_loss = loss_fn(y_true=tf.nn.relu(labels), y_pred=logits)
        ns_loss_mask = tf.cast(labels != -100, dtype=unmasked_ns_loss.dtype)
        # Just zero out samples where label is -100, no reduction
        masked_ns_loss = unmasked_ns_loss * ns_loss_mask

        return masked_ns_loss
//  [python]
def _compute_loss(self, prediction_tensor, target_tensor, weights):
        """Compute loss function.

        Args:
          prediction_tensor: A float tensor of shape [batch_size, num_anchors,
            num_classes] representing the predicted logits for each class
          target_tensor: A float tensor of shape [batch_size, num_anchors,
            num_classes] representing one-hot encoded classification targets
          weights: a float tensor of shape [batch_size, num_anchors]

        Returns:
          loss: a float tensor of shape [batch_size, num_anchors, num_classes]
            representing the value of the loss function.
        """
        if self._bootstrap_type == 'soft':
            bootstrap_target_tensor = self._alpha * target_tensor + (
                1.0 - self._alpha) * tf.sigmoid(prediction_tensor)
        else:
            bootstrap_target_tensor = self._alpha * target_tensor + (
                1.0 - self._alpha) * tf.to_float(
                    tf.sigmoid(prediction_tensor) > 0.5)
        per_entry_cross_ent = (tf.nn.sigmoid_cross_entropy_with_logits(
            labels=bootstrap_target_tensor, logits=prediction_tensor))
        return per_entry_cross_ent * tf.expand_dims(weights, 2)
//  [python]
def _compute_loss(self, prediction_tensor, target_tensor, weights):
        """Compute loss function.

        Args:
          prediction_tensor: A float tensor of shape [batch_size, num_anchors,
            num_classes] representing the predicted logits for each class
          target_tensor: A float tensor of shape [batch_size, num_anchors,
            num_classes] representing one-hot encoded classification targets
          weights: a float tensor of shape [batch_size, num_anchors]

        Returns:
          loss: a float tensor of shape [batch_size, num_anchors, num_classes]
            representing the value of the loss function.
        """
        if self._bootstrap_type == 'soft':
            bootstrap_target_tensor = self._alpha * target_tensor + (
                1.0 - self._alpha) * tf.sigmoid(prediction_tensor)
        else:
            bootstrap_target_tensor = self._alpha * target_tensor + (
                1.0 - self._alpha) * tf.cast(
                    tf.sigmoid(prediction_tensor) > 0.5, tf.float32)
        per_entry_cross_ent = (tf.nn.sigmoid_cross_entropy_with_logits(
            labels=bootstrap_target_tensor, logits=prediction_tensor))
        return per_entry_cross_ent * tf.expand_dims(weights, 2)
//  [python]
def hf_compute_loss(self, labels: tf.Tensor, logits: tf.Tensor) -> tf.Tensor:
        loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(
            from_logits=True, reduction=tf.keras.losses.Reduction.NONE
        )

        # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway
        unmasked_lm_losses = loss_fn(y_true=tf.nn.relu(labels["labels"]), y_pred=logits[0])
        # make sure only labels that are not equal to -100
        # are taken into account for the loss computation
        lm_loss_mask = tf.to_float(labels["labels"] != -100)
        masked_lm_losses = unmasked_lm_losses * lm_loss_mask
        reduced_masked_lm_loss = tf.reduce_sum(masked_lm_losses) / tf.reduce_sum(lm_loss_mask)

        # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway
        unmasked_ns_loss = loss_fn(y_true=tf.nn.relu(labels["next_sentence_label"]), y_pred=logits[1])
        ns_loss_mask = tf.to_float(labels["next_sentence_label"] != -100)
        masked_ns_loss = unmasked_ns_loss * ns_loss_mask

        reduced_masked_ns_loss = tf.reduce_sum(masked_ns_loss) / tf.reduce_sum(ns_loss_mask)

        return tf.reshape(reduced_masked_lm_loss + reduced_masked_ns_loss, (1,))
//  [python]
def hf_compute_loss(self, labels: tf.Tensor, logits: tf.Tensor) -> tf.Tensor:
        loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(
            from_logits=True, reduction=tf.keras.losses.Reduction.NONE
        )

        # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway
        unmasked_lm_losses = loss_fn(y_true=tf.nn.relu(labels["labels"]), y_pred=logits[0])
        # make sure only labels that are not equal to -100
        # are taken into account for the loss computation
        lm_loss_mask = tf.cast(labels["labels"] != -100, dtype=unmasked_lm_losses.dtype)
        masked_lm_losses = unmasked_lm_losses * lm_loss_mask
        reduced_masked_lm_loss = tf.reduce_sum(masked_lm_losses) / tf.reduce_sum(lm_loss_mask)

        # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway
        unmasked_ns_loss = loss_fn(y_true=tf.nn.relu(labels["next_sentence_label"]), y_pred=logits[1])
        ns_loss_mask = tf.cast(labels["next_sentence_label"] != -100, dtype=unmasked_ns_loss.dtype)
        masked_ns_loss = unmasked_ns_loss * ns_loss_mask

        reduced_masked_ns_loss = tf.reduce_sum(masked_ns_loss) / tf.reduce_sum(ns_loss_mask)

        return tf.reshape(reduced_masked_lm_loss + reduced_masked_ns_loss, (1,))
//  [python]
def cross_replica_mean(t, num_shards_per_group=None):
  """Calculates the average value of input tensor across TPU replicas."""
  num_shards = tpu_function.get_tpu_context().number_of_shards
  if not num_shards_per_group:
    return tf.compat.v1.tpu.cross_replica_sum(t) / tf.cast(num_shards, t.dtype)

  group_assignment = None
  if num_shards_per_group > 1:
    if num_shards % num_shards_per_group != 0:
      raise ValueError('num_shards: %d mod shards_per_group: %d, should be 0' %
                       (num_shards, num_shards_per_group))
    num_groups = num_shards // num_shards_per_group
    group_assignment = [[
        x for x in range(num_shards) if x // num_shards_per_group == y
    ] for y in range(num_groups)]
  return tf.compat.v1.tpu.cross_replica_sum(t, group_assignment) / t.dtype
//  [python]
def cross_replica_mean(t, num_shards_per_group=None):
  """Calculates the average value of input tensor across TPU replicas."""
  num_shards = tpu_function.get_tpu_context().number_of_shards
  if not num_shards_per_group:
    return tf.compat.v1.tpu.cross_replica_sum(t) / tf.cast(num_shards, t.dtype)

  group_assignment = None
  if num_shards_per_group > 1:
    if num_shards % num_shards_per_group != 0:
      raise ValueError('num_shards: %d mod shards_per_group: %d, should be 0' %
                       (num_shards, num_shards_per_group))
    num_groups = num_shards // num_shards_per_group
    group_assignment = [[
        x for x in range(num_shards) if x // num_shards_per_group == y
    ] for y in range(num_groups)]
  return tf.compat.v1.tpu.cross_replica_sum(t, group_assignment) / tf.cast(
      num_shards_per_group, t.dtype)
//  [python]
def get_img_points(sample_pc, trans_mat_right, pred_xyshift, FLAGS):
    # sample_pc B*N*3
    size_lst = sample_pc.get_shape().as_list()
    homo_pc = tf.concat((sample_pc, tf.ones((size_lst[0], size_lst[1], 1),dtype=np.float32)),axis= -1)
    print("homo_pc.get_shape()", homo_pc.get_shape())
    pc_xyz = tf.matmul(homo_pc, trans_mat_right)
    print("pc_xyz.get_shape()", pc_xyz.get_shape()) # B * N * 3
    pc_xy = tf.div(pc_xyz[:,:,:2], tf.expand_dims(pc_xyz[:,:,2], axis = 2))
    if FLAGS.shift:
        pc_xy = pc_xy + tf.tile(tf.expand_dims(pred_xyshift / 2 * FLAGS.img_h, axis=1), (1,FLAGS.num_points,1))
    mintensor = tf.constant([0.0,0.0], dtype=tf.float32)
    maxtensor = tf.constant([136.0,136.0], dtype=tf.float32)
    return tf.minimum(maxtensor, tf.maximum(mintensor, pc_xy)), pc_xy
//  [python]
def get_img_points(sample_pc, trans_mat_right, pred_xyshift, FLAGS):
    # sample_pc B*N*3
    size_lst = sample_pc.get_shape().as_list()
    homo_pc = tf.concat((sample_pc, tf.ones((size_lst[0], size_lst[1], 1),dtype=np.float32)),axis= -1)
    print("homo_pc.get_shape()", homo_pc.get_shape())
    pc_xyz = tf.matmul(homo_pc, trans_mat_right)
    print("pc_xyz.get_shape()", pc_xyz.get_shape()) # B * N * 3
    pc_xy = tf.cast(tf.divide(pc_xyz[:,:,:2], tf.expand_dims(pc_xyz[:,:,2], axis = 2)), dtype=tf.float32)
    if FLAGS.shift:
        pc_xy = pc_xy + tf.tile(tf.expand_dims(pred_xyshift / 2 * FLAGS.img_h, axis=1), (1,FLAGS.num_points,1))
    mintensor = tf.constant([0.0,0.0], dtype=tf.float32)
    maxtensor = tf.constant([136.0,136.0], dtype=tf.float32)
    return tf.minimum(maxtensor, tf.maximum(mintensor, pc_xy)), pc_xy
//  [python]
def __call__(self, step):
    if self.lr_decay_type == 'exponential':
      assert self.steps_per_epoch is not None
      decay_steps = self.steps_per_epoch * self.decay_epochs
      lr = tf.keras.optimizers.schedules.ExponentialDecay(
          self.initial_lr, decay_steps, self.decay_factor, staircase=True)(
              step)
    elif self.lr_decay_type == 'cosine':
      assert self.total_steps is not None
      lr = 0.5 * self.initial_lr * (
          1 + tf.cos(np.pi * t.float32(step) / self.total_steps))
    elif self.lr_decay_type == 'linear':
      assert self.total_steps is not None
      lr = (1.0 -
            t.float32(step) / self.total_steps) * self.initial_lr
    elif self.lr_decay_type == 'constant':
      lr = self.initial_lr
    else:
      assert False, 'Unknown lr_decay_type : %s' % self.lr_decay_type

    if self.minimal_lr:
      lr = tf.math.maximum(lr, self.minimal_lr)

    if self.warmup_epochs:
      warmup_steps = int(self.warmup_epochs * self.steps_per_epoch)
      warmup_lr = (
          self.initial_lr * t.float32(step) /
          t.float32(warmup_steps))
      lr = tf.cond(step < warmup_steps, lambda: warmup_lr, lambda: lr)

    return lr
//  [python]
def __call__(self, step):
    if self.lr_decay_type == 'exponential':
      assert self.steps_per_epoch is not None
      decay_steps = self.steps_per_epoch * self.decay_epochs
      lr = tf.keras.optimizers.schedules.ExponentialDecay(
          self.initial_lr, decay_steps, self.decay_factor, staircase=True)(
              step)
    elif self.lr_decay_type == 'cosine':
      assert self.total_steps is not None
      lr = 0.5 * self.initial_lr * (
          1 + tf.cos(np.pi * tf.cast(step, tf.float32) / self.total_steps))
    elif self.lr_decay_type == 'linear':
      assert self.total_steps is not None
      lr = (1.0 -
            tf.cast(step, tf.float32) / self.total_steps) * self.initial_lr
    elif self.lr_decay_type == 'constant':
      lr = self.initial_lr
    else:
      assert False, 'Unknown lr_decay_type : %s' % self.lr_decay_type

    if self.minimal_lr:
      lr = tf.math.maximum(lr, self.minimal_lr)

    if self.warmup_epochs:
      warmup_steps = int(self.warmup_epochs * self.steps_per_epoch)
      warmup_lr = (
          self.initial_lr * tf.cast(step, tf.float32) /
          tf.cast(warmup_steps, tf.float32))
      lr = tf.cond(step < warmup_steps, lambda: warmup_lr, lambda: lr)

    return lr
//  [python]
def predict_setup(self):
		# Create queue coordinator.
		self.coord = tf.train.Coordinator()

		# Load reader
		with tf.name_scope("create_inputs"):
			reader = ImageReader(
				self.conf.data_dir,
				self.conf.test_data_list,
				None, # the images have different sizes
				False, # no data-aug
				False, # no data-aug
				self.conf.ignore_label,
				IMG_MEAN,
				self.coord)
			image, label = reader.image, reader.label # [h, w, 3 or 1]
		# Add one batch dimension [1, h, w, 3 or 1]
		image_batch, label_batch = tf.expand_dims(image, dim=0), tf.expand_dims(label, dim=0)

		# Create network
		if self.conf.encoder_name not in ['res101', 'res50']:
			print('encoder_name ERROR!')
			print("Please input: res101, res50")
			sys.exit(-1)
		else:
			net = ResNet_segmentation(image_batch, self.conf.num_classes, False, self.conf.encoder_name)

		# Predictions.
		raw_output = net.outputs
		raw_output = tf.image.resize_bilinear(raw_output, tf.shape(image_batch)[1:3,])
		raw_output = tf.argmax(raw_output, axis=3)
		self.pred = tf.expand_dims(raw_output, dim=3)

		# Create directory
		if not os.path.exists(self.conf.out_dir):
			os.makedirs(self.conf.out_dir)
			os.makedirs(self.conf.out_dir + '/prediction')
			if self.conf.visual:
				os.makedirs(self.conf.out_dir + '/visual_prediction')

		# Loader for loading the checkpoint
		self.loader = tf.train.Saver(var_list=tf.global_variables())
//  [python]
def predict_setup(self):
		# Create queue coordinator.
		self.coord = tf.train.Coordinator()

		# Load reader
		with tf.name_scope("create_inputs"):
			reader = ImageReader(
				self.conf.data_dir,
				self.conf.test_data_list,
				None, # the images have different sizes
				False, # no data-aug
				False, # no data-aug
				self.conf.ignore_label,
				IMG_MEAN,
				self.coord)
			image, label = reader.image, reader.label # [h, w, 3 or 1]
		# Add one batch dimension [1, h, w, 3 or 1]
		image_batch, label_batch = tf.expand_dims(image, dim=0), tf.expand_dims(label, dim=0)

		# Create network
		if self.conf.encoder_name not in ['res101', 'res50']:
			print('encoder_name ERROR!')
			print("Please input: res101, res50")
			sys.exit(-1)
		else:
			net = ResNet_segmentation(image_batch, self.conf.num_classes, False, self.conf.encoder_name)

		# Predictions.
		raw_output = net.outputs
		raw_output = tf.image.resize_bilinear(raw_output, tf.shape(image_batch)[1:3,])
		raw_output = tf.argmax(raw_output, axis=3)
		self.pred = tf.cast(tf.expand_dims(raw_output, dim=3), tf.uint8)

		# Create directory
		if not os.path.exists(self.conf.out_dir):
			os.makedirs(self.conf.out_dir)
			os.makedirs(self.conf.out_dir + '/prediction')
			if self.conf.visual:
				os.makedirs(self.conf.out_dir + '/visual_prediction')

		# Loader for loading the checkpoint
		self.loader = tf.train.Saver(var_list=tf.global_variables())
//  [python]
def build_model_with_precision(pp, mm, ii, tt, *args, **kwargs):
  """Build model with its inputs/params for a specified precision context.

  This is highly specific to this codebase, and not intended to be general API.
  Advanced users only. DO NOT use it if you don't know what it does.
  NOTE: short argument names are intended to avoid conficts with kwargs.

  Args:
    pp: A string, precision policy name, such as "mixed_float16".
    mm: A function, for rmodel builder.
    ii: A tensor, for model inputs.
    tt: A bool, If true, it is for training; otherwise, it is for eval.
    *args: A list of model arguments.
    **kwargs: A dict, extra model parameters.

  Returns:
    the output of mm model.
  """
  if pp == 'mixed_bfloat16':
    set_precision_policy(pp)
    inputs = t.cast(ii, tf.bfloat16)
    with tf.compat.v1.tpu.bfloat16_scope():
      outputs = mm(inputs, *args, **kwargs)
    set_precision_policy('float32')
  elif pp == 'mixed_float16':
    set_precision_policy(pp, loss_scale=tt)
    inputs = t.cast(ii, tf.float16)
    with float16_scope():
      outputs = mm(inputs, *args, **kwargs)
    set_precision_policy('float32')
  elif not pp or pp == 'float32':
    outputs = mm(ii, *args, **kwargs)
  else:
    raise ValueError('Unknow precision name {}'.format(pp))

  # Users are responsible to convert the dtype of all outputs.
  return outputs
//  [python]
def build_model_with_precision(pp, mm, ii, tt, *args, **kwargs):
  """Build model with its inputs/params for a specified precision context.

  This is highly specific to this codebase, and not intended to be general API.
  Advanced users only. DO NOT use it if you don't know what it does.
  NOTE: short argument names are intended to avoid conficts with kwargs.

  Args:
    pp: A string, precision policy name, such as "mixed_float16".
    mm: A function, for rmodel builder.
    ii: A tensor, for model inputs.
    tt: A bool, If true, it is for training; otherwise, it is for eval.
    *args: A list of model arguments.
    **kwargs: A dict, extra model parameters.

  Returns:
    the output of mm model.
  """
  if pp == 'mixed_bfloat16':
    set_precision_policy(pp)
    inputs = tf.cast(ii, tf.bfloat16)
    with tf.compat.v1.tpu.bfloat16_scope():
      outputs = mm(inputs, *args, **kwargs)
    set_precision_policy('float32')
  elif pp == 'mixed_float16':
    set_precision_policy(pp, loss_scale=tt)
    inputs = tf.cast(ii, tf.float16)
    with float16_scope():
      outputs = mm(inputs, *args, **kwargs)
    set_precision_policy('float32')
  elif not pp or pp == 'float32':
    outputs = mm(ii, *args, **kwargs)
  else:
    raise ValueError('Unknow precision name {}'.format(pp))

  # Users are responsible to convert the dtype of all outputs.
  return outputs
//  [python]
def is_current_phase(self, x):
        unit = 54.5
        min_ts = t.cast(x.phase, tf.float32) * unit
        max_ts = min_ts + (unit * 4) + 1
        return tf.logical_and(tf.greater_equal(x.q_ts, min_ts), tf.less_equal(x.q_ts, max_ts))
//  [python]
def is_current_phase(self, x):
        unit = 54.5
        min_ts = tf.cast(x.phase, tf.float32) * unit
        max_ts = min_ts + (unit * 4) + 1
        return tf.logical_and(tf.greater_equal(x.q_ts, min_ts), tf.less_equal(x.q_ts, max_ts))
//  [python]
def non_linear_feat_emb(self, feat):
        # [N,]
        feat = t.cast(feat, tf.float32)
        f1 = self.float201(feat)
        f2 = self.float201(tf.log(feat))
        f3 = self.float201(tf.sqrt(feat))
        f = tf.stack([f1, f2, f3], axis=-1)
        return f
//  [python]
def non_linear_feat_emb(self, feat):
        # [N,]
        feat = tf.cast(feat, tf.float32)
        f1 = self.float201(feat)
        f2 = self.float201(tf.log(feat))
        f3 = self.float201(tf.sqrt(feat))
        f = tf.stack([f1, f2, f3], axis=-1)
        return f
//  [python]
def update_state(self, y_true, y_pred, sample_weight=None) -> None:
        if not hasattr(self, "squared_sum"):
            self.squared_sum = self.add_weight(
                name="squared_sum",
                shape=y_true.shape[1:],
                initializer="zeros",
                dtype=self._dtype,
            )
        if not hasattr(self, "sum"):
            self.sum = self.add_weight(
                name="sum",
                shape=y_true.shape[1:],
                initializer="zeros",
                dtype=self._dtype,
            )
        if not hasattr(self, "res"):
            self.res = self.add_weight(
                name="residual",
                shape=y_true.shape[1:],
                initializer="zeros",
                dtype=self._dtype,
            )
        if not hasattr(self, "count"):
            self.count = self.add_weight(
                name="count",
                shape=y_true.shape[1:],
                initializer="zeros",
                dtype=self._dtype,
            )

        y_true = t.cast(y_true, dtype=self._dtype)
        y_pred = t.cast(y_pred, dtype=self._dtype)
        if sample_weight is None:
            sample_weight = 1
        sample_weight = t.cast(sample_weight, dtype=self._dtype)
        sample_weight = weights_broadcast_ops.broadcast_weights(
            weights=sample_weight, values=y_true
        )

        weighted_y_true = y_true * sample_weight
        self.sum.assign_add(tf.reduce_sum(weighted_y_true, axis=0))
        self.squared_sum.assign_add(tf.reduce_sum(y_true * weighted_y_true, axis=0))
        self.res.assign_add(
            tf.reduce_sum((y_true - y_pred) ** 2 * sample_weight, axis=0)
        )
        self.count.assign_add(tf.reduce_sum(sample_weight, axis=0))
        self.num_samples.assign_add(tf.size(y_true))
//  [python]
def update_state(self, y_true, y_pred, sample_weight=None) -> None:
        if not hasattr(self, "squared_sum"):
            self.squared_sum = self.add_weight(
                name="squared_sum",
                shape=y_true.shape[1:],
                initializer="zeros",
                dtype=self._dtype,
            )
        if not hasattr(self, "sum"):
            self.sum = self.add_weight(
                name="sum",
                shape=y_true.shape[1:],
                initializer="zeros",
                dtype=self._dtype,
            )
        if not hasattr(self, "res"):
            self.res = self.add_weight(
                name="residual",
                shape=y_true.shape[1:],
                initializer="zeros",
                dtype=self._dtype,
            )
        if not hasattr(self, "count"):
            self.count = self.add_weight(
                name="count",
                shape=y_true.shape[1:],
                initializer="zeros",
                dtype=self._dtype,
            )

        y_true = tf.cast(y_true, dtype=self._dtype)
        y_pred = tf.cast(y_pred, dtype=self._dtype)
        if sample_weight is None:
            sample_weight = 1
        sample_weight = tf.cast(sample_weight, dtype=self._dtype)
        sample_weight = weights_broadcast_ops.broadcast_weights(
            weights=sample_weight, values=y_true
        )

        weighted_y_true = y_true * sample_weight
        self.sum.assign_add(tf.reduce_sum(weighted_y_true, axis=0))
        self.squared_sum.assign_add(tf.reduce_sum(y_true * weighted_y_true, axis=0))
        self.res.assign_add(
            tf.reduce_sum((y_true - y_pred) ** 2 * sample_weight, axis=0)
        )
        self.count.assign_add(tf.reduce_sum(sample_weight, axis=0))
        self.num_samples.assign_add(tf.size(y_true))
//  [python]
def is_current_day(self, x):
        unit = 54.5
        day = t.cast(x.phase + 3, t.float32)
        min_ts = 10.0 + day * unit
        max_ts = min_ts + unit + 1
        return tf.logical_and(tf.greater_equal(x.q_ts, min_ts), tf.less_equal(x.q_ts, max_ts))
//  [python]
def is_current_day(self, x):
        unit = 54.5
        day = tf.cast(x.phase + 3, tf.float32)
        min_ts = 10.0 + day * unit
        max_ts = min_ts + unit + 1
        return tf.logical_and(tf.greater_equal(x.q_ts, min_ts), tf.less_equal(x.q_ts, max_ts))
//  [python]
def result(self) -> tf.Tensor:
        mean = self.sum / self.count
        total = self.squared_sum - self.sum * mean
        raw_scores = 1 - (self.res / total)
        raw_scores = tf.where(tf.math.is_inf(raw_scores), 0.0, raw_scores)

        if self.multioutput == "raw_values":
            r2_score = raw_scores
        elif self.multioutput == "uniform_average":
            r2_score = tf.reduce_mean(raw_scores)
        elif self.multioutput == "variance_weighted":
            r2_score = _reduce_average(raw_scores, weights=total)
        else:
            raise RuntimeError(
                "The multioutput attribute must be one of {}, but was: {}".format(
                    _VALID_MULTIOUTPUT, self.multioutput
                )
            )

        if self.num_regressors < 0:
            raise ValueError(
                "num_regressors parameter should be greater than or equal to zero"
            )

        if self.num_regressors != 0:
            if self.num_regressors > self.num_samples - 1:
                UserWarning(
                    "More independent predictors than datapoints in adjusted r2 score. Falls back to standard r2 "
                    "score."
                )
            elif self.num_regressors == self.num_samples - 1:
                UserWarning(
                    "Division by zero in adjusted r2 score. Falls back to standard r2 score."
                )
            else:
                n = tf.to_float(self.num_samples)
                p = tf.to_float(self.num_regressors)

                num = tf.multiply(tf.subtract(1.0, r2_score), tf.subtract(n, 1.0))
                den = tf.subtract(tf.subtract(n, p), 1.0)
                r2_score = tf.subtract(1.0, tf.divide(num, den))

        return r2_score
//  [python]
def result(self) -> tf.Tensor:
        mean = self.sum / self.count
        total = self.squared_sum - self.sum * mean
        raw_scores = 1 - (self.res / total)
        raw_scores = tf.where(tf.math.is_inf(raw_scores), 0.0, raw_scores)

        if self.multioutput == "raw_values":
            r2_score = raw_scores
        elif self.multioutput == "uniform_average":
            r2_score = tf.reduce_mean(raw_scores)
        elif self.multioutput == "variance_weighted":
            r2_score = _reduce_average(raw_scores, weights=total)
        else:
            raise RuntimeError(
                "The multioutput attribute must be one of {}, but was: {}".format(
                    _VALID_MULTIOUTPUT, self.multioutput
                )
            )

        if self.num_regressors < 0:
            raise ValueError(
                "num_regressors parameter should be greater than or equal to zero"
            )

        if self.num_regressors != 0:
            if self.num_regressors > self.num_samples - 1:
                UserWarning(
                    "More independent predictors than datapoints in adjusted r2 score. Falls back to standard r2 "
                    "score."
                )
            elif self.num_regressors == self.num_samples - 1:
                UserWarning(
                    "Division by zero in adjusted r2 score. Falls back to standard r2 score."
                )
            else:
                n = tf.cast(self.num_samples, dtype=tf.float32)
                p = tf.cast(self.num_regressors, dtype=tf.float32)

                num = tf.multiply(tf.subtract(1.0, r2_score), tf.subtract(n, 1.0))
                den = tf.subtract(tf.subtract(n, p), 1.0)
                r2_score = tf.subtract(1.0, tf.divide(num, den))

        return r2_score
//  [python]
def get_loss(self, labels: tf.Tensor, logits: tf.Tensor, **kwargs: Any) -> tf.Tensor:
    with tf.name_scope("weights"):
      disliked_mask = generate_disliked_mask(labels)

      labels = tf.reshape(labels[:, 0:2], shape=[-1, 2])

      labels = labels * tf.logical_not(disliked_mask)

      with tf.name_scope("task_weight"):
        task_weights = self.get_task_weights(labels)

      with tf.name_scope("batch_size"):
        batch_size = tf.cast(tf.shape(labels)[0], dtype=tf.float32, name="batch_size")

      weights = task_weights / batch_size

    with tf.name_scope("loss"):
      loss = tf.reduce_sum(
        tf.nn.sigmoid_cross_entropy_with_logits(labels=labels, logits=logits) * weights,
      )

    return loss
//  [python]
def get_loss(self, labels: tf.Tensor, logits: tf.Tensor, **kwargs: Any) -> tf.Tensor:
    with tf.name_scope("weights"):
      disliked_mask = generate_disliked_mask(labels)

      labels = tf.reshape(labels[:, 0:2], shape=[-1, 2])

      labels = labels * tf.cast(tf.logical_not(disliked_mask), dtype=labels.dtype)

      with tf.name_scope("task_weight"):
        task_weights = self.get_task_weights(labels)

      with tf.name_scope("batch_size"):
        batch_size = tf.cast(tf.shape(labels)[0], dtype=tf.float32, name="batch_size")

      weights = task_weights / batch_size

    with tf.name_scope("loss"):
      loss = tf.reduce_sum(
        tf.nn.sigmoid_cross_entropy_with_logits(labels=labels, logits=logits) * weights,
      )

    return loss
//  [python]
def get_positive_weights():
  """Computes positive weights used for class imbalance from training data."""
  label_weights_df = utils.get_label_weights(
      "tos-data-media-full",
      project_id="twttr-abusive-interact-prod",
      dataset_id="tos_policy"
  )
  pos_weight_tensor = t.cast(
      label_weights_df.sort_values(by='label').positive_class_weight,
      dtype=t.float32
  )
  return pos_weight_tensor
//  [python]
def get_positive_weights():
  """Computes positive weights used for class imbalance from training data."""
  label_weights_df = utils.get_label_weights(
      "tos-data-media-full",
      project_id="twttr-abusive-interact-prod",
      dataset_id="tos_policy"
  )
  pos_weight_tensor = tf.cast(
      label_weights_df.sort_values(by='label').positive_class_weight,
      dtype=tf.float32
  )
  return pos_weight_tensor
//  [python]
def serve(self, images):
    """Cast image to float and run inference.

    Args:
      images: uint8 Tensor of shape [batch_size, None, None, 3]
    Returns:
      Tensor holding classification output logits.
    """
    with tf.device('cpu:0'):
      images = tf.to_float(images)

      images = tf.nest.map_structure(
          tf.identity,
          tf.map_fn(
              self._build_inputs, elems=images,
              fn_output_signature=tf.TensorSpec(
                  shape=self._input_image_size + [3], dtype=tf.float32),
              parallel_iterations=32
              )
          )

    masks = self.inference_step(images)
    keys = sorted(masks.keys())
    output = tf.image.resize(
        masks[keys[-1]],
        self._input_image_size, method='bilinear')

    return dict(predicted_masks=output)
//  [python]
def serve(self, images):
    """Cast image to float and run inference.

    Args:
      images: uint8 Tensor of shape [batch_size, None, None, 3]
    Returns:
      Tensor holding classification output logits.
    """
    with tf.device('cpu:0'):
      images = tf.cast(images, dtype=tf.float32)

      images = tf.nest.map_structure(
          tf.identity,
          tf.map_fn(
              self._build_inputs, elems=images,
              fn_output_signature=tf.TensorSpec(
                  shape=self._input_image_size + [3], dtype=tf.float32),
              parallel_iterations=32
              )
          )

    masks = self.inference_step(images)
    keys = sorted(masks.keys())
    output = tf.image.resize(
        masks[keys[-1]],
        self._input_image_size, method='bilinear')

    return dict(predicted_masks=output)
//  [python]
def positional_encoding(self):
        '''
        calculate the positional encoding given the window length
        :return: positional encoding (1, window_length, 1, d_model)
        '''
        angle_rads = self.get_angles(np.arange(self.sequence_length)[:, np.newaxis], np.arange(self.d_model)[np.newaxis, :])

        # apply sin to even indices in the array; 2i
        angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])

        # apply cos to odd indices in the array; 2i+1
        angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])

        pos_encoding = angle_rads[np.newaxis, :, np.newaxis, :]

        return tf.cast(pos_encoding, dtype=tf.float32)
//  [python]
def sample_neg(self, ans, pred_score):
        BS = tf.shape(ans)[0]
        def sample(minval, maxval):
            _neg = tf.random_uniform(
                shape=(BS, args.nb_neg),
                minval=minval,
                maxval=maxval,
                dtype=tf.int32,
                name='sample_neg',
            )
            # DEBUG.print_when_fit(_neg, block=True)
            return _neg

        if args.mode_neg_sample == 'all':
            neg = sample(3, args.nb_items)
        else:
            if args.mode_neg_sample == 'all2':
                base_score = 0.0
            elif args.mode_neg_sample == 'pred':
                base_score = t.cast(self.data.dp.score_mask, tf.float32)
            elif args.mode_neg_sample == 'rare':
                base_score = t.cast(self.data.is_rare, tf.float32)
            elif args.mode_neg_sample == 'not_rare':
                base_score = t.cast(tf.logical_not(self.data.is_rare), tf.float32)
            elif args.mode_neg_sample == 'pop':
                base_score = t.cast(self.data.item_pop, tf.float32)
                base_score = tf.log(base_score + 1.0)
            elif args.mode_neg_sample == 'deg':
                base_score = t.cast(self.data.item_deg, tf.float32)
                base_score = tf.log(base_score + 1.0)
            elif args.mode_neg_sample == 'deg_self':
                base_score = t.cast(self.data.item_deg_self, tf.float32)
                base_score = tf.log(base_score + 1.0)
            elif args.mode_neg_sample == 'pred_top':
                base_score = pred_score
                base_score = tf.nn.sigmoid(base_score)
            else:
                raise Exception
            random_score = tf.random_uniform(shape=tf.shape(pred_score), minval=0.0, maxval=args.alpha_neg_sample, dtype=tf.float32)
            neg = tf.nn.top_k(base_score + random_score, args.nb_neg).indices
        pos = tf.expand_dims(ans, -1)
        # neg: [BS, nb_neg]
        # pos: [BS, 1]
        neg = tf.where(tf.equal(neg, pos), tf.zeros_like(neg, dtype=tf.int32), neg)
        return neg
//  [python]
def sample_neg(self, ans, pred_score):
        BS = tf.shape(ans)[0]
        def sample(minval, maxval):
            _neg = tf.random_uniform(
                shape=(BS, args.nb_neg),
                minval=minval,
                maxval=maxval,
                dtype=tf.int32,
                name='sample_neg',
            )
            # DEBUG.print_when_fit(_neg, block=True)
            return _neg

        if args.mode_neg_sample == 'all':
            neg = sample(3, args.nb_items)
        else:
            if args.mode_neg_sample == 'all2':
                base_score = 0.0
            elif args.mode_neg_sample == 'pred':
                base_score = tf.cast(self.data.dp.score_mask, tf.float32)
            elif args.mode_neg_sample == 'rare':
                base_score = tf.cast(self.data.is_rare, tf.float32)
            elif args.mode_neg_sample == 'not_rare':
                base_score = tf.cast(tf.logical_not(self.data.is_rare), tf.float32)
            elif args.mode_neg_sample == 'pop':
                base_score = tf.cast(self.data.item_pop, tf.float32)
                base_score = tf.log(base_score + 1.0)
            elif args.mode_neg_sample == 'deg':
                base_score = tf.cast(self.data.item_deg, tf.float32)
                base_score = tf.log(base_score + 1.0)
            elif args.mode_neg_sample == 'deg_self':
                base_score = tf.cast(self.data.item_deg_self, tf.float32)
                base_score = tf.log(base_score + 1.0)
            elif args.mode_neg_sample == 'pred_top':
                base_score = pred_score
                base_score = tf.nn.sigmoid(base_score)
            else:
                raise Exception
            random_score = tf.random_uniform(shape=tf.shape(pred_score), minval=0.0, maxval=args.alpha_neg_sample, dtype=tf.float32)
            neg = tf.nn.top_k(base_score + random_score, args.nb_neg).indices
        pos = tf.expand_dims(ans, -1)
        # neg: [BS, nb_neg]
        # pos: [BS, 1]
        neg = tf.where(tf.equal(neg, pos), tf.zeros_like(neg, dtype=tf.int32), neg)
        return neg
//  [python]
def masked_bce(y_true, y_pred):
  y_true = t.cast(y_true, dtype=tf.float32)
  mask = y_true != -1
  
  return tf.keras.metrics.binary_crossentropy(tf.boolean_mask(y_true, mask), 
                                              tf.boolean_mask(y_pred, mask))
//  [python]
def masked_bce(y_true, y_pred):
  y_true = tf.cast(y_true, dtype=tf.float32)
  mask = y_true != -1
  
  return tf.keras.metrics.binary_crossentropy(tf.boolean_mask(y_true, mask), 
                                              tf.boolean_mask(y_pred, mask))
//  [python]
def create_look_ahead_mask(self):
        '''
        create a look ahead mask given a certain window length
        :return: the mask (window_length, window_length)
        '''
        ahead_mask = 1 - tf.linalg.band_part(tf.ones((self.sequence_length, self.sequence_length)), -1, 0)

        if self.random_window_min > 0 and self.is_training:
            random_win_len = tf.random.uniform([self.sequence_length], self.random_window_min, self.window_len, dtype=tf.int32)
            window_mask_padding = tf.maximum(0, tf.range(1, self.sequence_length + 1) - random_win_len)
        else:
            window_mask_padding = tf.maximum(0, tf.range(1, self.sequence_length+1) - self.window_len)
        window_mask = tf.sequence_mask(window_mask_padding, dtype=tf.float32, maxlen=self.sequence_length)

        mask = tf.maximum(ahead_mask, window_mask)
        
        if self.temporal_mask_drop > 0 and self.is_training:
            drop_mask = tf.to_float(tf.random.uniform((self.sequence_length, self.sequence_length), 0, 1) < self.temporal_mask_drop)
            mask_ = tf.maximum(mask, drop_mask)
            # Ensure that the first timestep is not masked.
            mask = tf.concat([mask[:, 0:1], mask_[:, 1:]], axis=1)
            
        return mask
//  [python]
def create_look_ahead_mask(self):
        '''
        create a look ahead mask given a certain window length
        :return: the mask (window_length, window_length)
        '''
        ahead_mask = 1 - tf.linalg.band_part(tf.ones((self.sequence_length, self.sequence_length)), -1, 0)

        if self.random_window_min > 0 and self.is_training:
            random_win_len = tf.random.uniform([self.sequence_length], self.random_window_min, self.window_len, dtype=tf.int32)
            window_mask_padding = tf.maximum(0, tf.range(1, self.sequence_length + 1) - random_win_len)
        else:
            window_mask_padding = tf.maximum(0, tf.range(1, self.sequence_length+1) - self.window_len)
        window_mask = tf.sequence_mask(window_mask_padding, dtype=tf.float32, maxlen=self.sequence_length)

        mask = tf.maximum(ahead_mask, window_mask)
        
        if self.temporal_mask_drop > 0 and self.is_training:
            drop_mask = tf.cast(tf.random.uniform((self.sequence_length, self.sequence_length), 0, 1) < self.temporal_mask_drop, tf.float32)
            mask_ = tf.maximum(mask, drop_mask)
            # Ensure that the first timestep is not masked.
            mask = tf.concat([mask[:, 0:1], mask_[:, 1:]], axis=1)
            
        return mask
//  [python]
def inv_kl_divergence(y_true, y_pred):
  y_pred = tf.convert_to_tensor(y_pred)
  y_true = tf.cast(y_true, y_pred.dtype)
  y_true = backend.clip(y_true, backend.epsilon(), 1)
  y_pred = backend.clip(y_pred, backend.epsilon(), 1)
  return tf.reduce_sum(y_pred * tf.math.log(y_pred / y_true), axis=-1)
//  [python]
def __init__(self, yh, y, oy, config):
    self.config = config
    self.yh = yh
    self.y = y
    self.oy = oy  # original training labels (when labels are permuted)
    self.loss = tf.reduce_mean(
        tf.nn.softmax_cross_entropy_with_logits_v2(logits=yh, labels=y))
    self.acc = tf.reduce_mean(
        tf.cast(tf.equal(tf.argmax(yh, 1), tf.argmax(y, 1)), tf.float32))

    inf = np.iinfo(np.int32).max
    self.learnt_first = np.full_like(oy, inf, dtype=np.int32)
    self.learnt_last = np.full_like(oy, inf, dtype=np.int32)
//  [python]
def call(self, inputs,mask=None):
        dis_entity,men_state,dis_entity_mark=inputs
        dis_entity_mark=t.cast(dis_entity_mark,tf.int32)
        men_state = tf.batch_gather(men_state, dis_entity_mark)
        out=concatenate([men_state, dis_entity])
        if mask is not None:
            mask=K.expand_dims(mask[0], 2)
            out *= K.cast(mask[0], K.floatx())
        return out