jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def negative_mse(states,
                 actions,
                 rewards,
                 next_states,
                 contexts,
                 state_scales=1.0,
                 goal_scales=1.0,
                 summarize=False):
  """Returns the negative mean square error between next_states and contexts.

  Args:
    states: A [batch_size, num_state_dims] Tensor representing a batch
        of states.
    actions: A [batch_size, num_action_dims] Tensor representing a batch
      of actions.
    rewards: A [batch_size] Tensor representing a batch of rewards.
    next_states: A [batch_size, num_state_dims] Tensor representing a batch
      of next states.
    contexts: A list of [batch_size, num_context_dims] Tensor representing
      a batch of contexts.
    state_scales: multiplicative scale for (next) states. A scalar or 1D tensor,
      must be broadcastable to number of state dimensions.
    goal_scales: multiplicative scale for contexts. A scalar or 1D tensor,
      must be broadcastable to number of goal dimensions.
    summarize: (boolean) enable summary ops.

  Returns:
    A new tf.float32 [batch_size] rewards Tensor, and
      tf.float32 [batch_size] discounts tensor.
  """
  del states, actions, rewards  # Unused
  mse = tf.reduce_mean(tf.squared_difference(next_states * state_scales,
                                             contexts[0] * goal_scales), -1)
  if summarize:
    with tf.name_scope('RewardFn/'):
      tf.summary.scalar('mean_mse', tf.reduce_mean(mse))
      tf.summary.histogram('mse', mse)
  rewards = tf.to_float(-mse)
  return rewards, tf.ones_like(rewards)
//  [python]
def negative_mse(states,
                 actions,
                 rewards,
                 next_states,
                 contexts,
                 state_scales=1.0,
                 goal_scales=1.0,
                 summarize=False):
  """Returns the negative mean square error between next_states and contexts.

  Args:
    states: A [batch_size, num_state_dims] Tensor representing a batch
        of states.
    actions: A [batch_size, num_action_dims] Tensor representing a batch
      of actions.
    rewards: A [batch_size] Tensor representing a batch of rewards.
    next_states: A [batch_size, num_state_dims] Tensor representing a batch
      of next states.
    contexts: A list of [batch_size, num_context_dims] Tensor representing
      a batch of contexts.
    state_scales: multiplicative scale for (next) states. A scalar or 1D tensor,
      must be broadcastable to number of state dimensions.
    goal_scales: multiplicative scale for contexts. A scalar or 1D tensor,
      must be broadcastable to number of goal dimensions.
    summarize: (boolean) enable summary ops.

  Returns:
    A new tf.float32 [batch_size] rewards Tensor, and
      tf.float32 [batch_size] discounts tensor.
  """
  del states, actions, rewards  # Unused
  mse = tf.reduce_mean(tf.squared_difference(next_states * state_scales,
                                             contexts[0] * goal_scales), -1)
  if summarize:
    with tf.name_scope('RewardFn/'):
      tf.summary.scalar('mean_mse', tf.reduce_mean(mse))
      tf.summary.histogram('mse', mse)
  rewards = tf.cast(-mse, tf.float32)
  return rewards, tf.ones_like(rewards)
//  [python]
def binary_indicator(states,
                     actions,
                     rewards,
                     next_states,
                     contexts,
                     termination_epsilon=1e-4,
                     offset=0,
                     epsilon=1e-10,
                     state_indices=None,
                     summarize=False):
  """Returns 0/1 by checking if next_states and contexts overlap.

  Args:
    states: A [batch_size, num_state_dims] Tensor representing a batch
        of states.
    actions: A [batch_size, num_action_dims] Tensor representing a batch
      of actions.
    rewards: A [batch_size] Tensor representing a batch of rewards.
    next_states: A [batch_size, num_state_dims] Tensor representing a batch
      of next states.
    contexts: A list of [batch_size, num_context_dims] Tensor representing
      a batch of contexts.
    termination_epsilon: terminate if dist is less than this quantity.
    offset: Offset the rewards.
    epsilon: small offset to ensure non-negative/zero distance.

  Returns:
    A new tf.float32 [batch_size] rewards Tensor, and
      tf.float32 [batch_size] discounts tensor.
  """
  del states, actions  # unused args
  next_states = index_states(next_states, state_indices)
  dist = tf.reduce_sum(tf.squared_difference(next_states, contexts[0]), -1)
  dist = tf.sqrt(dist + epsilon)
  discounts = dist > termination_epsilon
  rewards = tf.logical_not(discounts)
  rewards = tf.to_float(rewards) + offset
  return tf.to_float(rewards), tf.ones_like(tf.to_float(discounts))
//  [python]
def binary_indicator(states,
                     actions,
                     rewards,
                     next_states,
                     contexts,
                     termination_epsilon=1e-4,
                     offset=0,
                     epsilon=1e-10,
                     state_indices=None,
                     summarize=False):
  """Returns 0/1 by checking if next_states and contexts overlap.

  Args:
    states: A [batch_size, num_state_dims] Tensor representing a batch
        of states.
    actions: A [batch_size, num_action_dims] Tensor representing a batch
      of actions.
    rewards: A [batch_size] Tensor representing a batch of rewards.
    next_states: A [batch_size, num_state_dims] Tensor representing a batch
      of next states.
    contexts: A list of [batch_size, num_context_dims] Tensor representing
      a batch of contexts.
    termination_epsilon: terminate if dist is less than this quantity.
    offset: Offset the rewards.
    epsilon: small offset to ensure non-negative/zero distance.

  Returns:
    A new tf.float32 [batch_size] rewards Tensor, and
      tf.float32 [batch_size] discounts tensor.
  """
  del states, actions  # unused args
  next_states = index_states(next_states, state_indices)
  dist = tf.reduce_sum(tf.squared_difference(next_states, contexts[0]), -1)
  dist = tf.sqrt(dist + epsilon)
  discounts = dist > termination_epsilon
  rewards = tf.logical_not(discounts)
  rewards = tf.cast(rewards, tf.float32) + offset
  return tf.cast(rewards, tf.float32), tf.ones_like(tf.cast(discounts, tf.float32))
//  [python]
def timed_rewards(states,
                  actions,
                  rewards,
                  next_states,
                  contexts,
                  reward_fn=None,
                  dense=False,
                  timer_index=-1):
  """Return the timed rewards.

  Args:
    states: A [batch_size, num_state_dims] Tensor representing a batch
        of states.
    actions: A [batch_size, num_action_dims] Tensor representing a batch
      of actions.
    rewards: A [batch_size] Tensor representing a batch of rewards.
    next_states: A [batch_size, num_state_dims] Tensor representing a batch
      of next states.
    contexts: A list of [batch_size, num_context_dims] Tensor representing
      a batch of contexts.
    reward_fn: A reward function.
    dense: (boolean) Provide dense rewards or sparse rewards at time = 0.
    timer_index: (integer) The context list index that specifies timer.
  Returns:
    A new tf.float32 [batch_size] rewards Tensor, and
      tf.float32 [batch_size] discounts tensor.
  """
  assert contexts[timer_index].get_shape().as_list()[1] == 1
  timers = contexts[timer_index][:, 0]
  rewards, discounts = reward_fn(states, actions, rewards, next_states,
                                 contexts)
  terminates = tf.to_float(timers <= 0)  # if terminate set 1, else set 0
  for _ in range(rewards.shape.ndims - 1):
    terminates = tf.expand_dims(terminates, axis=-1)
  if not dense:
    rewards *= terminates  # if terminate, return rewards, else return 0
  discounts *= (tf.to_float(1.0) - terminates)
  return rewards, discounts
//  [python]
def timed_rewards(states,
                  actions,
                  rewards,
                  next_states,
                  contexts,
                  reward_fn=None,
                  dense=False,
                  timer_index=-1):
  """Return the timed rewards.

  Args:
    states: A [batch_size, num_state_dims] Tensor representing a batch
        of states.
    actions: A [batch_size, num_action_dims] Tensor representing a batch
      of actions.
    rewards: A [batch_size] Tensor representing a batch of rewards.
    next_states: A [batch_size, num_state_dims] Tensor representing a batch
      of next states.
    contexts: A list of [batch_size, num_context_dims] Tensor representing
      a batch of contexts.
    reward_fn: A reward function.
    dense: (boolean) Provide dense rewards or sparse rewards at time = 0.
    timer_index: (integer) The context list index that specifies timer.
  Returns:
    A new tf.float32 [batch_size] rewards Tensor, and
      tf.float32 [batch_size] discounts tensor.
  """
  assert contexts[timer_index].get_shape().as_list()[1] == 1
  timers = contexts[timer_index][:, 0]
  rewards, discounts = reward_fn(states, actions, rewards, next_states,
                                 contexts)
  terminates = tf.cast(timers <= 0, tf.float32)  # if terminate set 1, else set 0
  for _ in range(rewards.shape.ndims - 1):
    terminates = tf.expand_dims(terminates, axis=-1)
  if not dense:
    rewards *= terminates  # if terminate, return rewards, else return 0
  discounts *= (tf.cast(1.0, tf.float32) - terminates)
  return rewards, discounts
//  [python]
def rnn_classification(core_fn,
                       vocab_size=10000,
                       embed_dim=64,
                       aggregate_method="last"):
  """Helper for RNN based text classification tasks.

  Args:
    core_fn: callable callable that returns a sonnet RNN core
    vocab_size: int number of words to use for the embedding table. All index
      higher than this will be clipped
    embed_dim: int size of the embedding dim
    aggregate_method: str how to aggregate the sequence of features. If 'last'
      grab the last hidden features. If 'avg' compute the average over the full
      sequence.

  Returns:
    a callable that returns a sonnet module representing the loss.
  """

  def _build(batch):
    """Build the loss sonnet module."""
    # TODO(lmetz) these are dense updates.... so keeping this small for now.
    tokens = tf.minimum(batch["text"],
                        tf.to_int64(tf.reshape(vocab_size - 1, [1, 1])))
    embed = snt.Embed(vocab_size=vocab_size, embed_dim=embed_dim)
    embedded_tokens = embed(tokens)
    rnn = core_fn()
    bs = tokens.shape.as_list()[0]

    state = rnn.initial_state(bs, trainable=True)

    outputs, state = tf.nn.dynamic_rnn(
        rnn, embedded_tokens, initial_state=state)
    if aggregate_method == "last":
      last_output = outputs[:, -1]  # grab the last output
    elif aggregate_method == "avg":
      last_output = tf.reduce_mean(outputs, [1])  # average over length
    else:
      raise ValueError("not supported aggregate_method")

    logits = snt.Linear(2)(last_output)

    loss_vec = tf.nn.softmax_cross_entropy_with_logits_v2(
        labels=batch["label_onehot"], logits=logits)

    return tf.reduce_mean(loss_vec)

  return lambda: snt.Module(_build)
//  [python]
def rnn_classification(core_fn,
                       vocab_size=10000,
                       embed_dim=64,
                       aggregate_method="last"):
  """Helper for RNN based text classification tasks.

  Args:
    core_fn: callable callable that returns a sonnet RNN core
    vocab_size: int number of words to use for the embedding table. All index
      higher than this will be clipped
    embed_dim: int size of the embedding dim
    aggregate_method: str how to aggregate the sequence of features. If 'last'
      grab the last hidden features. If 'avg' compute the average over the full
      sequence.

  Returns:
    a callable that returns a sonnet module representing the loss.
  """

  def _build(batch):
    """Build the loss sonnet module."""
    # TODO(lmetz) these are dense updates.... so keeping this small for now.
    tokens = tf.minimum(batch["text"],
                        tf.cast(tf.reshape(vocab_size - 1, [1, 1]), tf.int64))
    embed = snt.Embed(vocab_size=vocab_size, embed_dim=embed_dim)
    embedded_tokens = embed(tokens)
    rnn = core_fn()
    bs = tokens.shape.as_list()[0]

    state = rnn.initial_state(bs, trainable=True)

    outputs, state = tf.nn.dynamic_rnn(
        rnn, embedded_tokens, initial_state=state)
    if aggregate_method == "last":
      last_output = outputs[:, -1]  # grab the last output
    elif aggregate_method == "avg":
      last_output = tf.reduce_mean(outputs, [1])  # average over length
    else:
      raise ValueError("not supported aggregate_method")

    logits = snt.Linear(2)(last_output)

    loss_vec = tf.nn.softmax_cross_entropy_with_logits_v2(
        labels=batch["label_onehot"], logits=logits)

    return tf.reduce_mean(loss_vec)

  return lambda: snt.Module(_build)
//  [python]
def vgg_graph(sess, phs):
  VGG = model.VGG()
  Y_score = VGG.build(phs['batch_images'], N_CLASS, phs['is_training_ph'])

  Y_hat = tf.nn.softmax(Y_score)
  Y_pred = tf.argmax(Y_hat, 1)
  Y_label = tf.to_float(tf.one_hot(phs['batch_labels'], N_CLASS))

  cross_entropy = tf.nn.softmax_cross_entropy_with_logits(logits = Y_score, labels = Y_label)
  loss_op = tf.reduce_mean(cross_entropy)
  correct_prediction = tf.equal(tf.argmax(Y_hat, 1), tf.argmax(Y_label, 1))
  acc_op = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

  update_op = tf.train.MomentumOptimizer(LR, MOMENTUM).minimize(loss_op, var_list=VGG.vars)

  return loss_op, acc_op, cross_entropy, Y_hat, update_op, Y_pred, VGG.vars
//  [python]
def vgg_graph(sess, phs):
  VGG = model.VGG()
  Y_score = VGG.build(phs['batch_images'], N_CLASS, phs['is_training_ph'])

  Y_hat = tf.nn.softmax(Y_score)
  Y_pred = tf.argmax(Y_hat, 1)
  Y_label = tf.cast(tf.one_hot(phs['batch_labels'], N_CLASS), tf.float32)

  cross_entropy = tf.nn.softmax_cross_entropy_with_logits(logits = Y_score, labels = Y_label)
  loss_op = tf.reduce_mean(cross_entropy)
  correct_prediction = tf.equal(tf.argmax(Y_hat, 1), tf.argmax(Y_label, 1))
  acc_op = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

  update_op = tf.train.MomentumOptimizer(LR, MOMENTUM).minimize(loss_op, var_list=VGG.vars)

  return loss_op, acc_op, cross_entropy, Y_hat, update_op, Y_pred, VGG.vars
//  [python]
def get_dropout(self, dropout_rate, is_training):
    return 1 - (tf.to_float(is_training) * dropout_rate)
//  [python]
def get_dropout(self, dropout_rate, is_training):
    return 1 - (tf.cast(is_training, tf.float32) * dropout_rate)
//  [python]
def softmax_loss(self, antecedent_scores, antecedent_labels):
    gold_scores = antecedent_scores + tf.log(tf.to_float(antecedent_labels)) # [k, max_ant + 1]
    marginalized_gold_scores = tf.reduce_logsumexp(gold_scores, [1]) # [k]
    log_norm = tf.reduce_logsumexp(antecedent_scores, [1]) # [k]
    return log_norm - marginalized_gold_scores
//  [python]
def softmax_loss(self, antecedent_scores, antecedent_labels):
    gold_scores = antecedent_scores + tf.log(tf.cast(antecedent_labels, dtype=tf.float32)) # [k, max_ant + 1]
    marginalized_gold_scores = tf.reduce_logsumexp(gold_scores, [1]) # [k]
    log_norm = tf.reduce_logsumexp(antecedent_scores, [1]) # [k]
    return log_norm - marginalized_gold_scores
//  [python]
def get_candidate_labels(self, candidate_starts, candidate_ends, labeled_starts, labeled_ends, labels):
    same_start = tf.equal(tf.expand_dims(labeled_starts, 1), tf.expand_dims(candidate_starts, 0)) # [num_labeled, num_candidates]
    same_end = tf.equal(tf.expand_dims(labeled_ends, 1), tf.expand_dims(candidate_ends, 0)) # [num_labeled, num_candidates]
    same_span = tf.logical_and(same_start, same_end) # [num_labeled, num_candidates]
    candidate_labels = tf.matmul(tf.expand_dims(labels, 0), tf.to_int32(same_span)) # [1, num_candidates]
    candidate_labels = tf.squeeze(candidate_labels, 0) # [num_candidates]
    return candidate_labels
//  [python]
def get_candidate_labels(self, candidate_starts, candidate_ends, labeled_starts, labeled_ends, labels):
    same_start = tf.equal(tf.expand_dims(labeled_starts, 1), tf.expand_dims(candidate_starts, 0)) # [num_labeled, num_candidates]
    same_end = tf.equal(tf.expand_dims(labeled_ends, 1), tf.expand_dims(candidate_ends, 0)) # [num_labeled, num_candidates]
    same_span = tf.logical_and(same_start, same_end) # [num_labeled, num_candidates]
    candidate_labels = tf.matmul(tf.expand_dims(labels, 0), tf.cast(same_span, tf.int32)) # [1, num_candidates]
    candidate_labels = tf.squeeze(candidate_labels, 0) # [num_candidates]
    return candidate_labels
//  [python]
def bucket_distance(self, distances):
    """
    Places the given values (designed for distances) into 10 semi-logscale buckets:
    [0, 1, 2, 3, 4, 5-7, 8-15, 16-31, 32-63, 64+].
    """
    logspace_idx = tf.to_int32(tf.floor(tf.log(tf.to_float(distances))/math.log(2))) + 3
    use_identity = tf.to_int32(distances <= 4)
    combined_idx = use_identity * distances + (1 - use_identity) * logspace_idx
    return tf.clip_by_value(combined_idx, 0, 9)
//  [python]
def bucket_distance(self, distances):
    """
    Places the given values (designed for distances) into 10 semi-logscale buckets:
    [0, 1, 2, 3, 4, 5-7, 8-15, 16-31, 32-63, 64+].
    """
    logspace_idx = tf.cast(tf.floor(tf.math.log(tf.cast(distances, tf.float32))/math.log(2)), tf.int32) + 3
    use_identity = tf.cast(distances <= 4, tf.int32)
    combined_idx = use_identity * distances + (1 - use_identity) * logspace_idx
    return tf.clip_by_value(combined_idx, 0, 9)
//  [python]
def coarse_to_fine_pruning(self, top_span_emb, top_span_mention_scores, c):
    k = util.shape(top_span_emb, 0)
    top_span_range = tf.range(k) # [k]
    antecedent_offsets = tf.expand_dims(top_span_range, 1) - tf.expand_dims(top_span_range, 0) # [k, k]
    antecedents_mask = antecedent_offsets >= 1 # [k, k]
    fast_antecedent_scores = tf.expand_dims(top_span_mention_scores, 1) + tf.expand_dims(top_span_mention_scores, 0) # [k, k]
    fast_antecedent_scores += tf.log(tf.to_float(antecedents_mask)) # [k, k]
    fast_antecedent_scores += self.get_fast_antecedent_scores(top_span_emb) # [k, k]

    _, top_antecedents = tf.nn.top_k(fast_antecedent_scores, c, sorted=False) # [k, c]
    top_antecedents_mask = util.batch_gather(antecedents_mask, top_antecedents) # [k, c]
    top_fast_antecedent_scores = util.batch_gather(fast_antecedent_scores, top_antecedents) # [k, c]
    top_antecedent_offsets = util.batch_gather(antecedent_offsets, top_antecedents) # [k, c]
    return top_antecedents, top_antecedents_mask, top_fast_antecedent_scores, top_antecedent_offsets
//  [python]
def coarse_to_fine_pruning(self, top_span_emb, top_span_mention_scores, c):
    k = util.shape(top_span_emb, 0)
    top_span_range = tf.range(k) # [k]
    antecedent_offsets = tf.expand_dims(top_span_range, 1) - tf.expand_dims(top_span_range, 0) # [k, k]
    antecedents_mask = antecedent_offsets >= 1 # [k, k]
    fast_antecedent_scores = tf.expand_dims(top_span_mention_scores, 1) + tf.expand_dims(top_span_mention_scores, 0) # [k, k]
    fast_antecedent_scores += tf.log(tf.cast(antecedents_mask, dtype=tf.float32)) # [k, k]
    fast_antecedent_scores += self.get_fast_antecedent_scores(top_span_emb) # [k, k]

    _, top_antecedents = tf.nn.top_k(fast_antecedent_scores, c, sorted=False) # [k, c]
    top_antecedents_mask = util.batch_gather(antecedents_mask, top_antecedents) # [k, c]
    top_fast_antecedent_scores = util.batch_gather(fast_antecedent_scores, top_antecedents) # [k, c]
    top_antecedent_offsets = util.batch_gather(antecedent_offsets, top_antecedents) # [k, c]
    return top_antecedents, top_antecedents_mask, top_fast_antecedent_scores, top_antecedent_offsets
//  [python]
def statistics_summaries(name, var):
  """Attach additional statistical summaries to the variable."""
  var = tf.to_float(var)
  with tf.variable_scope(name):
    tf.summary.scalar('mean', tf.reduce_mean(var))
    tf.summary.scalar('stddev', tf.math.reduce_std(var))
    tf.summary.scalar('max', tf.reduce_max(var))
    tf.summary.scalar('min', tf.reduce_min(var))
  tf.summary.histogram(name, var)
//  [python]
def statistics_summaries(name, var):
  """Attach additional statistical summaries to the variable."""
  var = tf.cast(var, tf.float32)
  with tf.variable_scope(name):
    tf.summary.scalar('mean', tf.reduce_mean(var))
    tf.summary.scalar('stddev', tf.math.reduce_std(var))
    tf.summary.scalar('max', tf.reduce_max(var))
    tf.summary.scalar('min', tf.reduce_min(var))
  tf.summary.histogram(name, var)
//  [python]
def _smooth_l1_loss_base(bbox_pred, bbox_targets, sigma=1.0):
    '''

    :param bbox_pred: [-1, 4] in RPN. [-1, cls_num+1, 4] in Fast-rcnn
    :param bbox_targets: shape is same as bbox_pred
    :param sigma:
    :return:
    '''
    sigma_2 = sigma**2

    box_diff = bbox_pred - bbox_targets

    abs_box_diff = tf.abs(box_diff)

    smoothL1_sign = tf.stop_gradient(
        tf.to_float(tf.less(abs_box_diff, 1. / sigma_2)))
    loss_box = tf.pow(box_diff, 2) * (sigma_2 / 2.0) * smoothL1_sign \
               + (abs_box_diff - (0.5 / sigma_2)) * (1.0 - smoothL1_sign)
    return loss_box
//  [python]
def _smooth_l1_loss_base(bbox_pred, bbox_targets, sigma=1.0):
    '''

    :param bbox_pred: [-1, 4] in RPN. [-1, cls_num+1, 4] in Fast-rcnn
    :param bbox_targets: shape is same as bbox_pred
    :param sigma:
    :return:
    '''
    sigma_2 = sigma**2

    box_diff = bbox_pred - bbox_targets

    abs_box_diff = tf.abs(box_diff)

    smoothL1_sign = tf.stop_gradient(
        tf.cast(tf.less(abs_box_diff, 1. / sigma_2), tf.float32))
    loss_box = tf.pow(box_diff, 2) * (sigma_2 / 2.0) * smoothL1_sign \
               + (abs_box_diff - (0.5 / sigma_2)) * (1.0 - smoothL1_sign)
    return loss_box
//  [python]
def generate_action_mask(features):
  """Computes the decode mask from "task" and "verb_refs"."""
  eos_positions = tf.to_int32(tf.expand_dims(
      tf.where(tf.equal(features["task"], 1))[:, 1], 1))
  decode_mask = tf.cumsum(tf.to_int32(
      tf.logical_and(
          tf.equal(features["verb_refs"][:, :, 0], eos_positions),
          tf.equal(features["verb_refs"][:, :, 1], eos_positions + 1))),
                          axis=-1)
  decode_mask = tf.sequence_mask(
      tf.reduce_sum(tf.to_int32(tf.less(decode_mask, 1)), -1),
      maxlen=tf.shape(decode_mask)[1])
  return decode_mask
//  [python]
def generate_action_mask(features):
  """Computes the decode mask from "task" and "verb_refs"."""
  eos_positions = tf.cast(tf.expand_dims(
      tf.where(tf.equal(features["task"], 1))[:, 1], 1), tf.int32)
  decode_mask = tf.cumsum(tf.cast(
      tf.logical_and(
          tf.equal(features["verb_refs"][:, :, 0], eos_positions),
          tf.equal(features["verb_refs"][:, :, 1], eos_positions + 1)),
                          tf.int32),
                          axis=-1)
  decode_mask = tf.sequence_mask(
      tf.reduce_sum(tf.cast(tf.less(decode_mask, 1), tf.int32), -1),
      maxlen=tf.shape(decode_mask)[1])
  return decode_mask
//  [python]
def annealed_annealer(gamma, step, iteration_threshold):
  """Linear annealing."""
  return tf.math.minimum(gamma * 1.,
                         gamma * 1. * tf.to_float(step) / iteration_threshold)
//  [python]
def annealed_annealer(gamma, step, iteration_threshold):
  """Linear annealing."""
  return tf.math.minimum(gamma * 1.,
                         gamma * 1. * tf.cast(step, tf.float32) / iteration_threshold)
//  [python]
def file_based_input_fn_builder(input_file, seq_length, is_training,
                                drop_remainder):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

  name_to_features = {
      "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "input_mask": tf.FixedLenFeature([seq_length], tf.int64),
      "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "label_ids": tf.FixedLenFeature([], tf.int64),
      "is_real_example": tf.FixedLenFeature([], tf.int64),
  }

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.to_int32(t)
      example[name] = t

    return example

  def input_fn(params):
    """The actual input function."""
    batch_size = params["batch_size"]

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    d = tf.data.TFRecordDataset(input_file)
    if is_training:
      d = d.repeat()
      d = d.shuffle(buffer_size=100)

    d = d.apply(
        tf.data.experimental.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder))

    return d

  return input_fn
//  [python]
def file_based_input_fn_builder(input_file, seq_length, is_training,
                                drop_remainder):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

  name_to_features = {
      "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "input_mask": tf.FixedLenFeature([seq_length], tf.int64),
      "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "label_ids": tf.FixedLenFeature([], tf.int64),
      "is_real_example": tf.FixedLenFeature([], tf.int64),
  }

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.cast(t, tf.int32)
      example[name] = t

    return example

  def input_fn(params):
    """The actual input function."""
    batch_size = params["batch_size"]

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    d = tf.data.TFRecordDataset(input_file)
    if is_training:
      d = d.repeat()
      d = d.shuffle(buffer_size=100)

    d = d.apply(
        tf.data.experimental.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder))

    return d

  return input_fn
//  [python]
def smooth_l1_loss_rcnn(bbox_pred, bbox_targets, label, num_classes, sigma=1.0):
    '''

    :param bbox_pred: [-1, (cfgs.CLS_NUM +1) * 4]
    :param bbox_targets:[-1, (cfgs.CLS_NUM +1) * 4]
    :param label:[-1]
    :param num_classes:
    :param sigma:
    :return:
    '''

    outside_mask = tf.stop_gradient(tf.to_float(tf.greater(label, 0)))

    bbox_pred = tf.reshape(bbox_pred, [-1, num_classes, 4])
    bbox_targets = tf.reshape(bbox_targets, [-1, num_classes, 4])

    value = _smooth_l1_loss_base(bbox_pred,
                                 bbox_targets,
                                 sigma=sigma)
    value = tf.reduce_sum(value, 2)
    value = tf.reshape(value, [-1, num_classes])

    inside_mask = tf.one_hot(tf.reshape(label, [-1, 1]),
                             depth=num_classes, axis=1)

    inside_mask = tf.stop_gradient(
        tf.to_float(tf.reshape(inside_mask, [-1, num_classes])))

    normalizer = tf.to_float(tf.shape(bbox_pred)[0])
    bbox_loss = tf.reduce_sum(
        tf.reduce_sum(value * inside_mask, 1)*outside_mask) / normalizer

    return bbox_loss
//  [python]
def smooth_l1_loss_rcnn(bbox_pred, bbox_targets, label, num_classes, sigma=1.0):
    '''

    :param bbox_pred: [-1, (cfgs.CLS_NUM +1) * 4]
    :param bbox_targets:[-1, (cfgs.CLS_NUM +1) * 4]
    :param label:[-1]
    :param num_classes:
    :param sigma:
    :return:
    '''

    outside_mask = tf.stop_gradient(tf.cast(tf.greater(label, 0), tf.float32))

    bbox_pred = tf.reshape(bbox_pred, [-1, num_classes, 4])
    bbox_targets = tf.reshape(bbox_targets, [-1, num_classes, 4])

    value = _smooth_l1_loss_base(bbox_pred,
                                 bbox_targets,
                                 sigma=sigma)
    value = tf.reduce_sum(value, 2)
    value = tf.reshape(value, [-1, num_classes])

    inside_mask = tf.one_hot(tf.reshape(label, [-1, 1]),
                             depth=num_classes, axis=1)

    inside_mask = tf.stop_gradient(
        tf.cast(tf.reshape(inside_mask, [-1, num_classes]), tf.float32))

    normalizer = tf.cast(tf.shape(bbox_pred)[0], tf.float32)
    bbox_loss = tf.reduce_sum(
        tf.reduce_sum(value * inside_mask, 1)*outside_mask) / normalizer

    return bbox_loss
//  [python]
def fine_tune_annealer(gamma, step, iteration_threshold):
  """Fine tuning.

  This annealer returns zero if step < iteration_threshold and gamma otherwise.

  Args:
    gamma: Weight of supervised loss.
    step: Current step of training.
    iteration_threshold: When to return gamma instead of zero.

  Returns:
    Either gamma or zero.
  """
  return gamma * tf.math.minimum(
      tf.to_float(1),
      tf.math.maximum(tf.to_float(0), tf.to_float(step - iteration_threshold)))
//  [python]
def fine_tune_annealer(gamma, step, iteration_threshold):
  """Fine tuning.

  This annealer returns zero if step < iteration_threshold and gamma otherwise.

  Args:
    gamma: Weight of supervised loss.
    step: Current step of training.
    iteration_threshold: When to return gamma instead of zero.

  Returns:
    Either gamma or zero.
  """
  return gamma * tf.math.minimum(
      tf.cast(1, tf.float32),
      tf.math.maximum(tf.cast(0, tf.float32), tf.cast(step - iteration_threshold, tf.float32)))
//  [python]
def exclude_ignore_index(self, logits, labels):
        """Ignore certain index."""
        logits = tf.transpose(logits, [0, 2, 3, 1])
        if vega.is_gpu_device():
            indices = tf.where(tf.not_equal(labels, self.ignore_index))
            labels = tf.cast(tf.gather_nd(labels, indices), tf.int32)
            logits = tf.gather_nd(logits, indices)
            return logits, labels, 1.0
        elif vega.is_npu_device():
            weights = tf.not_equal(labels, self.ignore_index)
            labels = tf.multiply(labels, tf.cast(weights, labels.dtype))
            return logits, labels, tf.to_float(weights)
//  [python]
def exclude_ignore_index(self, logits, labels):
        """Ignore certain index."""
        logits = tf.transpose(logits, [0, 2, 3, 1])
        if vega.is_gpu_device():
            indices = tf.where(tf.not_equal(labels, self.ignore_index))
            labels = tf.cast(tf.gather_nd(labels, indices), tf.int32)
            logits = tf.gather_nd(logits, indices)
            return logits, labels, 1.0
        elif vega.is_npu_device():
            weights = tf.not_equal(labels, self.ignore_index)
            labels = tf.multiply(labels, tf.cast(weights, labels.dtype))
            return logits, labels, tf.cast(weights, tf.float32)
//  [python]
def tensors_to_item(self, keys_to_tensors):
    """Maps the given dictionary of tensors to a concatenated list of bboxes.

    Args:
      keys_to_tensors: a mapping of TF-Example keys to parsed tensors.

    Returns:
      [time, num_boxes, 4] tensor of bounding box coordinates, in order
          [y_min, x_min, y_max, x_max]. Whether the tensor is a SparseTensor
          or a dense Tensor is determined by the return_dense parameter. Empty
          positions in the sparse tensor are filled with -1.0 values.
    """
    sides = []
    for key in self._full_keys:
      value = keys_to_tensors[key]
      expanded_dims = tf.concat(
          [tf.to_int64(tf.shape(value)),
           tf.constant([1], dtype=tf.int64)], 0)
      side = tf.sparse_reshape(value, expanded_dims)
      sides.append(side)
    bounding_boxes = tf.sparse_concat(2, sides)
    if self._return_dense:
      bounding_boxes = tf.sparse_tensor_to_dense(
          bounding_boxes, default_value=self._default_value)
    return bounding_boxes
//  [python]
def tensors_to_item(self, keys_to_tensors):
    """Maps the given dictionary of tensors to a concatenated list of bboxes.

    Args:
      keys_to_tensors: a mapping of TF-Example keys to parsed tensors.

    Returns:
      [time, num_boxes, 4] tensor of bounding box coordinates, in order
          [y_min, x_min, y_max, x_max]. Whether the tensor is a SparseTensor
          or a dense Tensor is determined by the return_dense parameter. Empty
          positions in the sparse tensor are filled with -1.0 values.
    """
    sides = []
    for key in self._full_keys:
      value = keys_to_tensors[key]
      expanded_dims = tf.concat(
          [tf.cast(tf.shape(value), tf.int64),
           tf.constant([1], dtype=tf.int64)], 0)
      side = tf.sparse_reshape(value, expanded_dims)
      sides.append(side)
    bounding_boxes = tf.sparse_concat(2, sides)
    if self._return_dense:
      bounding_boxes = tf.sparse_tensor_to_dense(
          bounding_boxes, default_value=self._default_value)
    return bounding_boxes
//  [python]
def mine(x, z, name_net="estimator_network"):
  """Computes I(X, Z).

  Uses the algorithm in "Mutual Information Neural Estimation"
  (https://arxiv.org/pdf/1801.04062.pdf).

  Args:
    x: Samples from x [batch_size, size_x].
    z: Samples from z [batch_size, size_z].
    name_net: Scope for the variables forming the network.

  Returns:
    Estimate of the mutual information and the update op for the optimizer.
  """
  z_shuffled = vae.shuffle_codes(z)

  concat_x_x = tf.concat([x, x], axis=0)
  concat_z_z_shuffled = tf.stop_gradient(tf.concat([z, z_shuffled], axis=0))

  with tf.variable_scope(name_net, reuse=tf.AUTO_REUSE):
    d1_x = tf.layers.dense(concat_x_x, 20, name="d1_x")
    d1_z = tf.layers.dense(concat_z_z_shuffled, 20, name="d1_z")
    d1 = tf.nn.elu(d1_x + d1_z, name="d1")
    d2 = tf.layers.dense(d1, 1, name="d2")

  batch_size = tf.shape(x)[0]
  pred_x_z = d2[:batch_size]
  pred_x_z_shuffled = d2[batch_size:]
  loss = -(
      tf.reduce_mean(pred_x_z, axis=0) + tf.math.log(tf.to_float(batch_size)) -
      tf.math.reduce_logsumexp(pred_x_z_shuffled))
  all_variables = tf.trainable_variables()
  mine_vars = [var for var in all_variables if "estimator_network" in var.name]
  mine_op = tf.train.AdamOptimizer(learning_rate=0.01).minimize(
      loss=loss, var_list=mine_vars)
  return -loss, mine_op
//  [python]
def mine(x, z, name_net="estimator_network"):
  """Computes I(X, Z).

  Uses the algorithm in "Mutual Information Neural Estimation"
  (https://arxiv.org/pdf/1801.04062.pdf).

  Args:
    x: Samples from x [batch_size, size_x].
    z: Samples from z [batch_size, size_z].
    name_net: Scope for the variables forming the network.

  Returns:
    Estimate of the mutual information and the update op for the optimizer.
  """
  z_shuffled = vae.shuffle_codes(z)

  concat_x_x = tf.concat([x, x], axis=0)
  concat_z_z_shuffled = tf.stop_gradient(tf.concat([z, z_shuffled], axis=0))

  with tf.variable_scope(name_net, reuse=tf.AUTO_REUSE):
    d1_x = tf.layers.dense(concat_x_x, 20, name="d1_x")
    d1_z = tf.layers.dense(concat_z_z_shuffled, 20, name="d1_z")
    d1 = tf.nn.elu(d1_x + d1_z, name="d1")
    d2 = tf.layers.dense(d1, 1, name="d2")

  batch_size = tf.shape(x)[0]
  pred_x_z = d2[:batch_size]
  pred_x_z_shuffled = d2[batch_size:]
  loss = -(
      tf.reduce_mean(pred_x_z, axis=0) + tf.math.log(tf.cast(batch_size, tf.float32)) -
      tf.math.reduce_logsumexp(pred_x_z_shuffled))
  all_variables = tf.trainable_variables()
  mine_vars = [var for var in all_variables if "estimator_network" in var.name]
  mine_op = tf.train.AdamOptimizer(learning_rate=0.01).minimize(
      loss=loss, var_list=mine_vars)
  return -loss, mine_op
//  [python]
def model_fn(self, features, labels, mode, params):
    """TPUEstimator compatible model function.

    Args:
      features: Batch of images [batch_size, 64, 64, 3].
      labels: Tuple with batch of features [batch_size, 64, 64, 3] and the
        labels [batch_size, labels_size].
      mode: Mode for the TPUEstimator.
      params: Dict with parameters.

    Returns:
      TPU estimator.
    """

    is_training = (mode == tf.estimator.ModeKeys.TRAIN)
    labelled_features = labels[0]
    labels = tf.to_float(labels[1])
    data_shape = features.get_shape().as_list()[1:]
    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
      z_mean, z_logvar = self.gaussian_encoder(
          features, is_training=is_training)
      z_mean_labelled, _ = self.gaussian_encoder(
          labelled_features, is_training=is_training)
    z_sampled = self.sample_from_latent_distribution(z_mean, z_logvar)
    reconstructions = self.decode(z_sampled, data_shape, is_training)
    per_sample_loss = losses.make_reconstruction_loss(features, reconstructions)
    reconstruction_loss = tf.reduce_mean(per_sample_loss)
    kl_loss = compute_gaussian_kl(z_mean, z_logvar)
    gamma_annealed = make_annealer(self.gamma_sup, tf.train.get_global_step())
    supervised_loss = make_supervised_loss(z_mean_labelled, labels,
                                           self.factor_sizes)
    regularizer = self.unsupervised_regularizer(
        kl_loss, z_mean, z_logvar, z_sampled) + gamma_annealed * supervised_loss
    loss = tf.add(reconstruction_loss, regularizer, name="loss")
    elbo = tf.add(reconstruction_loss, kl_loss, name="elbo")
    if mode == tf.estimator.ModeKeys.TRAIN:
      optimizer = optimizers.make_vae_optimizer()
      update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
      train_op = optimizer.minimize(
          loss=loss, global_step=tf.train.get_global_step())
      train_op = tf.group([train_op, update_ops])
      tf.summary.scalar("reconstruction_loss", reconstruction_loss)
      tf.summary.scalar("elbo", -elbo)

      logging_hook = tf.train.LoggingTensorHook({
          "loss": loss,
          "reconstruction_loss": reconstruction_loss,
          "elbo": -elbo,
          "supervised_loss": supervised_loss
      },
                                                every_n_iter=100)
      return TPUEstimatorSpec(
          mode=mode,
          loss=loss,
          train_op=train_op,
          training_hooks=[logging_hook])
    elif mode == tf.estimator.ModeKeys.EVAL:
      return TPUEstimatorSpec(
          mode=mode,
          loss=loss,
          eval_metrics=(make_metric_fn("reconstruction_loss", "elbo",
                                       "regularizer", "kl_loss",
                                       "supervised_loss"), [
                                           reconstruction_loss, -elbo,
                                           regularizer, kl_loss, supervised_loss
                                       ]))
    else:
      raise NotImplementedError("Eval mode not supported.")
//  [python]
def model_fn(self, features, labels, mode, params):
    """TPUEstimator compatible model function.

    Args:
      features: Batch of images [batch_size, 64, 64, 3].
      labels: Tuple with batch of features [batch_size, 64, 64, 3] and the
        labels [batch_size, labels_size].
      mode: Mode for the TPUEstimator.
      params: Dict with parameters.

    Returns:
      TPU estimator.
    """

    is_training = (mode == tf.estimator.ModeKeys.TRAIN)
    labelled_features = labels[0]
    labels = tf.cast(labels[1], tf.float32)
    data_shape = features.get_shape().as_list()[1:]
    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
      z_mean, z_logvar = self.gaussian_encoder(
          features, is_training=is_training)
      z_mean_labelled, _ = self.gaussian_encoder(
          labelled_features, is_training=is_training)
    z_sampled = self.sample_from_latent_distribution(z_mean, z_logvar)
    reconstructions = self.decode(z_sampled, data_shape, is_training)
    per_sample_loss = losses.make_reconstruction_loss(features, reconstructions)
    reconstruction_loss = tf.reduce_mean(per_sample_loss)
    kl_loss = compute_gaussian_kl(z_mean, z_logvar)
    gamma_annealed = make_annealer(self.gamma_sup, tf.train.get_global_step())
    supervised_loss = make_supervised_loss(z_mean_labelled, labels,
                                           self.factor_sizes)
    regularizer = self.unsupervised_regularizer(
        kl_loss, z_mean, z_logvar, z_sampled) + gamma_annealed * supervised_loss
    loss = tf.add(reconstruction_loss, regularizer, name="loss")
    elbo = tf.add(reconstruction_loss, kl_loss, name="elbo")
    if mode == tf.estimator.ModeKeys.TRAIN:
      optimizer = optimizers.make_vae_optimizer()
      update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
      train_op = optimizer.minimize(
          loss=loss, global_step=tf.train.get_global_step())
      train_op = tf.group([train_op, update_ops])
      tf.summary.scalar("reconstruction_loss", reconstruction_loss)
      tf.summary.scalar("elbo", -elbo)

      logging_hook = tf.train.LoggingTensorHook({
          "loss": loss,
          "reconstruction_loss": reconstruction_loss,
          "elbo": -elbo,
          "supervised_loss": supervised_loss
      },
                                                every_n_iter=100)
      return TPUEstimatorSpec(
          mode=mode,
          loss=loss,
          train_op=train_op,
          training_hooks=[logging_hook])
    elif mode == tf.estimator.ModeKeys.EVAL:
      return TPUEstimatorSpec(
          mode=mode,
          loss=loss,
          eval_metrics=(make_metric_fn("reconstruction_loss", "elbo",
                                       "regularizer", "kl_loss",
                                       "supervised_loss"), [
                                           reconstruction_loss, -elbo,
                                           regularizer, kl_loss, supervised_loss
                                       ]))
    else:
      raise NotImplementedError("Eval mode not supported.")
//  [python]
def model_fn(self, features, labels, mode, params):
    """TPUEstimator compatible model function.

    Args:
      features: Batch of images [batch_size, 64, 64, 3].
      labels: Tuple with batch of features [batch_size, 64, 64, 3] and the
        labels [batch_size, labels_size].
      mode: Mode for the TPUEstimator.
      params: Dict with parameters.

    Returns:
      TPU Estimator.
    """

    is_training = (mode == tf.estimator.ModeKeys.TRAIN)
    labelled_features = labels[0]
    labels = tf.to_float(labels[1])
    data_shape = features.get_shape().as_list()[1:]
    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
      z_mean, z_logvar = self.gaussian_encoder(
          features, is_training=is_training)
      z_mean_labelled, _ = self.gaussian_encoder(
          labelled_features, is_training=is_training)
    z_sampled = self.sample_from_latent_distribution(z_mean, z_logvar)
    reconstructions = self.decode(
        tf.stop_gradient(z_sampled), data_shape, is_training)
    per_sample_loss = losses.make_reconstruction_loss(features, reconstructions)
    reconstruction_loss = tf.reduce_mean(per_sample_loss)
    supervised_loss = make_supervised_loss(z_mean_labelled, labels,
                                           self.factor_sizes)
    regularizer = supervised_loss
    loss = tf.add(reconstruction_loss, regularizer, name="loss")
    if mode == tf.estimator.ModeKeys.TRAIN:
      optimizer = optimizers.make_vae_optimizer()
      update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
      train_op = optimizer.minimize(
          loss=loss, global_step=tf.train.get_global_step())
      train_op = tf.group([train_op, update_ops])
      tf.summary.scalar("reconstruction_loss", reconstruction_loss)

      logging_hook = tf.train.LoggingTensorHook({
          "loss": loss,
          "reconstruction_loss": reconstruction_loss,
          "supervised_loss": supervised_loss
      },
                                                every_n_iter=100)
      return TPUEstimatorSpec(
          mode=mode,
          loss=loss,
          train_op=train_op,
          training_hooks=[logging_hook])
    elif mode == tf.estimator.ModeKeys.EVAL:
      return TPUEstimatorSpec(
          mode=mode,
          loss=loss,
          eval_metrics=(make_metric_fn("reconstruction_loss", "regularizer",
                                       "supervised_loss"),
                        [reconstruction_loss, regularizer, supervised_loss]))
    else:
      raise NotImplementedError("Eval mode not supported.")
//  [python]
def model_fn(self, features, labels, mode, params):
    """TPUEstimator compatible model function.

    Args:
      features: Batch of images [batch_size, 64, 64, 3].
      labels: Tuple with batch of features [batch_size, 64, 64, 3] and the
        labels [batch_size, labels_size].
      mode: Mode for the TPUEstimator.
      params: Dict with parameters.

    Returns:
      TPU Estimator.
    """

    is_training = (mode == tf.estimator.ModeKeys.TRAIN)
    labelled_features = labels[0]
    labels = tf.cast(labels[1], tf.float32)
    data_shape = features.get_shape().as_list()[1:]
    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
      z_mean, z_logvar = self.gaussian_encoder(
          features, is_training=is_training)
      z_mean_labelled, _ = self.gaussian_encoder(
          labelled_features, is_training=is_training)
    z_sampled = self.sample_from_latent_distribution(z_mean, z_logvar)
    reconstructions = self.decode(
        tf.stop_gradient(z_sampled), data_shape, is_training)
    per_sample_loss = losses.make_reconstruction_loss(features, reconstructions)
    reconstruction_loss = tf.reduce_mean(per_sample_loss)
    supervised_loss = make_supervised_loss(z_mean_labelled, labels,
                                           self.factor_sizes)
    regularizer = supervised_loss
    loss = tf.add(reconstruction_loss, regularizer, name="loss")
    if mode == tf.estimator.ModeKeys.TRAIN:
      optimizer = optimizers.make_vae_optimizer()
      update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
      train_op = optimizer.minimize(
          loss=loss, global_step=tf.train.get_global_step())
      train_op = tf.group([train_op, update_ops])
      tf.summary.scalar("reconstruction_loss", reconstruction_loss)

      logging_hook = tf.train.LoggingTensorHook({
          "loss": loss,
          "reconstruction_loss": reconstruction_loss,
          "supervised_loss": supervised_loss
      },
                                                every_n_iter=100)
      return TPUEstimatorSpec(
          mode=mode,
          loss=loss,
          train_op=train_op,
          training_hooks=[logging_hook])
    elif mode == tf.estimator.ModeKeys.EVAL:
      return TPUEstimatorSpec(
          mode=mode,
          loss=loss,
          eval_metrics=(make_metric_fn("reconstruction_loss", "regularizer",
                                       "supervised_loss"),
                        [reconstruction_loss, regularizer, supervised_loss]))
    else:
      raise NotImplementedError("Eval mode not supported.")
//  [python]
def res_net(n_examples,
            input_shape,
            num_classes,
            batchnorm=False,
            variational='full'):
  """Wrapper for build_resnet_v1.

  Args:
    n_examples (int): number of training points.
    input_shape (list): input shape.
    num_classes (int): number of classes (CIFAR10 has 10).
    batchnorm (bool): use of batchnorm layers.
    variational (str): 'none', 'hybrid', 'full'. whether to use variational
      inference for zero, some, or all layers.

  Returns:
      model (Model): Keras model instance whose output is a
        tfp.distributions.Categorical distribution.
  """
  inputs = tf.keras.layers.Input(shape=input_shape)
  x = build_resnet_v1(
      inputs,
      depth=20,
      variational=variational,
      batchnorm=batchnorm,
      n_examples=n_examples)

  p_fn, q_fn = mean_field_fn(empirical_bayes=True)

  def normalized_kl_fn(q, p, _):
    return tfp.distributions.kl_divergence(q, p) / tf.to_float(n_examples)

  logits = tfp.layers.DenseLocalReparameterization(
      num_classes,
      kernel_prior_fn=p_fn,
      kernel_posterior_fn=q_fn,
      bias_prior_fn=p_fn,
      bias_posterior_fn=q_fn,
      kernel_divergence_fn=normalized_kl_fn,
      bias_divergence_fn=normalized_kl_fn)(x)
  outputs = tf.keras.layers.Lambda(lambda x: ed.Categorical(logits=x))(logits)
  return tf.keras.models.Model(inputs=inputs, outputs=outputs)