jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def init_optimizer(self):
        """
        Init self.optimizer.
        """
        self.learning_rate = tf.keras.optimizers.schedules.ExponentialDecay(self.learning_rate, self.max_iter, 0.1)
        self.optimizer = tf.keras.optimizers.Adam(learning_rate=self.learning_rate, amsgrad=False)
        if self.use_mixed_precision:
            self.optimizer = tf.keras.mixed_precision.LossScaleOptimizer(self.optimizer,
                                                                initial_scale=8, dynamic=True)
//  [python]
def get_optimizer(params):
  """Get optimizer."""
  learning_rate = learning_rate_schedule(params)
  momentum = params['momentum']
  if params['optimizer'].lower() == 'sgd':
    logging.info('Use SGD optimizer')
    optimizer = tf.keras.optimizers.SGD(learning_rate, momentum=momentum)
  elif params['optimizer'].lower() == 'adam':
    logging.info('Use Adam optimizer')
    optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=momentum)
  else:
    raise ValueError('optimizers should be adam or sgd')

  moving_average_decay = params['moving_average_decay']
  if moving_average_decay:
    # TODO(tanmingxing): potentially add dynamic_decay for new tfa release.
    from tensorflow_addons import optimizers as tfa_optimizers  # pylint: disable=g-import-not-at-top
    optimizer = tfa_optimizers.MovingAverage(
        optimizer, average_decay=moving_average_decay, dynamic_decay=True)
  precision = utils.get_precision(params['strategy'], params['mixed_precision'])
  if precision == 'mixed_float16' and params['loss_scale']:
    optimizer = tf.keras.mixed_precision.experimental.LossScaleOptimizer(
        optimizer,
        loss_scale=tf.mixed_precision.experimental.DynamicLossScale(
            params['loss_scale']))
  return optimizer
//  [python]
def get_optimizer(params):
  """Get optimizer."""
  learning_rate = learning_rate_schedule(params)
  momentum = params['momentum']
  if params['optimizer'].lower() == 'sgd':
    logging.info('Use SGD optimizer')
    optimizer = tf.keras.optimizers.SGD(learning_rate, momentum=momentum)
  elif params['optimizer'].lower() == 'adam':
    logging.info('Use Adam optimizer')
    optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=momentum)
  else:
    raise ValueError('optimizers should be adam or sgd')

  moving_average_decay = params['moving_average_decay']
  if moving_average_decay:
    # TODO(tanmingxing): potentially add dynamic_decay for new tfa release.
    from tensorflow_addons import optimizers as tfa_optimizers  # pylint: disable=g-import-not-at-top
    optimizer = tfa_optimizers.MovingAverage(
        optimizer, average_decay=moving_average_decay, dynamic_decay=True)
  precision = utils.get_precision(params['strategy'], params['mixed_precision'])
  if precision == 'mixed_float16' and params['loss_scale']:
    optimizer = tf.keras.mixed_precision.LossScaleOptimizer(
        optimizer,
        loss_scale=params['loss_scale'])
  return optimizer
//  [python]
def _decode(self, record: tf.Tensor):
    """Decodes a serialized tf.Example."""
    name_to_features = {
        'input_mask': tf.io.VarLenFeature(tf.int64),
        'masked_lm_positions': tf.io.VarLenFeature(tf.int64),
        'masked_lm_ids': tf.io.VarLenFeature(tf.int64),
        'masked_lm_weights': tf.io.VarLenFeature(tf.float32),
    }
    if self._params.use_v2_feature_names:
      input_ids_key = 'input_word_ids'
      segment_key = 'input_type_ids'
      name_to_features.update({
          input_ids_key: tf.io.VarLenFeature(tf.int64),
          segment_key: tf.io.VarLenFeature(tf.int64),
      })
    else:
      input_ids_key = 'input_ids'
      segment_key = 'segment_ids'
      name_to_features.update({
          input_ids_key: tf.io.VarLenFeature(tf.int64),
          segment_key: tf.io.VarLenFeature(tf.int64),
      })
    if self._use_next_sentence_label:
      name_to_features['next_sentence_labels'] = tf.io.FixedLenFeature([1],
                                                                       tf.int64)
    dynamic_keys = [input_ids_key, 'input_mask', segment_key]
    if self._use_position_id:
      name_to_features['position_ids'] = tf.io.VarLenFeature(tf.int64)
      dynamic_keys.append('position_ids')

    example = tf.io.parse_single_example(record, name_to_features)
    for key in dynamic_keys + self._mask_keys:
      example[key] = tf.sparse.to_dense(example[key])

    # Truncate padded data after the first non pad in the
    # sequence length dimension.
    # Pad before the first non pad from the back should not be removed.
    mask = tf.math.greater(
        tf.math.cumsum(example[input_ids_key], reverse=True), 0)
    for key in dynamic_keys:
      example[key] = tf.boolean_mask(example[key], mask)

    # masked_lm_ids should be 0 padded.
    # Change mask features to -1 padding so that we can differentiate
    # padding from data or from bucketizing.
    mask = tf.math.not_equal(example['masked_lm_ids'], 0)
    example['masked_lm_ids'] = tf.where(
        mask, example['masked_lm_ids'],
        -tf.ones(tf.shape(example['masked_lm_ids']), dtype=example[key].dtype))

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    # tf.data service uses dataset graph fingerprint to distinguish input
    # pipeline jobs, thus we sort the keys here to make sure they are generated
    # in a deterministic order each time the dataset function is traced.
    for name in sorted(list(example.keys())):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.cast(t, tf.int32)
      example[name] = t

    return example
//  [python]
def get_optimizer(
    optimizer_name,
    learning_rate,
    b1,
    b2,
    epsilon,
    l2_reg,
    vector_l2,
    accumulate_steps,
    mixed_precision,
    acc_grads_on_cpu,
):
    Optimizer = OPTIMIZERS.get(optimizer_name, None)
    if Optimizer is None:
        raise FinetuneError(
            "Optimizer must be in {}, not {}".format(
                list(OPTIMIZERS.keys()), optimizer_name
            )
        )

    if accumulate_steps > 1:
        Optimizer = get_grad_accumulation_optimizer(
            Optimizer, accumulate_steps, accumulate_on_cpu=acc_grads_on_cpu
        )

    decay_var_list = [
        v
        for v in tf.compat.v1.global_variables()
        if (len(v.get_shape()) > 1 or vector_l2) and "Transition_matrix" not in v.name
    ]

    opt = Optimizer(
        learning_rate=learning_rate,
        beta1=b1,
        beta2=b2,
        epsilon=epsilon,
        weight_decay=l2_reg * learning_rate,
        decay_var_list=decay_var_list,
    )

    if mixed_precision:
        opt = tf.compat.v1.train.experimental.MixedPrecisionLossScaleOptimizer(
            # Default is 2k which is too high for some of our small datasets. May be sitting at underflowing gradients for large amounts of training.
            opt,
            tf.compat.v1.train.experimental.DynamicLossScale(increment_period=100),
        )
    return opt
//  [python]
def get_optimizer(
    optimizer_name,
    learning_rate,
    b1,
    b2,
    epsilon,
    l2_reg,
    vector_l2,
    accumulate_steps,
    mixed_precision,
    acc_grads_on_cpu,
):
    Optimizer = OPTIMIZERS.get(optimizer_name, None)
    if Optimizer is None:
        raise FinetuneError(
            "Optimizer must be in {}, not {}".format(
                list(OPTIMIZERS.keys()), optimizer_name
            )
        )

    if accumulate_steps > 1:
        Optimizer = get_grad_accumulation_optimizer(
            Optimizer, accumulate_steps, accumulate_on_cpu=acc_grads_on_cpu
        )

    decay_var_list = [
        v
        for v in tf.compat.v1.global_variables()
        if (len(v.get_shape()) > 1 or vector_l2) and "Transition_matrix" not in v.name
    ]

    opt = Optimizer(
        learning_rate=learning_rate,
        beta1=b1,
        beta2=b2,
        epsilon=epsilon,
        weight_decay=l2_reg * learning_rate,
        decay_var_list=decay_var_list,
    )

    if mixed_precision:
        opt = tf.keras.mixed_precision.LossScaleOptimizer(
            opt, dynamic=True, initial_scale=2**16
        )
    return opt
//  [python]
def loss_scale_optimizer(
            cls,
            optimizer: keras.optimizers.Optimizer) -> mixedprecision.LossScaleOptimizer:
        """ Optimize loss scaling for mixed precision training.

        Parameters
        ----------
        optimizer: :class:`tf.keras.optimizers.Optimizer`
            The optimizer instance to wrap

        Returns
        --------
        :class:`tf.keras.mixed_precision.loss_scale_optimizer.LossScaleOptimizer`
            The original optimizer with loss scaling applied
        """
        return mixedprecision.LossScaleOptimizer(optimizer)
//  [python]
def loss_scale_optimizer(
            cls,
            optimizer: tf.keras.optimizers.Optimizer) -> mixedprecision.LossScaleOptimizer:
        """ Optimize loss scaling for mixed precision training.

        Parameters
        ----------
        optimizer: :class:`tf.keras.optimizers.Optimizer`
            The optimizer instance to wrap

        Returns
        --------
        :class:`tf.keras.mixed_precision.loss_scale_optimizer.LossScaleOptimizer`
            The original optimizer with loss scaling applied
        """
        return mixed_precision.LossScaleOptimizer(optimizer)
//  [python]
def loss_scale_optimizer(
            cls,
            optimizer: tf.keras.optimizers.Optimizer) -> mixedprecision.LossScaleOptimizer:
        """ Optimize loss scaling for mixed precision training.

        Parameters
        ----------
        optimizer: :class:`tf.keras.optimizers.Optimizer`
            The optimizer instance to wrap

        Returns
        --------
        :class:`tf.keras.mixed_precision.loss_scale_optimizer.LossScaleOptimizer`
            The original optimizer with loss scaling applied
        """
        return mixedprecision.LossScaleOptimizer(optimizer)
//  [python]
def __init__(self, args, env):
        super().__init__()
        
        self.n_veh = args["veh_count"]
        self.n_req_max = args["max_req_count"]
        self.batch_size = args["batch_size"]
        self.alpha = tf.exp(tf.constant(args["log_alpha"]))
        self.tau = tf.constant(args["tau"])
        self.huber_delta = tf.constant(args["huber_delta"])
        self.gradient_clipping = tf.constant(args["gradient_clipping"])
        self.clip_norm = tf.constant(args["clip_norm"])
        self.discount = tf.Variable(args["discount"], dtype=tf.float32)

        self.actor = Actor(args, env)
        self.qf1 = Critic(args, env, name="qf1")
        self.qf2 = Critic(args, env, name="qf2")
        self.qf1_target = Critic(args, env, name="qf1_target")
        self.qf2_target = Critic(args, env, name="qf2_target")
        
        if args["scheduled_lr"]:
            init_lr = args["lr"]
            decay_steps = round((args["lr_decay_steps"] - args["min_steps"]) / args["update_interval"])
            end_lr = args["lr_end"]
            self.actor_optimizer = Adam(PolynomialDecay(init_lr, decay_steps, end_lr))
            self.qf1_optimizer = Adam(PolynomialDecay(init_lr, decay_steps, end_lr))
            self.qf2_optimizer = Adam(PolynomialDecay(init_lr, decay_steps, end_lr))
        else:
            lr = args["lr"]
            self.actor_optimizer = Adam(lr)
            self.qf1_optimizer = Adam(lr)
            self.qf2_optimizer = Adam(lr)
        
        self.q1_update = tf.function(self.q_update)
        self.q2_update = tf.function(self.q_update)
//  [python]
def __init__(self, args, env):
        super().__init__()
        
        self.n_veh = args["veh_count"]
        self.n_req_max = args["max_req_count"]
        self.batch_size = args["batch_size"]
        self.alpha = tf.exp(tf.constant(args["log_alpha"]))
        self.tau = tf.constant(args["tau"])
        self.huber_delta = tf.constant(args["huber_delta"])
        self.gradient_clipping = tf.constant(args["gradient_clipping"])
        self.clip_norm = tf.constant(args["clip_norm"])
        self.discount = tf.Variable(args["discount"], dtype=tf.float32)

        self.actor = Actor(args, env)
        self.qf1 = Critic(args, env, name="qf1")
        self.qf2 = Critic(args, env, name="qf2")
        self.qf1_target = Critic(args, env, name="qf1_target")
        self.qf2_target = Critic(args, env, name="qf2_target")
        
        if args["scheduled_lr"]:
            init_lr = args["lr"]
            decay_steps = round((args["lr_decay_steps"] - args["min_steps"]) / args["update_interval"])
            end_lr = args["lr_end"]
            self.actor_optimizer = LossScaleOptimizer(Adam(PolynomialDecay(init_lr, decay_steps, end_lr)))
            self.qf1_optimizer = LossScaleOptimizer(Adam(PolynomialDecay(init_lr, decay_steps, end_lr)))
            self.qf2_optimizer = LossScaleOptimizer(Adam(PolynomialDecay(init_lr, decay_steps, end_lr)))
        else:
            lr = args["lr"]
            self.actor_optimizer = LossScaleOptimizer(Adam(lr))
            self.qf1_optimizer = LossScaleOptimizer(Adam(lr))
            self.qf2_optimizer = LossScaleOptimizer(Adam(lr))
        
        self.q1_update = tf.function(self.q_update)
        self.q2_update = tf.function(self.q_update)
//  [python]
def create_optimizers(flags):
    if flags.optimizer == 'sgd':
        embedding_optimizer = tf.keras.optimizers.legacy.SGD(learning_rate=flags.learning_rate, momentum=0)
        if flags.amp:
            embedding_optimizer = tf.train.experimental.LossScaleOptimizer(embedding_optimizer,
                                                     initial_scale=flags.loss_scale,
                                                     dynamic=False)
        mlp_optimizer = embedding_optimizer

    elif flags.optimizer == 'adam':
        embedding_optimizer = tfa.optimizers.LazyAdam(learning_rate=flags.learning_rate,
                                                      beta_1=flags.beta1, beta_2=flags.beta2)

        mlp_optimizer = tf.keras.optimizers.legacy.Adam(learning_rate=flags.learning_rate,
                                                 beta_1=flags.beta1, beta_2=flags.beta2)
        if flags.amp:
            # only wrap the mlp optimizer and not the embedding optimizer because the embeddings are not run in FP16
            mlp_optimizer = tf.train.experimental.LossScaleOptimizer(mlp_optimizer, initial_scale=flags.loss_scale, dynamic=False)

    return mlp_optimizer, embedding_optimizer
//  [python]
def create_optimizers(flags):
    if flags.optimizer == 'sgd':
        embedding_optimizer = tf.keras.optimizers.legacy.SGD(learning_rate=flags.learning_rate, momentum=0)
        if flags.amp:
            embedding_optimizer = LossScaleOptimizer(embedding_optimizer,
                                                     initial_scale=flags.loss_scale,
                                                     dynamic=False)
        mlp_optimizer = embedding_optimizer

    elif flags.optimizer == 'adam':
        embedding_optimizer = tfa.optimizers.LazyAdam(learning_rate=flags.learning_rate,
                                                      beta_1=flags.beta1, beta_2=flags.beta2)

        mlp_optimizer = tf.keras.optimizers.legacy.Adam(learning_rate=flags.learning_rate,
                                                 beta_1=flags.beta1, beta_2=flags.beta2)
        if flags.amp:
            # only wrap the mlp optimizer and not the embedding optimizer because the embeddings are not run in FP16
            mlp_optimizer = LossScaleOptimizer(mlp_optimizer, initial_scale=flags.loss_scale, dynamic=False)

    return mlp_optimizer, embedding_optimizer
//  [python]
def get_optimizer(params):
  """Get optimizer."""
  learning_rate = learning_rate_schedule(params)
  momentum = params['momentum']
  if params['optimizer'].lower() == 'sgd':
    logging.info('Use SGD optimizer')
    optimizer = tf.keras.optimizers.SGD(learning_rate, momentum=momentum)
  elif params['optimizer'].lower() == 'adam':
    logging.info('Use Adam optimizer')
    optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=momentum)
  else:
    raise ValueError('optimizers should be adam or sgd')

  moving_average_decay = params['moving_average_decay']
  if moving_average_decay:
    from tensorflow_addons import optimizers as tfa_optimizers  # pylint: disable=g-import-not-at-top
    optimizer = tfa_optimizers.MovingAverage(
        optimizer, average_decay=moving_average_decay, dynamic_decay=True)
  precision = utils.get_precision(params['strategy'], params['mixed_precision'])
  if precision == 'mixed_float16' and params['loss_scale']:
    from tensorflow.python.training import experimental as experimental_training
    optimizer = experimental_training.LossScaleOptimizer(
        optimizer,
        loss_scale=params['loss_scale'])
  return optimizer
//  [python]
def get_optimizer(params):
  """Get optimizer."""
  learning_rate = learning_rate_schedule(params)
  momentum = params['momentum']
  if params['optimizer'].lower() == 'sgd':
    logging.info('Use SGD optimizer')
    optimizer = tf.keras.optimizers.SGD(learning_rate, momentum=momentum)
  elif params['optimizer'].lower() == 'adam':
    logging.info('Use Adam optimizer')
    optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=momentum)
  else:
    raise ValueError('optimizers should be adam or sgd')

  moving_average_decay = params['moving_average_decay']
  if moving_average_decay:
    from tensorflow_addons import optimizers as tfa_optimizers  # pylint: disable=g-import-not-at-top
    optimizer = tfa_optimizers.MovingAverage(
        optimizer, average_decay=moving_average_decay, dynamic_decay=True)
  precision = utils.get_precision(params['strategy'], params['mixed_precision'])
  if precision == 'mixed_float16' and params['loss_scale']:
    optimizer = tf.keras.mixed_precision.LossScaleOptimizer(
        optimizer,
        initial_scale=params['loss_scale'])
  return optimizer
//  [python]
def get_optimizer(params):
  """Get optimizer."""
  learning_rate = learning_rate_schedule(params)
  momentum = params['momentum']
  if params['optimizer'].lower() == 'sgd':
    logging.info('Use SGD optimizer')
    optimizer = tf.keras.optimizers.SGD(learning_rate, momentum=momentum)
  elif params['optimizer'].lower() == 'adam':
    logging.info('Use Adam optimizer')
    optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=momentum)
  else:
    raise ValueError('optimizers should be adam or sgd')

  moving_average_decay = params['moving_average_decay']
  if moving_average_decay:
    # TODO(tanmingxing): potentially add dynamic_decay for new tfa release.
    from tensorflow_addons import optimizers as tfa_optimizers  # pylint: disable=g-import-not-at-top
    optimizer = tfa_optimizers.MovingAverage(
        optimizer, average_decay=moving_average_decay, dynamic_decay=True)
  precision = utils.get_precision(params['strategy'], params['mixed_precision'])
  if precision == 'mixed_float16' and params['loss_scale']:
    from tensorflow.python.training import experimental
    optimizer = experimental.LossScaleOptimizer(
        optimizer,
        initial_scale=params['loss_scale'])
  return optimizer
//  [python]
def get_optimizer(params):
  """Get optimizer."""
  learning_rate = learning_rate_schedule(params)
  momentum = params['momentum']
  if params['optimizer'].lower() == 'sgd':
    logging.info('Use SGD optimizer')
    optimizer = tf.keras.optimizers.SGD(learning_rate, momentum=momentum)
  elif params['optimizer'].lower() == 'adam':
    logging.info('Use Adam optimizer')
    optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=momentum)
  else:
    raise ValueError('optimizers should be adam or sgd')

  moving_average_decay = params['moving_average_decay']
  if moving_average_decay:
    # TODO(tanmingxing): potentially add dynamic_decay for new tfa release.
    from tensorflow_addons import optimizers as tfa_optimizers  # pylint: disable=g-import-not-at-top
    optimizer = tfa_optimizers.MovingAverage(
        optimizer, average_decay=moving_average_decay, dynamic_decay=True)
  precision = utils.get_precision(params['strategy'], params['mixed_precision'])
  if precision == 'mixed_float16' and params['loss_scale']:
    optimizer = tf.keras.mixed_precision.LossScaleOptimizer(
        optimizer,
        initial_scale=params['loss_scale'])
  return optimizer
//  [python]
def train_with_gradient_and_valuation(agent_list, grad, bi, lr, distr_type, iter_n, g_m):
    model_g = {
        'weights': g_m[0],
        'bias': g_m[1]
    }
    for i in range(iter_n-1, iter_n):
        # i->迭代轮数
        gradient_w = np.zeros([784, 10], dtype=np.float32)
        gradient_b = np.zeros([10], dtype=np.float32)
        for j in agent_list:
            gradient_w = np.add(np.multiply(grad[j][i], 1 / len(agent_list)), gradient_w)
            gradient_b = np.add(np.multiply(bi[j][i], 1 / len(agent_list)), gradient_b)
        model_g['weights'] = np.subtract(model_g['weights'], np.multiply(lr[0][i], gradient_w))
        model_g['bias'] = np.subtract(model_g['bias'], np.multiply(lr[0][i], gradient_b))

    test_images = readTestImagesFromFile(False)
    test_labels_onehot = readTestLabelsFromFile(False)
    m = np.dot(test_images, np.asarray(model_g['weights']))
    test_result = m + np.asarray(model_g['bias'])
    y = tf.nn.softmax(test_result)
    correct_prediction = tf.equal(tf.argmax(y, 1), tf.arg_max(test_labels_onehot, 1))
    #print(list(tf.argmax(y, 1).numpy()))
    #print(list(tf.arg_max(test_labels_onehot, 1).numpy()))
    #print(model_g['weights'])
    #print(model_g['bias'])
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
    return accuracy.numpy()
//  [python]
def train_with_gradient_and_valuation(agent_list, grad, bi, lr, distr_type, iter_n, g_m):
    model_g = {
        'weights': g_m[0],
        'bias': g_m[1]
    }
    for i in range(iter_n-1, iter_n):
        # i->迭代轮数
        gradient_w = np.zeros([784, 10], dtype=np.float32)
        gradient_b = np.zeros([10], dtype=np.float32)
        for j in agent_list:
            gradient_w = np.add(np.multiply(grad[j][i], 1 / len(agent_list)), gradient_w)
            gradient_b = np.add(np.multiply(bi[j][i], 1 / len(agent_list)), gradient_b)
        model_g['weights'] = np.subtract(model_g['weights'], np.multiply(lr[0][i], gradient_w))
        model_g['bias'] = np.subtract(model_g['bias'], np.multiply(lr[0][i], gradient_b))

    test_images = readTestImagesFromFile(False)
    test_labels_onehot = readTestLabelsFromFile(False)
    m = np.dot(test_images, np.asarray(model_g['weights']))
    test_result = m + np.asarray(model_g['bias'])
    y = tf.nn.softmax(test_result)
    correct_prediction = tf.equal(tf.argmax(y, 1), tf.math.argmax(test_labels_onehot, 1))
    #print(list(tf.argmax(y, 1).numpy()))
    #print(list(tf.math.argmax(test_labels_onehot, 1).numpy()))
    #print(model_g['weights'])
    #print(model_g['bias'])
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
    return accuracy.numpy()
//  [python]
def make_arg_min_max_tests(options):
  """Make a set of tests to do arg_max."""

  test_parameters = [{
      "input_dtype": [tf.float32, tf.int32],
      "input_shape": [[], [1, 1, 1, 3], [2, 3, 4, 5], [2, 3, 3], [5, 5], [10]],
      "output_type": [tf.int32, tf.int64],
      "is_arg_max": [True],
  }]

  def build_graph(parameters):
    """Build the topk op testing graph."""
    input_value = tf.placeholder(
        dtype=parameters["input_dtype"],
        name="input",
        shape=parameters["input_shape"])
    axis = random.randint(0, max(len(parameters["input_shape"]) - 1, 0))
    if parameters["is_arg_max"]:
      out = tf.arg_max(input_value, axis, output_type=parameters["output_type"])
    else:
      out = tf.arg_min(input_value, axis, output_type=parameters["output_type"])
    return [input_value], [out]

  def build_inputs(parameters, sess, inputs, outputs):
    input_value = create_tensor_data(parameters["input_dtype"],
                                     parameters["input_shape"])
    return [input_value], sess.run(
        outputs, feed_dict=dict(zip(inputs, [input_value])))

  make_zip_of_tests(
      options,
      test_parameters,
      build_graph,
      build_inputs,
      expected_tf_failures=4)
//  [python]
def make_arg_min_max_tests(options):
  """Make a set of tests to do arg_max."""

  test_parameters = [{
      "input_dtype": [tf.float32, tf.int32],
      "input_shape": [[], [1, 1, 1, 3], [2, 3, 4, 5], [2, 3, 3], [5, 5], [10]],
      "output_type": [tf.int32, tf.int64],
      "is_arg_max": [True],
  }]

  def build_graph(parameters):
    """Build the topk op testing graph."""
    input_value = tf.placeholder(
        dtype=parameters["input_dtype"],
        name="input",
        shape=parameters["input_shape"])
    axis = random.randint(0, max(len(parameters["input_shape"]) - 1, 0))
    if parameters["is_arg_max"]:
      out = tf.math.argmax(input_value, axis, output_type=parameters["output_type"])
    else:
      out = tf.math.argmin(input_value, axis, output_type=parameters["output_type"])
    return [input_value], [out]

  def build_inputs(parameters, sess, inputs, outputs):
    input_value = create_tensor_data(parameters["input_dtype"],
                                     parameters["input_shape"])
    return [input_value], sess.run(
        outputs, feed_dict=dict(zip(inputs, [input_value])))

  make_zip_of_tests(
      options,
      test_parameters,
      build_graph,
      build_inputs,
      expected_tf_failures=4)
//  [python]
def argmax2d(tensor):
  """Find the indices where the peaks value is on 2d maps.

  If there are multiple locations with the same value, return the top left one.

  Args:
    tensor: a 4-d tensor [BATCH, HEIGHT, WIDTH, N].

  Returns:
    [BATCH, N, 2] locations (row, col) for each channel and batch.
  """
  with tf.name_scope(None, 'argmax2d', [tensor]):
    shape = tf.shape(tensor)
    batch, width, channels = shape[0], shape[2], shape[3]
    flat_tensor = tf.reshape(tensor, (batch, -1, channels))
    index = tf.cast(tf.arg_max(flat_tensor, 1), tf.int32)
    y = index // width
    x = index % width
    locations = tf.stack([y, x], -1)
    return locations
//  [python]
def argmax2d(tensor):
  """Find the indices where the peaks value is on 2d maps.

  If there are multiple locations with the same value, return the top left one.

  Args:
    tensor: a 4-d tensor [BATCH, HEIGHT, WIDTH, N].

  Returns:
    [BATCH, N, 2] locations (row, col) for each channel and batch.
  """
  with tf.name_scope(None, 'argmax2d', [tensor]):
    shape = tf.shape(tensor)
    batch, width, channels = shape[0], shape[2], shape[3]
    flat_tensor = tf.reshape(tensor, (batch, -1, channels))
    index = tf.cast(tf.math.argmax(flat_tensor, 1), tf.int32)
    y = index // width
    x = index % width
    locations = tf.stack([y, x], -1)
    return locations
//  [python]
def translate(self, word):
        seq_data = [word, 'P' * len(word)]
        input_batch, output_batch, target_batch = Solution().make_batch([])
        prediction = tf.arg_max(self.model, 2)
        result = self.sess.run(prediction,
                        {self.enc_input: input_batch,
                        self.dec_input: output_batch,
                        self.targets: target_batch})
        decoded = [self.char_arr[i] for i in result[0]]
        end = decoded.index('E')
        translated = ' '.join(decoded[:end])
        return translated
//  [python]
def translate(self, word):
        seq_data = [word, 'P' * len(word)]
        input_batch, output_batch, target_batch = Solution().make_batch([])
        prediction = tf.math.argmax(self.model, 2)
        result = self.sess.run(prediction,
                        {self.enc_input: input_batch,
                        self.dec_input: output_batch,
                        self.targets: target_batch})
        decoded = [self.char_arr[i] for i in result[0]]
        end = decoded.index('E')
        translated = ' '.join(decoded[:end])
        return translated
//  [python]
def predict_keywords(self, batch):
        matching_score = self.forward_neural(batch['context_text_ids'], batch['context_length'])
        keywords_ids = self.kw_vocab.map_tokens_to_ids(batch['keywords_text'])
        kw_labels = tf.map_fn(lambda x: tf.sparse_to_dense(x, [self.kw_vocab.size], 1., 0., False),
                              keywords_ids, dtype=tf.float32, parallel_iterations=True)[:, 4:]
        loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=kw_labels, logits=matching_score)
        loss = tf.reduce_mean(loss)
        kw_ans = tf.arg_max(matching_score, -1)
        acc_label = tf.map_fn(lambda x: tf.gather(x[0], x[1]), (kw_labels, kw_ans), dtype=tf.float32)
        acc = tf.reduce_mean(acc_label)
        kws = tf.nn.top_k(matching_score, k=5)[1]
        kws = tf.reshape(kws,[-1])
        kws = tf.map_fn(lambda x: self.kw_list[x], kws, dtype=tf.int64)
        kws = tf.reshape(kws,[-1, 5])
        return loss, acc, kws
//  [python]
def predict_keywords(self, batch):
        matching_score = self.forward_neural(batch['context_text_ids'], batch['context_length'])
        keywords_ids = self.kw_vocab.map_tokens_to_ids(batch['keywords_text'])
        kw_labels = tf.map_fn(lambda x: tf.sparse_to_dense(x, [self.kw_vocab.size], 1., 0., False),
                              keywords_ids, dtype=tf.float32, parallel_iterations=True)[:, 4:]
        loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=kw_labels, logits=matching_score)
        loss = tf.reduce_mean(loss)
        kw_ans = tf.math.argmax(matching_score, -1)
        acc_label = tf.map_fn(lambda x: tf.gather(x[0], x[1]), (kw_labels, kw_ans), dtype=tf.float32)
        acc = tf.reduce_mean(acc_label)
        kws = tf.nn.top_k(matching_score, k=5)[1]
        kws = tf.reshape(kws,[-1])
        kws = tf.map_fn(lambda x: self.kw_list[x], kws, dtype=tf.int64)
        kws = tf.reshape(kws,[-1, 5])
        return loss, acc, kws
//  [python]
def forward(self, batch):
        source_embed = self.embedder(batch['source_text_ids'])
        target_embed = self.embedder(batch['target_text_ids'])
        target_embed = tf.reshape(target_embed, [-1, self.data_config._max_seq_len + 2, self.embedder.dim])
        source_code = self.source_encoder(source_embed,
                                          sequence_length_minor=batch['source_length'],
                                          sequence_length_major=batch['source_utterance_cnt'])[1]
        target_length = tf.reshape(batch['target_length'], [-1])
        target_code = self.target_encoder(target_embed, sequence_length=target_length)[1]
        target_code = tf.reshape(target_code, [-1, 20, self.config._code_len])
        source_code = tf.expand_dims(source_code, 1)
        source_code = tf.tile(source_code, [1, 20, 1])
        feature_code = target_code * source_code
        feature_code = tf.reshape(feature_code, [-1, self.config._code_len])
        logits = self.linear_matcher(feature_code)
        logits = tf.reshape(logits, [-1, 20])
        labels = tf.one_hot(batch['label'], 20)
        loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=labels, logits=logits))
        ans = tf.arg_max(logits, -1)
        acc = tx.evals.accuracy(batch['label'], ans)
        rank = tf.nn.top_k(logits, k=20)[1]
        return loss, acc, rank
//  [python]
def forward(self, batch):
        source_embed = self.embedder(batch['source_text_ids'])
        target_embed = self.embedder(batch['target_text_ids'])
        target_embed = tf.reshape(target_embed, [-1, self.data_config._max_seq_len + 2, self.embedder.dim])
        source_code = self.source_encoder(source_embed,
                                          sequence_length_minor=batch['source_length'],
                                          sequence_length_major=batch['source_utterance_cnt'])[1]
        target_length = tf.reshape(batch['target_length'], [-1])
        target_code = self.target_encoder(target_embed, sequence_length=target_length)[1]
        target_code = tf.reshape(target_code, [-1, 20, self.config._code_len])
        source_code = tf.expand_dims(source_code, 1)
        source_code = tf.tile(source_code, [1, 20, 1])
        feature_code = target_code * source_code
        feature_code = tf.reshape(feature_code, [-1, self.config._code_len])
        logits = self.linear_matcher(feature_code)
        logits = tf.reshape(logits, [-1, 20])
        labels = tf.one_hot(batch['label'], 20)
        loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=labels, logits=logits))
        ans = tf.math.argmax(logits, -1)
        acc = tx.evals.accuracy(batch['label'], ans)
        rank = tf.nn.top_k(logits, k=20)[1]
        return loss, acc, rank
//  [python]
def predict_keywords(self, batch):
        keywords_ids = self.kw_vocab.map_tokens_to_ids(batch['keywords_text'])
        matching_score = tf.map_fn(lambda x: self.forward_matrix(x), batch['context_text_ids'],
             dtype=tf.float32, parallel_iterations=True)
        kw_labels = tf.map_fn(lambda x: tf.sparse_to_dense(x, [self.kw_vocab.size], 1., 0., False),
                              keywords_ids, dtype=tf.float32, parallel_iterations=True)[:, 4:]
        kw_ans = tf.arg_max(matching_score, -1)
        acc_label = tf.map_fn(lambda x: tf.gather(x[0], x[1]), (kw_labels, kw_ans), dtype=tf.float32)
        acc = tf.reduce_mean(acc_label)
        kws = tf.nn.top_k(matching_score, k=5)[1]
        kws = tf.reshape(kws,[-1])
        kws = tf.map_fn(lambda x: self.kw_list[x], kws, dtype=tf.int64)
        kws = tf.reshape(kws,[-1, 5])
        return acc, kws
//  [python]
def predict_keywords(self, batch):
        keywords_ids = self.kw_vocab.map_tokens_to_ids(batch['keywords_text'])
        matching_score = tf.map_fn(lambda x: self.forward_matrix(x), batch['context_text_ids'],
             dtype=tf.float32, parallel_iterations=True)
        kw_labels = tf.map_fn(lambda x: tf.sparse_to_dense(x, [self.kw_vocab.size], 1., 0., False),
                              keywords_ids, dtype=tf.float32, parallel_iterations=True)[:, 4:]
        kw_ans = tf.math.argmax(matching_score, -1)
        acc_label = tf.map_fn(lambda x: tf.gather(x[0], x[1]), (kw_labels, kw_ans), dtype=tf.float32)
        acc = tf.reduce_mean(acc_label)
        kws = tf.nn.top_k(matching_score, k=5)[1]
        kws = tf.reshape(kws,[-1])
        kws = tf.map_fn(lambda x: self.kw_list[x], kws, dtype=tf.int64)
        kws = tf.reshape(kws,[-1, 5])
        return acc, kws
//  [python]
def __call__(self, sess, epoch, iter, model, loss, current_feed_dict=None):
        if iter == 0 and epoch % self.at_every_epoch == 0:
            total = 0
            correct = 0
            truth_all = []
            pred_all = []
            for values in self.batcher:
                total += len(values[-1])
                feed_dict = {}
                for i in range(0, len(self.placeholders)):
                    feed_dict[self.placeholders[i]] = values[i]
                truth = np.argmax(values[-1], 1)
                predicted = sess.run(tf.arg_max(tf.nn.softmax(model), 1), feed_dict=feed_dict)
                correct += sum(truth == predicted)
                truth_all.extend(truth)
                pred_all.extend(predicted)
            print(classification_report(truth_all, pred_all, digits=4))
//  [python]
def __call__(self, sess, epoch, iter, model, loss, current_feed_dict=None):
        if iter == 0 and epoch % self.at_every_epoch == 0:
            total = 0
            correct = 0
            truth_all = []
            pred_all = []
            for values in self.batcher:
                total += len(values[-1])
                feed_dict = {}
                for i in range(0, len(self.placeholders)):
                    feed_dict[self.placeholders[i]] = values[i]
                truth = np.argmax(values[-1], 1)
                predicted = sess.run(tf.math.argmax(tf.nn.softmax(model), 1), feed_dict=feed_dict)
                correct += sum(truth == predicted)
                truth_all.extend(truth)
                pred_all.extend(predicted)
            print(classification_report(truth_all, pred_all, digits=4))
//  [python]
def predict_keywords(self, batch):
        keywords_ids = self.kw_vocab.map_tokens_to_ids(batch['keywords_text'])
        matching_score = tf.map_fn(lambda kw_embed: self.forward_kernel(kw_embed, batch['context_text_ids']),
            self.keywords_embed, dtype=tf.float32, parallel_iterations=True)
        matching_score = tf.transpose(matching_score)
        matching_score = tf.nn.softmax(matching_score)
        kw_labels = tf.map_fn(lambda x: tf.sparse_to_dense(x, [self.kw_vocab.size], 1., 0., False),
            keywords_ids, dtype=tf.float32, parallel_iterations=True)[:, 4:]
        loss = tf.reduce_sum(-tf.log(matching_score) * kw_labels) / tf.reduce_sum(kw_labels)
        kw_ans = tf.arg_max(matching_score, -1)
        acc_label = tf.map_fn(lambda x: tf.gather(x[0], x[1]), (kw_labels, kw_ans), dtype=tf.float32)
        acc = tf.reduce_mean(acc_label)
        kws = tf.nn.top_k(matching_score, k=5)[1]
        kws = tf.reshape(kws,[-1])
        kws = tf.map_fn(lambda x: self.kw_list[x], kws, dtype=tf.int64)
        kws = tf.reshape(kws,[-1, 5])
        return loss, acc, kws
//  [python]
def predict_keywords(self, batch):
        keywords_ids = self.kw_vocab.map_tokens_to_ids(batch['keywords_text'])
        matching_score = tf.map_fn(lambda kw_embed: self.forward_kernel(kw_embed, batch['context_text_ids']),
            self.keywords_embed, dtype=tf.float32, parallel_iterations=True)
        matching_score = tf.transpose(matching_score)
        matching_score = tf.nn.softmax(matching_score)
        kw_labels = tf.map_fn(lambda x: tf.sparse_to_dense(x, [self.kw_vocab.size], 1., 0., False),
            keywords_ids, dtype=tf.float32, parallel_iterations=True)[:, 4:]
        loss = tf.reduce_sum(-tf.log(matching_score) * kw_labels) / tf.reduce_sum(kw_labels)
        kw_ans = tf.math.argmax(matching_score, -1)
        acc_label = tf.map_fn(lambda x: tf.gather(x[0], x[1]), (kw_labels, kw_ans), dtype=tf.float32)
        acc = tf.reduce_mean(acc_label)
        kws = tf.nn.top_k(matching_score, k=5)[1]
        kws = tf.reshape(kws,[-1])
        kws = tf.map_fn(lambda x: self.kw_list[x], kws, dtype=tf.int64)
        kws = tf.reshape(kws,[-1, 5])
        return loss, acc, kws
//  [python]
def train_with_gradient_and_valuation(agent_list, grad, bi, lr, distr_type, datanum, iter_n, g_m):
    model_g = {
        'weights': g_m[0],
        'bias': g_m[1]
    }
    data_sum = 0
    for i in agent_list:
        data_sum += datanum[i]
    agents_w = [0 for _ in range(NUM_AGENT)]
    for i in agent_list:
        agents_w[i] = datanum[i] / data_sum
    for i in range(iter_n - 1, iter_n):
        # i->迭代轮数
        gradient_w = np.zeros([784, 10], dtype=np.float32)
        gradient_b = np.zeros([10], dtype=np.float32)
        for j in agent_list:
            gradient_w = np.add(np.multiply(grad[j][i], agents_w[j]), gradient_w)
            gradient_b = np.add(np.multiply(bi[j][i], agents_w[j]), gradient_b)
        model_g['weights'] = np.subtract(model_g['weights'], np.multiply(lr[0][i], gradient_w))
        model_g['bias'] = np.subtract(model_g['bias'], np.multiply(lr[0][i], gradient_b))

    test_images = readTestImagesFromFile(False)
    test_labels_onehot = readTestLabelsFromFile(False)
    m = np.dot(test_images, np.asarray(model_g['weights']))
    test_result = m + np.asarray(model_g['bias'])
    y = tf.nn.softmax(test_result)
    correct_prediction = tf.equal(tf.argmax(y, 1), tf.arg_max(test_labels_onehot, 1))
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
    return accuracy.numpy()
//  [python]
def train_with_gradient_and_valuation(agent_list, grad, bi, lr, distr_type, datanum, iter_n, g_m):
    model_g = {
        'weights': g_m[0],
        'bias': g_m[1]
    }
    data_sum = 0
    for i in agent_list:
        data_sum += datanum[i]
    agents_w = [0 for _ in range(NUM_AGENT)]
    for i in agent_list:
        agents_w[i] = datanum[i] / data_sum
    for i in range(iter_n - 1, iter_n):
        # i->迭代轮数
        gradient_w = np.zeros([784, 10], dtype=np.float32)
        gradient_b = np.zeros([10], dtype=np.float32)
        for j in agent_list:
            gradient_w = np.add(np.multiply(grad[j][i], agents_w[j]), gradient_w)
            gradient_b = np.add(np.multiply(bi[j][i], agents_w[j]), gradient_b)
        model_g['weights'] = np.subtract(model_g['weights'], np.multiply(lr[0][i], gradient_w))
        model_g['bias'] = np.subtract(model_g['bias'], np.multiply(lr[0][i], gradient_b))

    test_images = readTestImagesFromFile(False)
    test_labels_onehot = readTestLabelsFromFile(False)
    m = np.dot(test_images, np.asarray(model_g['weights']))
    test_result = m + np.asarray(model_g['bias'])
    y = tf.nn.softmax(test_result)
    correct_prediction = tf.equal(tf.argmax(y, 1), tf.math.argmax(test_labels_onehot, 1))
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
    return accuracy.numpy()
//  [python]
def __call__(self, sess, epoch, iteration, model, loss):
        if iteration == 0 and epoch % self.at_every_epoch == 0:
            total = 0
            correct = 0
            for values in self.batcher:
                total += len(values[-1])
                feed_dict = {}
                for i in range(0, len(self.placeholders)):
                    # print(values[i].shape)
                    if(len(values[i].shape)==1):
                        values[i] = values[i].reshape(values[i].shape[0],1)
                    feed_dict[self.placeholders[i]] = values[i]
                truth = np.argmax(values[-1], 1)
                predicted = sess.run(tf.arg_max(tf.nn.softmax(model), 1),
                                     feed_dict=feed_dict)
                correct += sum(truth == predicted)
            acc = float(correct) / total
            self.update_summary(sess, iteration, ACCURACY_TRACE_TAG, acc)
            print("Epoch " + str(epoch) +
                  "\tAcc " + str(acc) +
                  "\tCorrect " + str(correct) + "\tTotal " + str(total))
//  [python]
def __call__(self, sess, epoch, iteration, model, loss):
        if iteration == 0 and epoch % self.at_every_epoch == 0:
            total = 0
            correct = 0
            for values in self.batcher:
                total += len(values[-1])
                feed_dict = {}
                for i in range(0, len(self.placeholders)):
                    # print(values[i].shape)
                    if(len(values[i].shape)==1):
                        values[i] = values[i].reshape(values[i].shape[0],1)
                    feed_dict[self.placeholders[i]] = values[i]
                truth = np.argmax(values[-1], 1)
                predicted = sess.run(tf.math.argmax(tf.nn.softmax(model), 1),
                                     feed_dict=feed_dict)
                correct += sum(truth == predicted)
            acc = float(correct) / total
            self.update_summary(sess, iteration, ACCURACY_TRACE_TAG, acc)
            print("Epoch " + str(epoch) +
                  "\tAcc " + str(acc) +
                  "\tCorrect " + str(correct) + "\tTotal " + str(total))
//  [python]
def train_with_gradient_and_valuation(agent_list, grad, bi, lr, distr_type, iter_n, g_m):
    model_g = {
        'weights': g_m[0],
        'bias': g_m[1]
    }

    for i in range(iter_n-1, iter_n):
        # i->迭代轮数
        gradient_w = np.zeros([784, 10], dtype=np.float32)
        gradient_b = np.zeros([10], dtype=np.float32)
        for j in agent_list:
            gradient_w = np.add(np.multiply(grad[j][i], 1 / len(agent_list)), gradient_w)
            gradient_b = np.add(np.multiply(bi[j][i], 1 / len(agent_list)), gradient_b)
        model_g['weights'] = np.subtract(model_g['weights'], np.multiply(lr[0][i], gradient_w))
        model_g['bias'] = np.subtract(model_g['bias'], np.multiply(lr[0][i], gradient_b))

    test_images = readTestImagesFromFile(False)
    test_labels_onehot = readTestLabelsFromFile(False)
    m = np.dot(test_images, np.asarray(model_g['weights']))
    test_result = m + np.asarray(model_g['bias'])
    y = tf.nn.softmax(test_result)
    correct_prediction = tf.equal(tf.argmax(y, 1), tf.arg_max(test_labels_onehot, 1))
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
    return accuracy.numpy()
//  [python]
def train_with_gradient_and_valuation(agent_list, grad, bi, lr, distr_type, iter_n, g_m):
    model_g = {
        'weights': g_m[0],
        'bias': g_m[1]
    }

    for i in range(iter_n-1, iter_n):
        # i->迭代轮数
        gradient_w = np.zeros([784, 10], dtype=np.float32)
        gradient_b = np.zeros([10], dtype=np.float32)
        for j in agent_list:
            gradient_w = np.add(np.multiply(grad[j][i], 1 / len(agent_list)), gradient_w)
            gradient_b = np.add(np.multiply(bi[j][i], 1 / len(agent_list)), gradient_b)
        model_g['weights'] = np.subtract(model_g['weights'], np.multiply(lr[0][i], gradient_w))
        model_g['bias'] = np.subtract(model_g['bias'], np.multiply(lr[0][i], gradient_b))

    test_images = readTestImagesFromFile(False)
    test_labels_onehot = readTestLabelsFromFile(False)
    m = np.dot(test_images, np.asarray(model_g['weights']))
    test_result = m + np.asarray(model_g['bias'])
    y = tf.nn.softmax(test_result)
    correct_prediction = tf.equal(tf.argmax(y, 1), tf.math.argmax(test_labels_onehot, 1))
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
    return accuracy.numpy()
//  [python]
def mean_score_classification_accuracy(y_true, y_pred):
    """
    Use the mean score of an image against all the samples from the same class to get a score per class for each image.
    """
    y_pred = tf.linalg.normalize(
        tf.linalg.matmul(y_pred, tf.math.divide_no_nan(y_true, tf.reduce_sum(y_true, axis=0))), ord=1, axis=1
    )[0]
    y_pred = tf.map_fn(lambda x: y_true[x], tf.argmax(y_pred, axis=1), dtype=y_pred.dtype)
    return tf.reduce_mean(tf.reduce_sum(y_true * y_pred, axis=1))
//  [python]
def mean_score_classification_accuracy(y_true, y_pred):
    """
    Use the mean score of an image against all the samples from the same class to get a score per class for each image.
    """
    y_pred = tf.linalg.normalize(
        tf.linalg.matmul(y_pred, tf.math.divide_no_nan(y_true, tf.reduce_sum(y_true, axis=0))), ord=1, axis=1
    )[0]
    y_pred = tf.map_fn(lambda x: y_true[x], tf.math.argmax(y_pred, axis=1), dtype=y_pred.dtype)
    return tf.reduce_mean(tf.reduce_sum(y_true * y_pred, axis=1))
//  [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.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